Update app.py

app.py

@@ -1,24 +1,41 @@
-import spaces
-import gradio as gr
+import os
+import json
+import copy
+import time
+import random
+import logging
+import numpy as np
+from typing import Any, Dict, List, Optional, Union
 import torch
 from PIL import Image
+import gradio as gr
+import spaces
 from diffusers import DiffusionPipeline
-import random
-import uuid
-from typing import Union, List, Optional
-import numpy as np
-import time
-import zipfile
-import os
+from huggingface_hub import (
+    hf_hub_download,
+    HfFileSystem,
+    ModelCard,
+    snapshot_download)
+from diffusers.utils import load_image
 import requests
 from urllib.parse import urlparse
 import tempfile
 import shutil
+import uuid
+import zipfile
 
-# Description for the app
-DESCRIPTION = """## Qwen Image Hpc/."""
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.16,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
 
-# Helper functions
 def save_image(img):
     unique_name = str(uuid.uuid4()) + ".png"
     img.save(unique_name)
@@ -29,25 +46,138 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 2048
+# Qwen Image pipeline with live preview capability
+@torch.inference_mode()
+def qwen_pipe_call_that_returns_an_iterable_of_images(
+    self,
+    prompt: Union[str, List[str]] = None,
+    negative_prompt: Optional[Union[str, List[str]]] = None,
+    height: Optional[int] = None,
+    width: Optional[int] = None,
+    num_inference_steps: int = 50,
+    guidance_scale: float = 4.0,
+    num_images_per_prompt: Optional[int] = 1,
+    generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+    output_type: Optional[str] = "pil",
+):
+    height = height or 1024
+    width = width or 1024
+    
+    batch_size = 1 if isinstance(prompt, str) else len(prompt)
+    device = self._execution_device
+    
+    # Generate intermediate images during the process
+    for i in range(num_inference_steps):
+        if i % 5 == 0:  # Show progress every 5 steps
+            # Generate partial result
+            temp_result = self(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                height=height,
+                width=width,
+                guidance_scale=guidance_scale,
+                num_inference_steps=max(1, i + 1),
+                num_images_per_prompt=num_images_per_prompt,
+                generator=generator,
+                output_type=output_type,
+            ).images[0]
+            yield temp_result
+            torch.cuda.empty_cache()
+    
+    # Final high-quality result
+    final_result = self(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        height=height,
+        width=width,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        num_images_per_prompt=num_images_per_prompt,
+        generator=generator,
+        output_type=output_type,
+    ).images[0]
+    
+    yield final_result
 
-# Load Qwen/Qwen-Image pipeline
+loras = [
+    # Sample Qwen-compatible LoRAs
+    {
+        "image": "https://huggingface.co/prithivMLmods/Qwen-Image-Studio-Realism/resolve/main/images/2.png",
+        "title": "Studio Realism",
+        "repo": "prithivMLmods/Qwen-Image-Studio-Realism",
+        "weights": "qwen-studio-realism.safetensors",
+        "trigger_word": "Studio Realism"
+    },
+    {
+        "image": "https://huggingface.co/prithivMLmods/Qwen-Image-Sketch-Smudge/resolve/main/images/1.png",
+        "title": "Sketch Smudge",
+        "repo": "prithivMLmods/Qwen-Image-Sketch-Smudge",
+        "weights": "qwen-sketch-smudge.safetensors",
+        "trigger_word": "Sketch Smudge"
+    },
+    {
+        "image": "https://huggingface.co/prithivMLmods/Qwen-Image-Anime-LoRA/resolve/main/images/1.png",
+        "title": "Qwen Anime",
+        "repo": "prithivMLmods/Qwen-Image-Anime-LoRA",
+        "weights": "qwen-anime.safetensors",
+        "trigger_word": "Qwen Anime"
+    },
+    {
+        "image": "https://huggingface.co/prithivMLmods/Qwen-Image-Synthetic-Face/resolve/main/images/2.png",
+        "title": "Synthetic Face",
+        "repo": "prithivMLmods/Qwen-Image-Synthetic-Face",
+        "weights": "qwen-synthetic-face.safetensors",
+        "trigger_word": "Synthetic Face"
+    },
+    {
+        "image": "huggingface.co/prithivMLmods/Qwen-Image-Fragmented-Portraiture/resolve/main/images/3.png",
+        "title": "Fragmented Portraiture",
+        "repo": "prithivMLmods/Qwen-Image-Fragmented-Portraiture",
+        "weights": "qwen-fragmented-portraiture.safetensors",
+        "trigger_word": "Fragmented Portraiture"
+    },
+]
+
+#--------------------------------------------------Model Initialization-----------------------------------------------------------------------------------------#
 dtype = torch.bfloat16
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+base_model = "Qwen/Qwen-Image"
 
-# --- Model Loading ---
-pipe_qwen = DiffusionPipeline.from_pretrained("Qwen/Qwen-Image", torch_dtype=dtype).to(device)
+# Load Qwen Image pipeline
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype).to(device)
 
-# Aspect ratios
+# Add aspect ratios for Qwen
 aspect_ratios = {
-    "1:1": (1328, 1328),
-    "16:9": (1664, 928),
-    "9:16": (928, 1664),
-    "4:3": (1472, 1140),
-    "3:4": (1140, 1472)
+    "1:1": (1024, 1024),
+    "16:9": (1344, 768),
+    "9:16": (768, 1344),
+    "4:3": (1152, 896),
+    "3:4": (896, 1152),
+    "3:2": (1216, 832),
+    "2:3": (832, 1216)
 }
 
+MAX_SEED = 2**32-1
+
+# Add the custom method to the pipeline
+pipe.qwen_pipe_call_that_returns_an_iterable_of_images = qwen_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
+
+class calculateDuration:
+    def __init__(self, activity_name=""):
+        self.activity_name = activity_name
+
+    def __enter__(self):
+        self.start_time = time.time()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.end_time = time.time()
+        self.elapsed_time = self.end_time - self.start_time
+        if self.activity_name:
+            print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds")
+        else:
+            print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
+
 def load_lora_opt(pipe, lora_input):
     lora_input = lora_input.strip()
     if not lora_input:
@@ -60,7 +190,6 @@ def load_lora_opt(pipe, lora_input):
 
     if lora_input.startswith("http"):
         url = lora_input
-
         # Repo page (no blob/resolve)
         if "huggingface.co" in url and "/blob/" not in url and "/resolve/" not in url:
             repo_id = urlparse(url).path.strip("/")
@@ -74,7 +203,6 @@ def load_lora_opt(pipe, lora_input):
         # Download direct file
         tmp_dir = tempfile.mkdtemp()
         local_path = os.path.join(tmp_dir, os.path.basename(urlparse(url).path))
-
         try:
             print(f"Downloading LoRA from {url}...")
             resp = requests.get(url, stream=True)
@@ -87,272 +215,315 @@ def load_lora_opt(pipe, lora_input):
         finally:
             shutil.rmtree(tmp_dir, ignore_errors=True)
 
-# Generation function for Qwen/Qwen-Image
-@spaces.GPU(duration=120)
-def generate_qwen(
-    prompt: str,
-    negative_prompt: str = "",
-    seed: int = 0,
-    width: int = 1024,
-    height: int = 1024,
-    guidance_scale: float = 4.0,
-    randomize_seed: bool = False,
-    num_inference_steps: int = 50,
-    num_images: int = 1,
-    zip_images: bool = False,
-    lora_input: str = "",
-    lora_scale: float = 1.0,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device).manual_seed(seed)
-    
-    start_time = time.time()
-
-    current_adapters = pipe_qwen.get_list_adapters()
-    for adapter in current_adapters:
-        pipe_qwen.delete_adapters(adapter)
-    pipe_qwen.disable_lora()
-
-    use_lora = False
-    if lora_input and lora_input.strip() != "":
-        load_lora_opt(pipe_qwen, lora_input)
-        pipe_qwen.set_adapters(["default"], adapter_weights=[lora_scale])
-        use_lora = True
-    
-    images = pipe_qwen(
-        prompt=prompt,
-        negative_prompt=negative_prompt if negative_prompt else "",
-        height=height,
-        width=width,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        num_images_per_prompt=num_images,
-        generator=generator,
-        output_type="pil",
-    ).images
+def update_selection(evt: gr.SelectData, width, height):
+    selected_lora = loras[evt.index]
+    new_placeholder = f"Type a prompt for {selected_lora['title']}"
+    lora_repo = selected_lora["repo"]
+    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✅"
     
-    end_time = time.time()
-    duration = end_time - start_time
+    if "aspect" in selected_lora:
+        if selected_lora["aspect"] == "portrait":
+            width = 768
+            height = 1024
+        elif selected_lora["aspect"] == "landscape":
+            width = 1024
+            height = 768
+        else:
+            width = 1024
+            height = 1024
     
-    image_paths = [save_image(img) for img in images]
-    zip_path = None
-    if zip_images:
-        zip_name = str(uuid.uuid4()) + ".zip"
-        with zipfile.ZipFile(zip_name, 'w') as zipf:
-            for i, img_path in enumerate(image_paths):
-                zipf.write(img_path, arcname=f"Img_{i}.png")
-        zip_path = zip_name
-
-    # Clean up adapters
-    current_adapters = pipe_qwen.get_list_adapters()
-    for adapter in current_adapters:
-        pipe_qwen.delete_adapters(adapter)
-    pipe_qwen.disable_lora()
+    return (
+        gr.update(placeholder=new_placeholder),
+        updated_text,
+        evt.index,
+        width,
+        height,
+    )
+
+@spaces.GPU(duration=120)
+def generate_image(prompt_mash, negative_prompt, steps, seed, cfg_scale, width, height, lora_scale, progress):
+    pipe.to("cuda")
+    generator = torch.Generator(device="cuda").manual_seed(seed)
     
-    return image_paths, seed, f"{duration:.2f}", zip_path
+    with calculateDuration("Generating image"):
+        # Generate image with live preview
+        for img in pipe.qwen_pipe_call_that_returns_an_iterable_of_images(
+            prompt=prompt_mash,
+            negative_prompt=negative_prompt,
+            num_inference_steps=steps,
+            guidance_scale=cfg_scale,
+            width=width,
+            height=height,
+            generator=generator,
+        ):
+            yield img
+
+def set_dimensions(ar):
+    w, h = aspect_ratios[ar]
+    return gr.update(value=w), gr.update(value=h)
 
-# Wrapper function to handle UI logic
 @spaces.GPU(duration=120)
-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,
-    lora_input: str,
-    lora_scale: float,
-    progress=gr.Progress(track_tqdm=True),
-):
+def run_lora(prompt, negative_prompt, use_negative_prompt, aspect_ratio, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
+    if selected_index is None:
+        raise gr.Error("You must select a LoRA before proceeding.🧨")
+
+    selected_lora = loras[selected_index]
+    lora_path = selected_lora["repo"]
+    trigger_word = selected_lora["trigger_word"]
+
+    # Set dimensions based on aspect ratio
+    width, height = aspect_ratios[aspect_ratio]
+
+    if trigger_word:
+        if "trigger_position" in selected_lora:
+            if selected_lora["trigger_position"] == "prepend":
+                prompt_mash = f"{trigger_word} {prompt}"
+            else:
+                prompt_mash = f"{prompt} {trigger_word}"
+        else:
+            prompt_mash = f"{trigger_word} {prompt}"
+    else:
+        prompt_mash = prompt
+
+    # Handle negative prompt
     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,
-        lora_input=lora_input,
-        lora_scale=lora_scale,
-        progress=progress,
-    )
 
-# Examples
-examples = [
-    "A decadent slice of layered chocolate cake on a ceramic plate with a drizzle of chocolate syrup and powdered sugar dusted on top. photographed from a slightly low angle with high resolution, natural soft lighting, rich contrast, shallow depth of field, and professional color grading to highlight the dessert’s textures --ar 85:128 --v 6.0 --style raw",
-    "A beautifully decorated round chocolate birthday cake with rich chocolate frosting and elegant piping, topped with the name 'Qwen' written in white icing. placed on a wooden cake stand with scattered chocolate shavings around, softly lit with natural light, high resolution, professional food photography, clean background, no branding --ar 85:128 --v 6.0 --style raw",
-    "Realistic still life photography style: A single, fresh apple, resting on a clean, soft-textured surface. The apple is slightly off-center, softly backlit to highlight its natural gloss and subtle color gradients—deep crimson red blending into light golden hues. Fine details such as small blemishes, dew drops, and a few light highlights enhance its lifelike appearance. A shallow depth of field gently blurs the neutral background, drawing full attention to the apple. Hyper-detailed 8K resolution, studio lighting, photorealistic render, emphasizing texture and form.",
-    "一幅精致细腻的工笔画，画面中心是一株蓬勃生长的红色牡丹，花朵繁茂，既有盛开的硕大花瓣，也有含苞待放的花蕾，层次丰富，色彩艳丽而不失典雅。牡丹枝叶舒展，叶片浓绿饱满，脉络清晰可见，与红花相映成趣。一只蓝紫色蝴蝶仿佛被画中花朵吸引，停驻在画面中央的一朵盛开牡丹上，流连忘返，蝶翼轻展，细节逼真，仿佛随时会随风飞舞。整幅画作笔触工整严谨，色彩浓郁鲜明，展现出中国传统工笔画的精妙与神韵，画面充满生机与灵动之感。",
-    "A young girl wearing school uniform stands in a classroom, writing on a chalkboard. The text Introducing Qwen-Image, a foundational image generation model that excels in complex text rendering and precise image editing appears in neat white chalk at the center of the blackboard. Soft natural light filters through windows, casting gentle shadows. The scene is rendered in a realistic photography style with fine details, shallow depth of field, and warm tones. The girl's focused expression and chalk dust in the air add dynamism. Background elements include desks and educational posters, subtly blurred to emphasize the central action. Ultra-detailed 32K resolution, DSLR-quality, soft bokeh effect, documentary-style composition",
-    "手绘风格的水循环示意图，整体画面呈现出一幅生动形象的水循环过程图解。画面中央是一片起伏的山脉和山谷，山谷中流淌着一条清澈的河流，河流最终汇入一片广阔的海洋。山体和陆地上绘制有绿色植被。画面下方为地下水层，用蓝色渐变色块表现，与地表水形成层次分明的空间关系。太阳位于画面右上角，促使地表水蒸发，用上升的曲线箭头表示蒸发过程。云朵漂浮在空中，由白色棉絮状绘制而成，部分云层厚重，表示水汽凝结成雨，用向下箭头连接表示降雨过程。雨水以蓝色线条和点状符号表示，从云中落下，补充河流与地下水。整幅图以卡通手绘风格呈现，线条柔和，色彩明亮，标注清晰。背景为浅黄色纸张质感，带有轻微的手绘纹理。"
-]
+    with calculateDuration("Unloading LoRA"):
+        # Clear existing adapters
+        current_adapters = pipe.get_list_adapters() if hasattr(pipe, 'get_list_adapters') else []
+        for adapter in current_adapters:
+            if hasattr(pipe, 'delete_adapters'):
+                pipe.delete_adapters(adapter)
+        if hasattr(pipe, 'disable_lora'):
+            pipe.disable_lora()
+
+    # Load new LoRA weights
+    with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
+        weight_name = selected_lora.get("weights", None)
+        load_lora_opt(pipe, lora_path)
+        if hasattr(pipe, 'set_adapters'):
+            pipe.set_adapters(["default"], adapter_weights=[lora_scale])
+
+    with calculateDuration("Randomizing seed"):
+        if randomize_seed:
+            seed = random.randint(0, MAX_SEED)
+
+    image_generator = generate_image(prompt_mash, final_negative_prompt, steps, seed, cfg_scale, width, height, lora_scale, progress)
+
+    final_image = None
+    step_counter = 0
+    for image in image_generator:
+        step_counter += 1
+        final_image = image
+        progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'
+        yield image, seed, gr.update(value=progress_bar, visible=True)
+
+    yield final_image, seed, gr.update(value=progress_bar, visible=False)
+
+def get_huggingface_safetensors(link):
+    split_link = link.split("/")
+    if len(split_link) == 2:
+        model_card = ModelCard.load(link)
+        base_model = model_card.data.get("base_model")
+        print(base_model)
+
+        # Allow Qwen models
+        if base_model and "qwen" not in base_model.lower():
+            raise Exception("Qwen-compatible LoRA Not Found!")
+
+        image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
+        trigger_word = model_card.data.get("instance_prompt", "")
+        image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
+        
+        fs = HfFileSystem()
+        try:
+            list_of_files = fs.ls(link, detail=False)
+            for file in list_of_files:
+                if file.endswith(".safetensors"):
+                    safetensors_name = file.split("/")[-1]
+                if not image_url and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp")):
+                    image_elements = file.split("/")
+                    image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"
+        except Exception as e:
+            print(e)
+            gr.Warning(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
+            raise Exception(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
+        
+        return split_link[1], link, safetensors_name, trigger_word, image_url
+
+def check_custom_model(link):
+    if link.startswith("https://"):
+        if link.startswith("https://huggingface.co") or link.startswith("https://www.huggingface.co"):
+            link_split = link.split("huggingface.co/")
+            return get_huggingface_safetensors(link_split[1])
+    else: 
+        return get_huggingface_safetensors(link)
+
+def add_custom_lora(custom_lora):
+    global loras
+    if custom_lora:
+        try:
+            title, repo, path, trigger_word, image = check_custom_model(custom_lora)
+            print(f"Loaded custom LoRA: {repo}")
+            card = f'''
+            <div class="custom_lora_card">
+              <span>Loaded custom LoRA:</span>
+              <div class="card_internal">
+                <img src="{image}" />
+                <div>
+                    <h3>{title}</h3>
+                    <small>{"Using: <code><b>"+trigger_word+"</code></b> as the trigger word" if trigger_word else "No trigger word found. If there's a trigger word, include it in your prompt"}<br></small>
+                </div>
+              </div>
+            </div>
+            '''
+            existing_item_index = next((index for (index, item) in enumerate(loras) if item['repo'] == repo), None)
+            if not existing_item_index:
+                new_item = {
+                    "image": image,
+                    "title": title,
+                    "repo": repo,
+                    "weights": path,
+                    "trigger_word": trigger_word
+                }
+                print(new_item)
+                existing_item_index = len(loras)
+                loras.append(new_item)
+
+            return gr.update(visible=True, value=card), gr.update(visible=True), gr.Gallery(selected_index=None), f"Custom: {path}", existing_item_index, trigger_word
+        except Exception as e:
+            gr.Warning(f"Invalid LoRA: either you entered an invalid link, or a non-Qwen compatible LoRA")
+            return gr.update(visible=True, value=f"Invalid LoRA: either you entered an invalid link, a non-Qwen compatible LoRA"), gr.update(visible=False), gr.update(), "", None, ""
+    else:
+        return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""
+
+def remove_custom_lora():
+    return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""
+
+run_lora.zerogpu = True
 
 css = '''
-.gradio-container {
-    max-width: 590px !important;
-    margin: 0 auto !important;
-}
-h1 {
-    text-align: center;
-}
-footer {
-    visibility: hidden;
-}
+#gen_btn{height: 100%}
+#gen_column{align-self: stretch}
+#title{text-align: center}
+#title h1{font-size: 3em; display:inline-flex; align-items:center}
+#title img{width: 100px; margin-right: 0.5em}
+#gallery .grid-wrap{height: 10vh}
+#lora_list{background: var(--block-background-fill);padding: 0 1em .3em; font-size: 90%}
+.card_internal{display: flex;height: 100px;margin-top: .5em}
+.card_internal img{margin-right: 1em}
+.styler{--form-gap-width: 0px !important}
+#progress{height:30px}
+#progress .generating{display:none}
+.progress-container {width: 100%;height: 30px;background-color: #f0f0f0;border-radius: 15px;overflow: hidden;margin-bottom: 20px}
+.progress-bar {height: 100%;background-color: #4f46e5;width: calc(var(--current) / var(--total) * 100%);transition: width 0.5s ease-in-out}
 '''
 
-# Gradio interface
-with gr.Blocks(css=css, theme="bethecloud/storj_theme", delete_cache=(240, 240)) as demo:
-    gr.Markdown(DESCRIPTION)
-    with gr.Row():
-        prompt = gr.Text(
-            label="Prompt",
-            show_label=False,
-            max_lines=1,
-            placeholder="✦︎ Enter your prompt",
-            container=False,
-        )
-        run_button = gr.Button("Run", scale=0, variant="primary")
-    result = gr.Gallery(label="Result", columns=1, show_label=False, preview=True)
+with gr.Blocks(theme="bethecloud/storj_theme", css=css, delete_cache=(120, 120)) as app:
+    title = gr.HTML("""<h1>Qwen Image LoRA DLC🥳</h1>""", elem_id="title",)
+    selected_index = gr.State(None)
     
     with gr.Row():
-        aspect_ratio = gr.Dropdown(
-            label="Aspect Ratio",
-            choices=list(aspect_ratios.keys()),
-            value="1:1",
-        )
-        lora = gr.Textbox(label="qwen image lora (optional)", placeholder="enter the path...")
-    with gr.Accordion("Additional Options", open=False):
-        use_negative_prompt = gr.Checkbox(
-            label="Use negative prompt",
-            value=True,
-            visible=True
-        )
-        negative_prompt = gr.Text(
-            label="Negative prompt",
-            max_lines=1,
-            placeholder="Enter a negative prompt",
-            value="text, watermark, copyright, blurry, low resolution",
-            visible=True,
-        )
-        seed = gr.Slider(
-            label="Seed",
-            minimum=0,
-            maximum=MAX_SEED,
-            step=1,
-            value=0,
-        )
-        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-        with gr.Row():
-            width = gr.Slider(
-                label="Width",
-                minimum=512,
-                maximum=2048,
-                step=64,
-                value=1024,
-            )
-            height = gr.Slider(
-                label="Height",
-                minimum=512,
-                maximum=2048,
-                step=64,
-                value=1024,
-            )
-        guidance_scale = gr.Slider(
-            label="Guidance Scale",
-            minimum=0.0,
-            maximum=20.0,
-            step=0.1,
-            value=4.0,
-        )
-        num_inference_steps = gr.Slider(
-            label="Number of inference steps",
-            minimum=1,
-            maximum=100,
-            step=1,
-            value=50,
-        )
-        num_images = gr.Slider(
-            label="Number of images",
-            minimum=1,
-            maximum=5,
-            step=1,
-            value=1,
-        )
-        zip_images = gr.Checkbox(label="Zip generated images", value=False)
-        with gr.Row(): 
-            lora_scale = gr.Slider(
-                label="LoRA Scale",
-                minimum=0,
-                maximum=2,
-                step=0.01,
-                value=1,
+        with gr.Column(scale=3):
+            prompt = gr.Textbox(label="Prompt", lines=1, placeholder="✦︎ Choose the LoRA and type the prompt")
+        with gr.Column(scale=1, elem_id="gen_column"):
+            generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
+
+    with gr.Row():
+        with gr.Column():
+            selected_info = gr.Markdown("")
+            gallery = gr.Gallery(
+                [(item["image"], item["title"]) for item in loras],
+                label="Qwen LoRA Collection",
+                allow_preview=False,
+                columns=3,
+                elem_id="gallery",
+                show_share_button=False
             )
-        
-        gr.Markdown("### Output Information")
-        seed_display = gr.Textbox(label="Seed used", interactive=False)
-        generation_time = gr.Textbox(label="Generation time (seconds)", interactive=False)
-        zip_file = gr.File(label="Download ZIP")
-
-    # Update aspect ratio
-    def set_dimensions(ar):
-        w, h = aspect_ratios[ar]
-        return gr.update(value=w), gr.update(value=h)
-    
+            
+            with gr.Group():
+                custom_lora = gr.Textbox(label="Enter Custom Qwen LoRA", placeholder="prithivMLmods/Qwen-Image-Sketch-Smudge")
+                gr.Markdown("[Check the list of Qwen-compatible LoRAs](https://huggingface.co/models?search=qwen+lora)", elem_id="lora_list")
+            
+            custom_lora_info = gr.HTML(visible=False)
+            custom_lora_button = gr.Button("Remove custom LoRA", visible=False)
+            
+        with gr.Column():
+            progress_bar = gr.Markdown(elem_id="progress", visible=False)
+            result = gr.Image(label="Generated Image", format="png")
+
+            with gr.Row():
+                aspect_ratio = gr.Dropdown(
+                    label="Aspect Ratio",
+                    choices=list(aspect_ratios.keys()),
+                    value="1:1",
+                )
+
+    with gr.Row():
+        with gr.Accordion("Advanced Settings", open=False):
+                
+            with gr.Row():
+                use_negative_prompt = gr.Checkbox(
+                    label="Use negative prompt", value=True, visible=True
+                )
+                negative_prompt = gr.Text(
+                    label="Negative prompt",
+                    max_lines=1,
+                    placeholder="Enter a negative prompt",
+                    value="text, watermark, copyright, blurry, low resolution",
+                    visible=True,
+                )
+            
+            with gr.Column():
+                with gr.Row():
+                    cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=4.0)
+                    steps = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=50)
+                
+                with gr.Row():
+                    width = gr.Slider(label="Width", minimum=256, maximum=2048, step=64, value=1024)
+                    height = gr.Slider(label="Height", minimum=256, maximum=2048, step=64, value=1024)
+                
+                with gr.Row():
+                    randomize_seed = gr.Checkbox(True, label="Randomize seed")
+                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
+                    lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=2, step=0.01, value=1.0)
+
+    # Event handlers
+    gallery.select(
+        update_selection,
+        inputs=[width, height],
+        outputs=[prompt, selected_info, selected_index, width, height]
+    )
+
     aspect_ratio.change(
         fn=set_dimensions,
         inputs=aspect_ratio,
         outputs=[width, height]
     )
 
-    # Negative prompt visibility
     use_negative_prompt.change(
         fn=lambda x: gr.update(visible=x),
         inputs=use_negative_prompt,
         outputs=negative_prompt
     )
 
-    # Run button and prompt submit
-    gr.on(
-        triggers=[prompt.submit, run_button.click],
-        fn=generate,
-        inputs=[
-            prompt,
-            negative_prompt,
-            use_negative_prompt,
-            seed,
-            width,
-            height,
-            guidance_scale,
-            randomize_seed,
-            num_inference_steps,
-            num_images,
-            zip_images,
-            lora,
-            lora_scale,
-        ],
-        outputs=[result, seed_display, generation_time, zip_file],
-        api_name="run",
+    custom_lora.input(
+        add_custom_lora,
+        inputs=[custom_lora],
+        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, prompt]
     )
 
-    # Examples
-    gr.Examples(
-        examples=examples,
-        inputs=prompt,
-        outputs=[result, seed_display, generation_time, zip_file],
-        fn=generate,
-        cache_examples=False,
+    custom_lora_button.click(
+        remove_custom_lora,
+        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, custom_lora]
+    )
+
+    gr.on(
+        triggers=[generate_button.click, prompt.submit],
+        fn=run_lora,
+        inputs=[prompt, negative_prompt, use_negative_prompt, aspect_ratio, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale],
+        outputs=[result, seed, progress_bar]
     )
 
-if __name__ == "__main__":
-    demo.queue(max_size=50).launch(share=False, mcp_server=True, ssr_mode=False, debug=True, show_error=True)
+app.queue()
+app.launch(share=False, ssr_mode=False, show_error=True)