Update app.py

app.py

@@ -298,7 +298,7 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
     pipe.to("cuda")
     generator = torch.Generator(device="cuda").manual_seed(seed)
     with calculateDuration("Generating image"):
-        # Generate image
+        # Generate images
         for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
             prompt=prompt_mash,
             num_inference_steps=steps,
@@ -311,7 +311,6 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
             good_vae=good_vae,
         ):
             yield img
-    return img
 
 @spaces.GPU(duration=75)
 def run_lora(prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state, progress=gr.Progress(track_tqdm=True)):
@@ -332,11 +331,11 @@ def run_lora(prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scal
                 appends.append(trigger_word)
     prompt_mash = " ".join(prepends + [prompt] + appends)
     print("Prompt Mash: ", prompt_mash)
+
     # Unload previous LoRA weights
     with calculateDuration("Unloading LoRA"):
         pipe.unload_lora_weights()
 
-        
     print(pipe.get_active_adapters())
 
     lora_names = []
@@ -369,13 +368,13 @@ def run_lora(prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scal
         image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
 
         step_counter = 0
+        images = []
         for image in image_generator:
             step_counter += 1
-            final_image = image
+            images.append(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 images, seed, gr.update(value=progress_bar, visible=True)
+            
 run_lora.zerogpu = True
 
 def get_huggingface_safetensors(link):
@@ -563,13 +562,8 @@ with gr.Blocks(css=css, delete_cache=(240, 240)) as app:
                 interactive=False
             )
         with gr.Column():
-            progress_bar = gr.Markdown(elem_id="progress", visible=False)
-            result = gr.Image(label="Image 1", interactive=False, show_share_button=False)
-            result = gr.Image(label="Image 2", interactive=False, show_share_button=False)
-            result = gr.Image(label="Image 3", interactive=False, show_share_button=False)
-            result = gr.Image(label="Image 4", interactive=False, show_share_button=False)
-#            with gr.Accordion("History", open=False):
-#                history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", interactive=False)  
+    progress_bar = gr.Markdown(elem_id="progress", visible=False)
+    results = [gr.Image(label=f"Image {i+1}", interactive=False, show_share_button=False) for i in range(4)]
 
     gallery.select(
         update_selection,