Refactor: Improve error handling and add image loading from URLs in app.py

app.py

@@ -1,8 +1,6 @@
 import gradio as gr
 import torch
 from PIL import Image
-from io import BytesIO
-import requests
 from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
 
 # 加载模型和ControlNet
@@ -12,43 +10,35 @@ controlnet_id = "lllyasviel/sd-controlnet-canny"
 pipe = StableDiffusionControlNetPipeline.from_pretrained(model_id, controlnet=ControlNetModel.from_pretrained(controlnet_id))
 pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
-def load_image_from_url(url):
-    try:
-        response = requests.get(url)
-        response.raise_for_status()
-        return Image.open(BytesIO(response.content)).convert("RGB")
-    except Exception as e:
-        return None
+def is_valid_pil_image(image):
+    return isinstance(image, Image.Image)
 
-def generate_image(control_image_path, additional_prompt, reference_image1_path, reference_image2_path):
-    control_image = load_image_from_url(control_image_path["path"])
-    reference_image1 = load_image_from_url(reference_image1_path["path"])
-    reference_image2 = load_image_from_url(reference_image2_path["path"])
-    
-    if control_image is None or reference_image1 is None or reference_image2 is None:
-        return "Error: One or more input images could not be loaded."
+def generate_image(control_image, additional_prompt, reference_image1, reference_image2):
+    # 检查输入图像是否为有效的 PIL 图像
+    if not all(is_valid_pil_image(img) for img in [control_image, reference_image1, reference_image2]):
+        return "Error: One or more input images are not valid PIL images."
 
     try:
         # 第一步：使用 canny 模型生成初始图像
         initial_image = pipe(prompt=additional_prompt, control_image=control_image).images[0]
-        if initial_image is None:
-            return "Error in Step 1: Generated image is None."
+        if not is_valid_pil_image(initial_image):
+            return "Error in Step 1: Generated image is not a valid PIL image."
     except Exception as e:
         return f"Error in Step 1: {str(e)}"
     
     try:
         # 第二步：使用第一张参考图像进行 reference_adain+attn 生成
         step2_image = pipe(prompt=additional_prompt, control_image=reference_image1).images[0]
-        if step2_image is None:
-            return "Error in Step 2: Generated image is None."
+        if not is_valid_pil_image(step2_image):
+            return "Error in Step 2: Generated image is not a valid PIL image."
     except Exception as e:
         return f"Error in Step 2: {str(e)}"
     
     try:
         # 第三步：使用第二张参考图像进行 reference_adain+attn 生成
         final_image = pipe(prompt=additional_prompt, control_image=reference_image2).images[0]
-        if final_image is None:
-            return "Error in Step 3: Generated image is None."
+        if not is_valid_pil_image(final_image):
+            return "Error in Step 3: Generated image is not a valid PIL image."
         return final_image
     except Exception as e:
         return f"Error in Step 3: {str(e)}"
@@ -57,10 +47,10 @@ def generate_image(control_image_path, additional_prompt, reference_image1_path,
 interface = gr.Interface(
     fn=generate_image,
     inputs=[
-        gr.components.Textbox(label="Control Image URL"),
+        gr.components.Image(type="pil", label="Control Image"),
         gr.components.Textbox(label="Additional Prompt"),
-        gr.components.Textbox(label="Reference Image 1 URL"),
-        gr.components.Textbox(label="Reference Image 2 URL")
+        gr.components.Image(type="pil", label="Reference Image 1"),
+        gr.components.Image(type="pil", label="Reference Image 2")
     ],
     outputs=gr.components.Image(type="pil", label="Generated Image"),
     title="Stable Diffusion with Multi-Step ControlNet",