Refactor: Improve error handling and add logging in app.py

.gitignore

@@ -0,0 +1 @@
+flagged

app.py

@@ -2,6 +2,10 @@ import gradio as gr
 import torch
 from PIL import Image
 from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+import logging
+
+# 配置日志记录
+logging.basicConfig(level=logging.DEBUG)
 
 # 加载模型和ControlNet
 model_id = "CompVis/stable-diffusion-v1-4"
@@ -14,34 +18,55 @@ def is_valid_pil_image(image):
     return isinstance(image, Image.Image)
 
 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:
+        logging.debug("Starting image generation process...")
+        
+        # 打印传入的图像对象类型
+        logging.debug(f"Control Image type: {type(control_image)}")
+        logging.debug(f"Reference Image 1 type: {type(reference_image1)}")
+        logging.debug(f"Reference Image 2 type: {type(reference_image2)}")
+        
+        # 检查输入图像是否为有效的 PIL 图像
+        if not all(is_valid_pil_image(img) for img in [control_image, reference_image1, reference_image2]):
+            error_message = "Error: One or more input images are not valid PIL images."
+            logging.error(error_message)
+            return error_message
+
         # 第一步：使用 canny 模型生成初始图像
         initial_image = pipe(prompt=additional_prompt, control_image=control_image).images[0]
         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 生成
+            error_message = "Error in Step 1: Generated image is not a valid PIL image."
+            logging.error(error_message)
+            return error_message
+
+        # 第二步：使用第一张参考图像生成图像
         step2_image = pipe(prompt=additional_prompt, control_image=reference_image1).images[0]
+        if step2_image is None:
+            error_message = "Error in Step 2: Generated image is None (NoneType)."
+            logging.error(error_message)
+            return error_message
         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 生成
+            error_message = "Error in Step 2: Generated image is not a valid PIL image."
+            logging.error(error_message)
+            return error_message
+
+        # 第三步：使用第二张参考图像生成图像
         final_image = pipe(prompt=additional_prompt, control_image=reference_image2).images[0]
+        if final_image is None:
+            error_message = "Error in Step 3: Generated image is None (NoneType)."
+            logging.error(error_message)
+            return error_message
         if not is_valid_pil_image(final_image):
-            return "Error in Step 3: Generated image is not a valid PIL image."
+            error_message = "Error in Step 3: Generated image is not a valid PIL image."
+            logging.error(error_message)
+            return error_message
+
+        logging.debug("Image generation process completed successfully.")
         return final_image
     except Exception as e:
-        return f"Error in Step 3: {str(e)}"
+        error_message = f"Error: {str(e)}"
+        logging.error(error_message)
+        return error_message
 
 # 定义Gradio接口
 interface = gr.Interface(