Improve image generation and display

app.py

@@ -4,6 +4,8 @@ import random
 import logging
 import sys
 import os
+from PIL import Image as PILImage
+import io
 
 # 设置日志记录
 logging.basicConfig(level=logging.INFO, 
@@ -54,6 +56,12 @@ try:
 except Exception as e:
     logger.error(f"Failed to patch Gradio: {str(e)}")
 
+# 创建一个简单的示例图像，在模型加载失败或生成失败时使用
+def create_dummy_image():
+    # 创建一个256x256的红色图像
+    img = PILImage.new('RGB', (256, 256), color = (255, 0, 0))
+    return img
+
 # 加载模型
 try:
     from diffusers import DiffusionPipeline
@@ -85,7 +93,11 @@ def generate_image(prompt):
     try:
         if pipe is None:
             logger.error("Model not loaded, cannot generate image")
-            return None
+            return create_dummy_image()
+            
+        if not prompt or prompt.strip() == "":
+            prompt = "A beautiful landscape"
+            logger.info(f"Empty prompt, using default: {prompt}")
             
         logger.info(f"Generating image for prompt: {prompt}")
         seed = random.randint(0, MAX_SEED)
@@ -99,26 +111,60 @@ def generate_image(prompt):
                 num_inference_steps=2,
                 generator=generator
             ).images[0]
+            
+            # 确保图像是 PIL.Image 类型
+            if not isinstance(image, PILImage.Image):
+                logger.warning(f"Converting image from {type(image)} to PIL.Image")
+                if hasattr(image, 'numpy'):
+                    image = PILImage.fromarray(image.numpy())
+                else:
+                    image = PILImage.fromarray(np.array(image))
+                    
+            # 转换为 RGB 模式，确保兼容性
+            if image.mode != 'RGB':
+                image = image.convert('RGB')
+                
+            logger.info("Image generation successful")
+            # 保存图像以供调试
+            debug_path = "debug_image.jpg"
+            image.save(debug_path)
+            logger.info(f"Debug image saved to {debug_path}")
+            
+            return image
         except IndexError as e:
             logger.error(f"Index error in pipe: {str(e)}")
-            return None
+            return create_dummy_image()
         
-        logger.info("Image generation successful")
-        return image
     except Exception as e:
         logger.error(f"Error generating image: {str(e)}")
-        return None
+        return create_dummy_image()
 
 # 创建简单的 Gradio 界面，禁用示例缓存
-demo = gr.Interface(
-    fn=generate_image,
-    inputs=gr.Textbox(label="Enter your prompt"),
-    outputs=gr.Image(label="Generated Image"),
-    title="SDXL Turbo Text-to-Image",
-    description="Enter a text prompt to generate an image.",
-    examples=["A cute cat", "Sunset over mountains"],  # 减少示例数量
-    cache_examples=False  # 禁用示例缓存以避免文件访问错误
-)
+with gr.Blocks(title="SDXL Turbo Text-to-Image") as demo:
+    gr.Markdown("# SDXL Turbo Text-to-Image Generator")
+    
+    with gr.Row():
+        with gr.Column():
+            prompt_input = gr.Textbox(label="Enter your prompt", placeholder="Type your text prompt here...")
+            generate_button = gr.Button("Generate Image")
+        
+        with gr.Column():
+            image_output = gr.Image(label="Generated Image", type="pil")
+    
+    # 示例
+    gr.Examples(
+        examples=["A cute cat", "Sunset over mountains"],
+        inputs=prompt_input,
+        outputs=image_output,
+        fn=generate_image,
+        cache_examples=False
+    )
+    
+    # 绑定生成按钮
+    generate_button.click(fn=generate_image, inputs=prompt_input, outputs=image_output)
+    
+    # 直接绑定文本框的提交事件
+    prompt_input.submit(fn=generate_image, inputs=prompt_input, outputs=image_output)
 
 # 启动应用
 if __name__ == "__main__":
@@ -126,9 +172,10 @@ if __name__ == "__main__":
         logger.info("Starting Gradio app")
         # 添加更多启动选项，以提高稳定性
         demo.launch(
+            debug=True,
             show_error=True,
-            max_threads=1,  # 降低并发以避免竞态条件
-            share=False
+            server_name="0.0.0.0",  # 确保可以从外部访问
+            max_threads=1  # 降低并发以避免竞态条件
         )
     except Exception as e:
         logger.error(f"Error launching app: {str(e)}")