app.py

import gradio as gr
import numpy as np
import random
import logging
import sys
import os
from PIL import Image as PILImage

# 设置日志记录
logging.basicConfig(level=logging.INFO, 
                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
                    stream=sys.stdout)
logger = logging.getLogger(__name__)

# 创建一个简单的示例图像
def create_dummy_image():
    logger.info("Creating dummy image")
    img = PILImage.new('RGB', (256, 256), color = (255, 100, 100))
    return img

# 使用小型模型减轻负担
def get_model():
    try:
        from diffusers import StableDiffusionPipeline
        import torch
        
        logger.info("Loading smaller model: runwayml/stable-diffusion-v1-5")
        
        model_id = "runwayml/stable-diffusion-v1-5"
        device = "cuda" if torch.cuda.is_available() else "cpu"
        logger.info(f"Using device: {device}")
        
        if torch.cuda.is_available():
            pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
            pipe = pipe.to(device)
            pipe.enable_attention_slicing()
        else:
            pipe = StableDiffusionPipeline.from_pretrained(model_id)
            pipe = pipe.to(device)
        
        logger.info("Model loaded successfully")
        return pipe
    except Exception as e:
        logger.error(f"Failed to load model: {e}")
        return None

# 全局变量
pipe = None

# 非常简单的功能实现
def generate(prompt):
    global pipe
    
    # 如果提示为空，使用默认提示
    if not prompt or prompt.strip() == "":
        prompt = "a beautiful landscape"
    
    logger.info(f"Received prompt: {prompt}")
    
    # 第一次调用时加载模型
    if pipe is None:
        pipe = get_model()
        # 如果模型加载失败，返回示例图像
        if pipe is None:
            return create_dummy_image()
    
    try:
        logger.info("Starting image generation")
        # 使用最小的计算量生成图像
        image = pipe(
            prompt=prompt, 
            num_inference_steps=1,
            guidance_scale=7.0,
            height=256,
            width=256
        ).images[0]
        
        logger.info("Image generated successfully")
        return image
    except Exception as e:
        logger.error(f"Generation failed: {e}")
        return create_dummy_image()

# 极简界面
iface = gr.Interface(
    fn=generate,
    inputs="text",
    outputs="image",
    title="Simple Text-to-Image",
    description="Type a prompt to generate an image.",
    examples=["a cat", "mountain landscape"]
)

# 启动应用
if __name__ == "__main__":
    try:
        logger.info("Starting Gradio interface")
        iface.launch(server_name="0.0.0.0")
    except Exception as e:
        logger.error(f"Failed to launch: {e}")