Update app.py

app.py

@@ -294,23 +294,38 @@ def remove_custom_lora(selected_indices, current_loras, gallery):
     )
 
 def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress):
-    print("Generating image...")
+    print("Generating images...")
     pipe.to("cuda")
     generator = torch.Generator(device="cuda").manual_seed(seed)
-    with calculateDuration("Generating image"):
-        # Generate images
-        for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
-            prompt=prompt_mash,
-            num_inference_steps=steps,
-            guidance_scale=cfg_scale,
-            width=width,
-            height=height,
-            generator=generator,
-            joint_attention_kwargs={"scale": 1.0},
-            output_type="pil",
-            good_vae=good_vae,
-        ):
-            yield img
+    images = []
+
+    with calculateDuration("Generating images"):
+        try:
+            for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
+                prompt=prompt_mash,
+                num_inference_steps=steps,
+                guidance_scale=cfg_scale,
+                width=width,
+                height=height,
+                generator=generator,
+                joint_attention_kwargs={"scale": 1.0},
+                output_type="pil",
+                good_vae=good_vae,
+            ):
+                images.append(img)
+                progress(percent=len(images) / 4 * 100)  # Adjust progress for 4 images
+                if len(images) == 4:  # Collect exactly 4 images
+                    break
+        except Exception as e:
+            print(f"Error during image generation: {e}")
+            raise
+
+    if len(images) < 4:
+        print("Fewer than 4 images generated. Padding with None.")
+    while len(images) < 4:
+        images.append(None)
+
+    return images
 
 @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)):
@@ -365,22 +380,9 @@ def run_lora(prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scal
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
 
-        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
-            images.append(image)
-            progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'
-            if len(images) == 4:  # Ensure four images are collected
-                break
-
-        # Pad to ensure exactly 4 images are returned
-        while len(images) < 4:
-            images.append(None)
+        images = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
 
-        return *images, seed, gr.update(value=progress_bar, visible=True)
+        return *images[:4], seed, gr.update(value="", visible=False)
             
 run_lora.zerogpu = True