Update app.py

app.py

@@ -2,6 +2,8 @@ import gradio as gr
 import torch
 from diffusers import DiffusionPipeline
 from PIL import Image
+import numpy as np
+from torchvision import transforms
 
 # Set up device
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -17,9 +19,12 @@ def process_image(prompt, image, style, upscale_factor, inpaint):
     try:
         # Ensure the image is in the correct format
         if image is not None:
-            image = image.convert("RGB")
-        else:
-            image = None
+            if isinstance(image, np.ndarray):
+                image = Image.fromarray(image)
+            elif isinstance(image, torch.Tensor):
+                image = transforms.ToPILImage()(image)
+            elif not isinstance(image, Image.Image):
+                return None, "Error: Unsupported image format."
 
         # Log the input parameters
         print(f"Prompt: {prompt}")
@@ -49,7 +54,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             prompt_input = gr.Textbox(label="Enter your prompt")
-            image_input = gr.Image(label="Image (for inpainting)", type="pil")  # Removed 'optional=True'
+            image_input = gr.Image(label="Image (for inpainting)", type="pil")
             style_input = gr.Dropdown(choices=["Fooocus Style", "SAI Anime"], label="Select Style")
             upscale_input = gr.Slider(minimum=1, maximum=4, step=1, label="Upscale Factor")
             inpaint_input = gr.Checkbox(label="Enable Inpainting")