Update app.py

app.py

@@ -2,16 +2,16 @@ import spaces
 import gradio as gr
 import torch
 from PIL import Image
-from diffusers import DiffusionPipeline, AutoencoderTiny
+from diffusers import DiffusionPipeline
 import random
 import uuid
-from typing import Tuple, Union, List, Optional, Any, Dict
+from typing import Union, List, Optional
 import numpy as np
 import time
 import zipfile
 
 # Description for the app
-DESCRIPTION = """## Qwen Image Hpc/."""
+DESCRIPTION = """## Qwen Image Generator"""
 
 # Helper functions
 def save_image(img):
@@ -27,11 +27,10 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
-# Load Qwen/Qwen-Image pipeline with taef1 VAE
+# Load Qwen/Qwen-Image pipeline
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
-taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
-pipe_qwen = DiffusionPipeline.from_pretrained("Qwen/Qwen-Image", torch_dtype=dtype, vae=taef1).to(device)
+pipe_qwen = DiffusionPipeline.from_pretrained("Qwen/Qwen-Image", torch_dtype=dtype).to(device)
 
 # Aspect ratios
 aspect_ratios = {
@@ -90,6 +89,37 @@ def generate_qwen(
     
     return image_paths, seed, f"{duration:.2f}", zip_path
 
+# Wrapper function to handle UI logic
+@spaces.GPU
+def generate(
+    prompt: str,
+    negative_prompt: str,
+    use_negative_prompt: bool,
+    seed: int,
+    width: int,
+    height: int,
+    guidance_scale: float,
+    randomize_seed: bool,
+    num_inference_steps: int,
+    num_images: int,
+    zip_images: bool,
+    progress=gr.Progress(track_tqdm=True),
+):
+    final_negative_prompt = negative_prompt if use_negative_prompt else ""
+    return generate_qwen(
+        prompt=prompt,
+        negative_prompt=final_negative_prompt,
+        seed=seed,
+        width=width,
+        height=height,
+        guidance_scale=guidance_scale,
+        randomize_seed=randomize_seed,
+        num_inference_steps=num_inference_steps,
+        num_images=num_images,
+        zip_images=zip_images,
+        progress=progress,
+    )
+
 # Examples
 examples = [
     "An attractive young woman with blue eyes lying face down on the bed, light white and light amber, timeless beauty, sunrays shine upon it",
@@ -134,6 +164,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         use_negative_prompt = gr.Checkbox(
             label="Use negative prompt",
             value=False,
+            visible=True
         )
         negative_prompt = gr.Text(
             label="Negative prompt",
@@ -213,10 +244,11 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
     # Run button and prompt submit
     gr.on(
         triggers=[prompt.submit, run_button.click],
-        fn=generate_qwen,
+        fn=generate,
         inputs=[
             prompt,
             negative_prompt,
+            use_negative_prompt,
             seed,
             width,
             height,
@@ -235,7 +267,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         examples=examples,
         inputs=prompt,
         outputs=[result, seed_display, generation_time, zip_file],
-        fn=generate_qwen,
+        fn=generate,
         cache_examples=False,
     )