Update app.py

app.py

@@ -454,6 +454,7 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
             output_type="pil",
             good_vae=good_vae,
         ):
+            # Debugging: Check image type before yielding
             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)}")
@@ -462,6 +463,7 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
             final_image = img  # Update final_image with the current image
         return final_image
 
+
 @spaces.GPU(duration=75)
 def run_lora(prompt, cfg_scale, steps, selected_info_1, selected_info_2, selected_info_3, selected_info_4, selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, lora_scale_4, randomize_seed, seed, width, height, loras_state, image_input=None, progress=gr.Progress(track_tqdm=True)):
     print("run_lora function called.")  # Debugging statement
@@ -551,12 +553,27 @@ def run_lora(prompt, cfg_scale, steps, selected_info_1, selected_info_2, selecte
                 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(f"Debug: final_image type: {type(final_image)}")  # Debugging
-                yield final_image, seed, gr.update(value=progress_bar, visible=False)
+            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(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)
+        
+        if final_image is None:
+            print("No final image generated.")  # Debugging statement
+        else:
+            print(f"Debug: final_image 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, but got a non-image value.")
+        print(f"Debug: final_image value: {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
         raise gr.Error("An error occurred during image generation.")