恢复AI图像生成功能，使用更轻量的方式

app.py

@@ -59,36 +59,128 @@ def create_dummy_image():
     img = PILImage.new('RGB', (256, 256), color = (255, 100, 100))
     return img
 
-# 使用极简方法
-def simple_demo():
-    # 定义一个极简单的生成函数
-    def generate(text):
-        logger.info(f"Received text: {text}")
-        # 创建一个简单的图像作为响应
-        image = PILImage.new('RGB', (256, 256), color=(100, 200, 100))
-        logger.info("Created simple green image")
-        return image
+# 全局变量
+pipe = None
+
+# 懒加载AI模型函数
+def get_model():
+    try:
+        import torch
+        from diffusers import StableDiffusionPipeline
+        
+        logger.info("开始加载模型...")
+        
+        # 使用较小的模型而不是SDXL
+        model_id = "runwayml/stable-diffusion-v1-5"
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        logger.info(f"使用设备: {device}")
+        
+        # 优化设置以减少内存使用
+        if torch.cuda.is_available():
+            # 使用半精度
+            pipe = StableDiffusionPipeline.from_pretrained(
+                model_id, 
+                torch_dtype=torch.float16,
+                safety_checker=None,  # 禁用安全检查器以节省内存
+                requires_safety_checker=False,
+                use_safetensors=True
+            )
+            pipe = pipe.to(device)
+            pipe.enable_attention_slicing()  # 减少显存使用
+            
+            # 释放不必要的内存
+            torch.cuda.empty_cache()
+        else:
+            # CPU版本，占用内存较大，但现在只用处理一个请求
+            pipe = StableDiffusionPipeline.from_pretrained(
+                model_id,
+                safety_checker=None,
+                requires_safety_checker=False,
+                use_safetensors=True
+            )
+            pipe = pipe.to(device)
+        
+        logger.info("模型加载成功")
+        return pipe
+    except Exception as e:
+        logger.error(f"模型加载失败: {e}")
+        return None
+
+# 生成图像函数
+def generate_image(prompt):
+    global pipe
+    
+    # 如果提示为空，使用默认提示
+    if not prompt or prompt.strip() == "":
+        prompt = "a beautiful landscape"
+        logger.info(f"输入为空，使用默认提示词: {prompt}")
     
-    logger.info("Setting up simple demo without model loading")
-    # 创建最小界面
+    logger.info(f"收到提示词: {prompt}")
+    
+    # 第一次调用时加载模型
+    if pipe is None:
+        pipe = get_model()
+        if pipe is None:
+            logger.error("模型加载失败，返回默认图像")
+            return create_dummy_image()
+    
+    try:
+        # 优化生成参数，减少内存需求
+        logger.info("开始生成图像...")
+        
+        # 设置随机种子以确保结果一致性
+        seed = random.randint(0, 2147483647)
+        generator = torch.Generator(device=pipe.device).manual_seed(seed)
+        
+        # 使用最轻量级的参数
+        image = pipe(
+            prompt=prompt,
+            num_inference_steps=3,  # 极少的步骤
+            guidance_scale=7.5,
+            height=256,  # 小尺寸
+            width=256,   # 小尺寸
+            generator=generator
+        ).images[0]
+        
+        # 释放缓存，避免内存增长
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            
+        logger.info(f"图像生成成功，种子: {seed}")
+        return image
+    except Exception as e:
+        logger.error(f"生成过程发生错误: {e}")
+        return create_dummy_image()
+
+# 创建Gradio界面
+def create_demo():
+    # 使用简单界面，避免复杂组件
     demo = gr.Interface(
-        fn=generate,
-        inputs=gr.Textbox(label="Prompt"),
-        outputs=gr.Image(type="pil", label="Generated Image"),
-        title="Simple Text-to-Image", 
-        examples=["a cat"],
-        cache_examples=False
+        fn=generate_image,
+        inputs=gr.Textbox(label="输入提示词"),
+        outputs=gr.Image(type="pil", label="生成的图像"),
+        title="文本到图像生成",
+        description="输入文本描述，AI将生成相应的图像（会加载较长时间，请耐心等待）",
+        examples=["a cute cat", "mountain landscape"],
+        cache_examples=False,
+        allow_flagging="never"  # 禁用标记功能以减少复杂性
     )
     return demo
 
-# 只使用极简应用，避免加载任何模型
-demo = simple_demo()
+# 创建演示界面
+demo = create_demo()
 
 # 启动应用
 if __name__ == "__main__":
     try:
-        logger.info("Starting Gradio interface")
+        logger.info("启动Gradio界面...")
         # 使用最小配置
-        demo.launch(server_name="0.0.0.0", show_api=False)
+        demo.launch(
+            server_name="0.0.0.0", 
+            show_api=False,      # 禁用API
+            share=False,         # 不创建公共链接
+            debug=False,         # 禁用调试��式
+            quiet=True           # 减少日志输出
+        )
     except Exception as e:
-        logger.error(f"Failed to launch: {e}")
+        logger.error(f"启动失败: {e}")