Update app.py

app.py

@@ -438,7 +438,6 @@ 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("Debug: Starting image generation")
     pipe.to("cuda")
     generator = torch.Generator(device="cuda").manual_seed(seed)
     final_image = None  # Initialize final_image
@@ -456,6 +455,9 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
             good_vae=good_vae,
         ):
             print(f"Debug: Yielding image of type {type(img)}")  # Debugging
+            if not isinstance(img, (PIL.Image.Image, torch.Tensor)):  # Check if it's an image
+                print(f"Error: Expected an image, but got {type(img)}")
+                raise ValueError("Expected an image, but got a non-image value.")
             yield img
             final_image = img  # Update final_image with the current image
         return final_image
@@ -531,29 +533,29 @@ def run_lora(prompt, cfg_scale, steps, selected_info_1, selected_info_2, selecte
     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)
-            print("Debug: generate_image_to_image output type:", type(final_image))  # Debugging
-            print("Debug: generate_image_to_image output value:", final_image)  # Debugging
+            print(f"Debug: generate_image_to_image output type: {type(final_image)}")  # Debugging
+            if not isinstance(final_image, (PIL.Image.Image, torch.Tensor)):  # Check if it's an image
+                print(f"Error: Expected an image, but got {type(final_image)}")
+                raise ValueError("Expected an image from generate_image_to_image, but got a non-image value.")
+            print(f"Debug: generate_image_to_image output value: {final_image}")  # Debugging
             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:
-                print("Debug: generate_image yielded type:", type(image))  # Debugging
-                print("Debug: generate_image yielded value:", image)  # Debugging
+                print(f"Debug: generate_image yielded type: {type(image)}")  # Debugging
+                print(f"Debug: generate_image yielded value: {image}")  # Debugging
                 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
             else:
-                print("Debug: final_image type:", type(final_image))  # Debugging
-                print("Debug: final_image value:", final_image)  # Debugging
-                
-            yield final_image, seed, gr.update(value=progress_bar, visible=False)
+                print(f"Debug: final_image type: {type(final_image)}")  # Debugging
+                yield final_image, seed, gr.update(value=progress_bar, visible=False)
 
     except Exception as e:
         print(f"Error during image generation: {e}")  # Error handling