22

app.py

@@ -59,8 +59,15 @@ pipe_lora = QwenImagePipeline.from_pretrained(
 )
 
 pipe_imagen = QwenImagePipeline.from_pretrained(
-    "Qwen/Qwen-Image",
-    device="cuda"
+    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,
 )
 
 
@@ -73,10 +80,10 @@ def generate_lora(
     ulid = str(ULID()).lower()[:12]
     print(f"ulid: {ulid}")
 
-    print("before: ", input_images)
+    # print("before: ", input_images)
     # input_images = [item[0] for item in input_images]
     input_images = [Image.open(filepath) for filepath, _ in input_images]
-    print("after: ", input_images)
+    # print("after: ", input_images)
     if not input_images:
         print("images are empty.")
     # Load images