Update app.py

app.py

@@ -468,6 +468,7 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
             output_type="pil",
             good_vae=good_vae,
         ):
+            print("Image generated successfully.")  # Debugging statement
             yield img
 
 @spaces.GPU(duration=75)
@@ -488,13 +489,15 @@ def run_lora(prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scal
             else:
                 appends.append(trigger_word)
     prompt_mash = " ".join(prepends + [prompt] + appends)
-    print("Prompt Mash: ", prompt_mash)
+    print("Prompt Mash: ", prompt_mash)  # Debugging statement
+
     # Unload previous LoRA weights
     with calculateDuration("Unloading LoRA"):
         pipe.unload_lora_weights()
         pipe_i2i.unload_lora_weights()
         
     print(pipe.get_active_adapters())
+    
     # Load LoRA weights with respective scales
     lora_names = []
     lora_weights = []
@@ -514,15 +517,15 @@ def run_lora(prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scal
                 low_cpu_mem_usage=True,
                 adapter_name=lora_name
             )
-#            if image_input is not None: pipe_i2i = pipe_to_use
-#            else: pipe = pipe_to_use
         print("Loaded LoRAs:", lora_names)
         print("Adapter weights:", lora_weights)
         if image_input is not None:
             pipe_i2i.set_adapters(lora_names, adapter_weights=lora_weights)
         else:
             pipe.set_adapters(lora_names, adapter_weights=lora_weights)
+    
     print(pipe.get_active_adapters())
+    
     # Set random seed for reproducibility
     with calculateDuration("Randomizing seed"):
         if randomize_seed:
@@ -530,20 +533,27 @@ def run_lora(prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scal
         print("Image Seed:", seed)
 
     # Generate image
-    if image_input is not None:
-        #final_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, seed)
-        yield final_image, seed, gr.update(visible=False)
-    else:
-        image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
-        # Consume the generator to get the final image
-        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)
+    try:
+        if image_input is not None:
+            #final_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, seed)
+            yield final_image, seed, gr.update(visible=False)
+        else:
+            image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
+            # Consume the generator to get the final image
+            final_image = None
+            step_counter = 0
+            for image in image_generator:
+                step_counter += 1
+                final_image = image
+                print(f"Yielding image {step_counter}/{steps}")  # Debugging statement
+                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)
+            if final_image is None:
+                print("No final image generated.")  # Debugging statement
+            yield final_image, seed, gr.update(value=progress_bar, visible=False)
+    except Exception as e:
+        print(f"Error during image generation: {e}")  # Error handling
+        raise gr.Error("An error occurred during image generation.")
 
 run_lora.zerogpu = True