22

app.py

@@ -31,7 +31,7 @@ vram_config_disk_offload = {
 }
 
 # Load models
-pipe = QwenImagePipeline.from_pretrained(
+pipe_lora = QwenImagePipeline.from_pretrained(
     torch_dtype=torch.bfloat16,
     device="cuda",
     model_configs=[
@@ -58,6 +58,11 @@ pipe = QwenImagePipeline.from_pretrained(
     vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
 )
 
+pipe_imagen = QwenImagePipeline.from_pretrained(
+    "Qwen/Qwen-Image", 
+    torch_dtype=DTYPE
+).to("cuda")
+
 
 @spaces.GPU
 def generate_lora(
@@ -86,8 +91,8 @@ def generate_lora(
 
     # Model inference
     with torch.no_grad():
-        embs = QwenImageUnit_Image2LoRAEncode().process(pipe, image2lora_images=input_images)
-        lora = QwenImageUnit_Image2LoRADecode().process(pipe, **embs)["lora"]
+        embs = QwenImageUnit_Image2LoRAEncode().process(pipe_lora, image2lora_images=input_images)
+        lora = QwenImageUnit_Image2LoRADecode().process(pipe_lora, **embs)["lora"]
 
     lora_name = f"{ulid}.safetensors"
     lora_path = f"loras/{lora_name}"
@@ -109,7 +114,27 @@ def generate_image(
     num_inference_steps=8,
     progress=gr.Progress(track_tqdm=True),
 ):
-    
+    lora_path = f"loras/{lora_name}"
+    pipe_imagen.load_lora(pipe_imagen.dit, lora_path)
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+
+    generator = torch.Generator().manual_seed(seed)
+
+    image = pipe_imagen(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        num_inference_steps=num_inference_steps,
+        width=width,
+        height=height,
+        generator=generator,
+        true_cfg_scale=guidance_scale,
+        guidance_scale=1.0  # Use a fixed default for distilled guidance
+    ).images[0]
+
+    return image, seed
+
     return True
 
 
@@ -156,7 +181,7 @@ with gr.Blocks() as demo:
             with gr.Column():
                 lora_name = gr.Textbox(label="Generated LoRA path",lines=2, interactive=False)
                 lora_download = gr.DownloadButton(label=f"Download LoRA", visible=False)
-        with gr.Column(visible=False, elem_classes='imagen-container') as imagen_container:
+        with gr.Column(visible=False, elem_id='imagen-container') as imagen_container:
             gr.Markdown("### After your LoRA is ready, you can try generate image here.")
             with gr.Row():
                 with gr.Column():
@@ -183,7 +208,7 @@ with gr.Blocks() as demo:
                             minimum=1,
                             maximum=50,
                             step=1,
-                            value=20,
+                            value=25,
                         )
                         with gr.Row():
                             width = gr.Slider(
@@ -244,7 +269,7 @@ with gr.Blocks() as demo:
             guidance_scale,
             num_inference_steps,
         ],
-        outputs=[lora_name, lora_download, imagen_container],
+        outputs=[output_image, seed],
     )