22

app.py

@@ -7,6 +7,7 @@ from diffsynth.pipelines.qwen_image import (
     QwenImageUnit_Image2LoRAEncode, QwenImageUnit_Image2LoRADecode
 )
 from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
+from transformers import Qwen2_5_VLForConditionalGeneration, QwenImageTransformer2DModel
 from safetensors.torch import save_file
 import torch, math
 from PIL import Image
@@ -20,28 +21,28 @@ URL_PUBLIC = "https://huggingface.co/spaces/AiSudo/Qwen-Image-to-LoRA/blob/main"
 DTYPE = torch.bfloat16
 MAX_SEED = np.iinfo(np.int32).max
 
-vram_config_disk_offload = {
-    "offload_dtype": "disk",
-    "offload_device": "disk",
-    "onload_dtype": "disk",
-    "onload_device": "disk",
-    "preparing_dtype": torch.bfloat16,
-    "preparing_device": "cuda",
-    "computation_dtype": torch.bfloat16,
-    "computation_device": "cuda",
-}
-
 # vram_config_disk_offload = {
-#     "offload_dtype": torch.bfloat16,
-#     "offload_device": "cuda",
-#     "onload_dtype": torch.bfloat16,
-#     "onload_device": "cuda",
+#     "offload_dtype": "disk",
+#     "offload_device": "disk",
+#     "onload_dtype": "disk",
+#     "onload_device": "disk",
 #     "preparing_dtype": torch.bfloat16,
 #     "preparing_device": "cuda",
 #     "computation_dtype": torch.bfloat16,
 #     "computation_device": "cuda",
 # }
 
+vram_config_disk_offload = {
+    "offload_dtype": torch.bfloat16,
+    "offload_device": "cuda",
+    "onload_dtype": torch.bfloat16,
+    "onload_device": "cuda",
+    "preparing_dtype": torch.bfloat16,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+
 # Load models
 pipe_lora = QwenImagePipeline.from_pretrained(
     torch_dtype=torch.bfloat16,
@@ -81,17 +82,17 @@ vram_config = {
     "computation_device": "cuda",
 }
 
-pipe_imagen = QwenImagePipeline.from_pretrained(
-    torch_dtype=torch.bfloat16,
-    device="cuda",
-    model_configs=[
-        ModelConfig(download_source="huggingface", model_id="Qwen/Qwen-Image", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors", **vram_config),
-        ModelConfig(download_source="huggingface", model_id="Qwen/Qwen-Image", origin_file_pattern="text_encoder/model*.safetensors", **vram_config),
-        ModelConfig(download_source="huggingface", model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors", **vram_config),
-    ],
-    tokenizer_config=ModelConfig(download_source="huggingface", model_id="Qwen/Qwen-Image", origin_file_pattern="tokenizer/"),
-    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
-)
+# pipe_imagen = QwenImagePipeline.from_pretrained(
+#     torch_dtype=torch.bfloat16,
+#     device="cuda",
+#     model_configs=[
+#         ModelConfig(download_source="huggingface", model_id="Qwen/Qwen-Image", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors", **vram_config),
+#         ModelConfig(download_source="huggingface", model_id="Qwen/Qwen-Image", origin_file_pattern="text_encoder/model*.safetensors", **vram_config),
+#         ModelConfig(download_source="huggingface", model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors", **vram_config),
+#     ],
+#     tokenizer_config=ModelConfig(download_source="huggingface", model_id="Qwen/Qwen-Image", origin_file_pattern="tokenizer/"),
+#     vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
+# )
 
 # scheduler_config = {
 #     "base_image_seq_len": 256,
@@ -109,14 +110,39 @@ pipe_imagen = QwenImagePipeline.from_pretrained(
 #     "use_exponential_sigmas": False,
 #     "use_karras_sigmas": False,
 # }
-# scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
+#
+quant_model = "AiSudo/Qwen-Image-fp8-4steps"
+
+transformer=QwenImageTransformer2DModel.from_pretrained(
+    quant_model,
+    subfolder='transformer',
+    torch_dtype=torch.bfloat16,
+    use_safetensors=False,
+    device_map="cuda"
+)
+
+text_encoder=Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    quant_model,
+    subfolder='text_encoder',
+    torch_dtype=torch.bfloat16,
+    use_safetensors=False,
+    device_map="cuda"
+)
+
+scheduler = FlowMatchEulerDiscreteScheduler.from_config(
+    quant_model,
+    subfolder="scheduler"
+)
+
+pipe_imagen = DiffusionPipeline.from_pretrained(
+    "Qwen/Qwen-Image",
+    transformer=transformer,
+    text_encoder=text_encoder,
+    scheduler=scheduler,
+    torch_dtype=torch.bfloat16,
+    low_cpu_mem_usage=False,
+).to("cuda")
 
-# pipe_imagen = DiffusionPipeline.from_pretrained(
-#     "Qwen/Qwen-Image",
-#     scheduler=scheduler,
-#     torch_dtype=torch.bfloat16,
-#     low_cpu_mem_usage=False,
-# ).to("cuda")
 # pipe_imagen.load_lora_weights(
 #     "lightx2v/Qwen-Image-Lightning",
 #     weight_name="Qwen-Image-Lightning-8steps-V2.0-bf16.safetensors",
@@ -126,7 +152,7 @@ pipe_imagen = QwenImagePipeline.from_pretrained(
 # pipe_imagen.fuse_lora(adapter_names=["lightning_steps"], lora_scale=1.25)
 # pipe_imagen.unload_lora_weights()
 
-# pipe_imagen.to("cuda")
+pipe_imagen.to("cuda")
 
 @spaces.GPU
 def generate_lora(
@@ -169,17 +195,17 @@ def generate_image(
     num_inference_steps=8,
     progress=gr.Progress(track_tqdm=True),
 ):
-    lora_path = f"loras/{lora_name}"
-    pipe_imagen.clear_lora()
-    pipe_imagen.load_lora(pipe_imagen.dit, lora_path)
-
-    # pipe_imagen.unload_lora_weights()
-    # pipe_imagen.load_lora_weights(
-    #     "loras",
-    #     weight_name=lora_name,
-    #     adapter_name="generated_lora"
-    # )
-    # pipe_imagen.set_adapters("generated_lora", adapter_weights=lora_strength)
+    # lora_path = f"loras/{lora_name}"
+    # pipe_imagen.clear_lora()
+    # pipe_imagen.load_lora(pipe_imagen.dit, lora_path)
+
+    pipe_imagen.unload_lora_weights()
+    pipe_imagen.load_lora_weights(
+        "loras",
+        weight_name=lora_name,
+        adapter_name="generated_lora"
+    )
+    pipe_imagen.set_adapters("generated_lora", adapter_weights=lora_strength)
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)

static/header.html

@@ -6,6 +6,6 @@
         <p>
             Generate LoRA from a few images.
             <br>
-            Demo by <a href="https://aisudo.com/" target="_blank">AiSudo</a> 😊 
+            Demo by <a href="https://aisudo.com/ai-models/qwen-2512-8-step-mdl032qw8s" target="_blank">AiSudo</a> 😊
     </div>
 </div>