Update app.py

app.py

@@ -24,80 +24,30 @@ import shutil
 import uuid
 import zipfile
 
-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)
     return unique_name
 
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        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
+dtype = torch.bfloat16
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load Qwen model
+pipe = DiffusionPipeline.from_pretrained("Qwen/Qwen-Image", torch_dtype=dtype).to(device)
+
+# Aspect ratios
+aspect_ratios = {
+    "1:1": (1328, 1328),
+    "16:9": (1664, 928),
+    "9:16": (928, 1664),
+    "4:3": (1472, 1140),
+    "3:4": (1140, 1472)
+}
 
 loras = [
     # Sample Qwen-compatible LoRAs
@@ -138,56 +88,16 @@ loras = [
     },
 ]
 
-#--------------------------------------------------Model Initialization-----------------------------------------------------------------------------------------#
-dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
-base_model = "Qwen/Qwen-Image"
-
-# Load Qwen Image pipeline
-pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype).to(device)
-
-# Add aspect ratios for Qwen
-aspect_ratios = {
-    "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:
         return
-
+    
     # If it's just an ID like "author/model"
     if "/" in lora_input and not lora_input.startswith("http"):
         pipe.load_lora_weights(lora_input, adapter_name="default")
         return
-
+    
     if lora_input.startswith("http"):
         url = lora_input
         # Repo page (no blob/resolve)
@@ -195,11 +105,11 @@ def load_lora_opt(pipe, lora_input):
             repo_id = urlparse(url).path.strip("/")
             pipe.load_lora_weights(repo_id, adapter_name="default")
             return
-
+        
         # Blob link → convert to resolve link
         if "/blob/" in url:
             url = url.replace("/blob/", "/resolve/")
-
+        
         # Download direct file
         tmp_dir = tempfile.mkdtemp()
         local_path = os.path.join(tmp_dir, os.path.basename(urlparse(url).path))
@@ -215,148 +125,47 @@ def load_lora_opt(pipe, lora_input):
         finally:
             shutil.rmtree(tmp_dir, ignore_errors=True)
 
-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}) ✅"
-    
-    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
-    
-    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)
-    
-    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)
-
-@spaces.GPU(duration=120)
-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 ""
-
-    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]}"
+            response = requests.get(f"https://huggingface.co/api/models/{link}")
+            response.raise_for_status()
+            model_info = response.json()
+            
+            # Check if it's a Qwen model
+            if "qwen" not in model_info.get("tags", []):
+                raise Exception("Not a Qwen LoRA model!")
+                
+            # Get image if available
+            image_url = None
+            if "cardData" in model_info and "widget" in model_info["cardData"]:
+                if len(model_info["cardData"]["widget"]) > 0:
+                    image_url = model_info["cardData"]["widget"][0].get("output", {}).get("url", None)
+            
+            # Try to find safetensors file
+            safetensors_name = None
+            try:
+                model_files = requests.get(f"https://huggingface.co/api/models/{link}/tree/main").json()
+                for file in model_files:
+                    if file.get("path", "").endswith(".safetensors"):
+                        safetensors_name = file["path"]
+                        break
+            except:
+                pass
+                
+            return split_link[1], link, safetensors_name, "trigger_word", image_url
         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
+            print(f"Error getting model info: {e}")
+            raise Exception(f"Failed to get model info: {e}")
+    return None, None, None, None, None
 
 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: 
+    else:
         return get_huggingface_safetensors(link)
 
 def add_custom_lora(custom_lora):
@@ -364,6 +173,9 @@ def add_custom_lora(custom_lora):
     if custom_lora:
         try:
             title, repo, path, trigger_word, image = check_custom_model(custom_lora)
+            if not title:
+                raise Exception("Invalid LoRA model")
+                
             print(f"Loaded custom LoRA: {repo}")
             card = f'''
             <div class="custom_lora_card">
@@ -377,6 +189,7 @@ def add_custom_lora(custom_lora):
               </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 = {
@@ -386,21 +199,161 @@ def add_custom_lora(custom_lora):
                     "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, ""
+            gr.Warning(f"Invalid LoRA: either you entered an invalid link, or a non-Qwen LoRA")
+            return gr.update(visible=True, value=f"Invalid LoRA: either you entered an invalid link, a non-Qwen 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
+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}) ✅"
+    
+    # Update aspect ratio based on LoRA if it has aspect info
+    if "aspect" in selected_lora:
+        if selected_lora["aspect"] == "portrait":
+            width = 928
+            height = 1664
+        elif selected_lora["aspect"] == "landscape":
+            width = 1664
+            height = 928
+        else:
+            width = 1328
+            height = 1328
+    
+    return (
+        gr.update(placeholder=new_placeholder),
+        updated_text,
+        evt.index,
+        width,
+        height,
+    )
+
+@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()
+    
+    # Clear any existing LoRA adapters
+    current_adapters = pipe.get_list_adapters()
+    for adapter in current_adapters:
+        pipe.delete_adapters(adapter)
+    pipe.disable_lora()
+    
+    use_lora = False
+    if lora_input and lora_input.strip() != "":
+        load_lora_opt(pipe, lora_input)
+        pipe.set_adapters(["default"], adapter_weights=[lora_scale])
+        use_lora = True
+    
+    images = pipe(
+        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
+    
+    end_time = time.time()
+    duration = end_time - start_time
+    
+    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.get_list_adapters()
+    for adapter in current_adapters:
+        pipe.delete_adapters(adapter)
+    pipe.disable_lora()
+    
+    return image_paths, seed, f"{duration:.2f}", zip_path
+
+@spaces.GPU(duration=120)
+def run_lora(
+    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,
+    selected_index: int,
+    lora_scale: float,
+    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_repo = selected_lora["repo"]
+    trigger_word = selected_lora["trigger_word"]
+    
+    if trigger_word:
+        prompt_mash = f"{trigger_word} {prompt}"
+    else:
+        prompt_mash = prompt
+    
+    final_negative_prompt = negative_prompt if use_negative_prompt else ""
+    
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    
+    return generate_qwen(
+        prompt=prompt_mash,
+        negative_prompt=final_negative_prompt,
+        seed=seed,
+        width=width,
+        height=height,
+        guidance_scale=guidance_scale,
+        randomize_seed=False,  # Already handled
+        num_inference_steps=num_inference_steps,
+        num_images=num_images,
+        zip_images=zip_images,
+        lora_input=lora_repo,
+        lora_scale=lora_scale,
+        progress=progress,
+    )
 
 css = '''
 #gen_btn{height: 100%}
@@ -420,7 +373,7 @@ css = '''
 '''
 
 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",)
+    title = gr.HTML("""<h1>Qwen Image LoRA DLC ❤️‍🔥</h1>""", elem_id="title")
     selected_index = gr.State(None)
     
     with gr.Row():
@@ -428,101 +381,127 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css, delete_cache=(120, 120))
             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",
+                label="Qwen LoRA DLC's",
                 allow_preview=False,
                 columns=3,
                 elem_id="gallery",
                 show_share_button=False
             )
-            
             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 = gr.Textbox(label="Enter Custom LoRA", placeholder="prithivMLmods/Qwen-Image-Sketch-Smudge")
+                gr.Markdown("[Check the list of Qwen LoRA's](https://huggingface.co/models?other=base_model:adapter:Qwen/Qwen-Image)", 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")
-
+            result = gr.Gallery(label="Generated Images", columns=1, show_label=False, preview=True)
             with gr.Row():
                 aspect_ratio = gr.Dropdown(
                     label="Aspect Ratio",
                     choices=list(aspect_ratios.keys()),
                     value="1:1",
                 )
+            with gr.Row():
+                steps = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=48)
 
     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
+                    label="Use negative prompt",
+                    value=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]
-    )
-
+            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=512, maximum=2048, step=64, value=1328)
+                height = gr.Slider(label="Height", minimum=512, maximum=2048, step=64, value=1328)
+            
+            with gr.Row():
+                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():
+                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)
+    
+    # Output information
+    with gr.Row():
+        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)
+    
     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
     )
-
+    
+    gallery.select(
+        update_selection,
+        inputs=[width, height],
+        outputs=[prompt, selected_info, selected_index, width, height]
+    )
+    
     custom_lora.input(
         add_custom_lora,
         inputs=[custom_lora],
         outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, prompt]
     )
-
+    
     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]
+        inputs=[
+            prompt,
+            negative_prompt,
+            use_negative_prompt,
+            seed,
+            width,
+            height,
+           #guidance_scale,
+            randomize_seed,
+            steps,
+            num_images,
+            zip_images,
+            selected_index,
+            lora_scale,
+        ],
+        outputs=[result, seed_display, generation_time, zip_file]
     )
 
 app.queue()