app.py

import gradio as gr
import numpy as np
import random
import logging
import sys
import os
import requests
import io
import json
import base64
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__)

# 补丁修复 Gradio JSON Schema 错误
try:
    import gradio_client.utils
    
    # 保存原始函数
    original_get_type = gradio_client.utils.get_type
    
    # 创建新的 get_type 函数，处理布尔值
    def patched_get_type(schema):
        if schema is True or schema is False or schema is None:
            return "any"
        if not isinstance(schema, dict):
            return "any"
        return original_get_type(schema)
    
    # 替换原始函数
    gradio_client.utils.get_type = patched_get_type
    logger.info("Successfully patched gradio_client.utils.get_type")
    
    # 同样修补 _json_schema_to_python_type 函数
    original_json_schema_to_python_type = gradio_client.utils._json_schema_to_python_type
    
    def patched_json_schema_to_python_type(schema, defs=None):
        if schema is True or schema is False:
            return "bool"
        if schema is None:
            return "None"
        if not isinstance(schema, dict):
            return "any"
        
        try:
            return original_json_schema_to_python_type(schema, defs)
        except Exception as e:
            logger.warning(f"Error in json_schema_to_python_type: {e}")
            return "any"
    
    gradio_client.utils._json_schema_to_python_type = patched_json_schema_to_python_type
    logger.info("Successfully patched gradio_client.utils._json_schema_to_python_type")
except Exception as e:
    logger.error(f"Failed to patch Gradio utils: {e}")

# 创建一个备用图像
def create_backup_image(prompt=""):
    logger.info(f"Creating backup image for: {prompt}")
    img = PILImage.new('RGB', (512, 512), color=(240, 240, 250))
    
    try:
        from PIL import ImageDraw, ImageFont
        draw = ImageDraw.Draw(img)
        font = ImageFont.load_default()
        
        # 使用英文消息避免编码问题
        draw.text((20, 20), f"Prompt: {prompt}", fill=(0, 0, 0), font=font)
        draw.text((20, 60), "Model loading failed. Showing placeholder image.", fill=(255, 0, 0), font=font)
        
    except Exception as e:
        logger.error(f"Error creating backup image: {e}")
    
    return img

# 预加载图像用于快速响应
PLACEHOLDER_IMAGE = create_backup_image("placeholder")

# 使用 Hugging Face Inference API 生成图像
def generate_image_with_api(prompt, api_url=None, api_key=None):
    """
    使用 Hugging Face Inference API 生成图像
    
    Parameters:
    - prompt: 文本提示
    - api_url: API 端点 URL (可选)
    - api_key: API 密钥 (可选)
    
    Returns:
    - PIL Image
    """
    logger.info(f"Generating image via API for: {prompt}")
    
    # 默认使用 Stable Diffusion API
    if api_url is None:
        api_url = "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5"
    
    # 尝试从环境变量中获取 API 密钥
    if api_key is None:
        api_key = os.environ.get("HF_API_KEY", "")
    
    # 如果没有 API 密钥，使用公共访问（可能会受到速率限制）
    headers = {}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    
    try:
        # 设置请求参数
        payload = {
            "inputs": prompt,
            "parameters": {
                "num_inference_steps": 10,  # 减少推理步骤以加快速度
                "guidance_scale": 7.5,
                "width": 512,
                "height": 512
            }
        }
        
        # 发送请求
        response = requests.post(api_url, headers=headers, json=payload)
        
        # 检查响应是否成功
        if response.status_code == 200:
            # 从响应中获取图像
            image = PILImage.open(io.BytesIO(response.content))
            logger.info("Successfully generated image via API")
            return image
        else:
            # 如果遇到错误，记录响应并返回备用图像
            error_text = response.text
            logger.error(f"API error: {response.status_code}, {error_text}")
            return None
        
    except Exception as e:
        logger.error(f"Failed to generate image via API: {e}")
        return None

# 使用 Hugging Face Spaces 内置模型生成图像
def generate_image_with_spaces(prompt):
    """使用同一个 Hugging Face Space 内的其他公共空间来生成图像"""
    logger.info(f"Generating image via Spaces for: {prompt}")
    
    try:
        # 一些公共可用的 Stable Diffusion 空间 URLs
        space_urls = [
            "https://huggingface-projects-stable-diffusion-demo.hf.space/api/predict",
            "https://huggingface-projects-text-to-image.hf.space/api/predict",
            "https://dataautogpt-playground.hf.space/api/predict"
        ]
        
        # 尝试每个 URL 直到成功
        for url in space_urls:
            try:
                payload = {
                    "data": [prompt, 7.5, 512, 512] 
                }
                response = requests.post(url, json=payload, timeout=30)
                
                if response.status_code == 200:
                    # 解析 JSON 响应
                    result = response.json()
                    
                    # 根据结果格式提取图像数据
                    if isinstance(result, dict) and 'data' in result:
                        image_data = result['data'][0]
                        
                        # 检查图像数据格式
                        if image_data.startswith('data:image'):
                            # 提取 base64 图像数据
                            image_b64 = image_data.split(',')[1]
                            image_bytes = base64.b64decode(image_b64)
                            image = PILImage.open(io.BytesIO(image_bytes))
                            logger.info(f"Successfully generated image via {url}")
                            return image
            except Exception as e:
                logger.warning(f"Failed to generate with {url}: {e}")
                continue
        
        logger.error("All space URLs failed")
        return None
    
    except Exception as e:
        logger.error(f"Failed to generate image via Spaces: {e}")
        return None

# 使用简单的规则生成图像作为备用方案
def generate_rule_based_image(prompt):
    """当AI模型不可用时使用规则生成图像"""
    logger.info(f"Using rule-based generator for: {prompt}")
    
    # 创建基础图像
    img = PILImage.new('RGB', (512, 512), color=(240, 240, 250))
    
    try:
        from PIL import ImageDraw, ImageFont
        draw = ImageDraw.Draw(img)
        
        # 提取关键词
        prompt_lower = prompt.lower()
        
        # 设置默认颜色和形状
        bg_color = (240, 240, 250)  # 浅蓝背景
        shape_color = (64, 64, 128)  # 深蓝形状
        
        # 基于关键词调整颜色
        if "red" in prompt_lower:
            shape_color = (200, 50, 50)
        elif "blue" in prompt_lower:
            shape_color = (50, 50, 200)
        elif "green" in prompt_lower:
            shape_color = (50, 200, 50)
        elif "yellow" in prompt_lower:
            shape_color = (200, 200, 50)
        
        # 画一个基本形状
        if "cat" in prompt_lower or "kitten" in prompt_lower:
            # 猫头
            draw.ellipse((156, 156, 356, 356), fill=shape_color)
            # 猫眼睛
            draw.ellipse((206, 206, 236, 236), fill=(255, 255, 255))
            draw.ellipse((276, 206, 306, 236), fill=(255, 255, 255))
            # 猫瞳孔
            draw.ellipse((216, 216, 226, 226), fill=(0, 0, 0))
            draw.ellipse((286, 216, 296, 226), fill=(0, 0, 0))
            # 猫鼻子
            draw.polygon([(256, 256), (246, 276), (266, 276)], fill=(255, 150, 150))
            # 猫耳朵
            draw.polygon([(156, 156), (176, 96), (216, 156)], fill=shape_color)
            draw.polygon([(356, 156), (336, 96), (296, 156)], fill=shape_color)
        elif "landscape" in prompt_lower or "mountain" in prompt_lower:
            # 天空
            draw.rectangle([(0, 0), (512, 300)], fill=(100, 150, 250))
            # 山脉
            draw.polygon([(0, 300), (150, 100), (300, 300)], fill=(100, 100, 100))
            draw.polygon([(200, 300), (400, 150), (512, 300)], fill=(80, 80, 80))
            # 地面
            draw.rectangle([(0, 300), (512, 512)], fill=(100, 200, 100))
        elif "castle" in prompt_lower or "building" in prompt_lower:
            # 天空
            draw.rectangle([(0, 0), (512, 200)], fill=(150, 200, 250))
            # 主塔
            draw.rectangle([(200, 200), (312, 400)], fill=shape_color)
            # 塔顶
            draw.polygon([(180, 200), (256, 100), (332, 200)], fill=(180, 0, 0))
            # 小塔
            draw.rectangle([(150, 300), (200, 400)], fill=shape_color)
            draw.rectangle([(312, 300), (362, 400)], fill=shape_color)
            # 城墙
            draw.rectangle([(100, 400), (412, 450)], fill=shape_color)
            # 地面
            draw.rectangle([(0, 450), (512, 512)], fill=(100, 150, 100))
        else:
            # 默认绘制几何形状
            draw.rectangle([(100, 100), (412, 412)], outline=(0, 0, 0), width=2)
            draw.ellipse((150, 150, 362, 362), fill=shape_color)
            draw.polygon([(256, 100), (412, 412), (100, 412)], fill=(shape_color[0]//2, shape_color[1]//2, shape_color[2]//2))
        
        # 添加提示词和说明
        font = ImageFont.load_default()
        draw.text((10, 10), f"Prompt: {prompt}", fill=(0, 0, 0), font=font)
        draw.text((10, 30), "Generated with rules (AI model unavailable)", fill=(100, 100, 100), font=font)
        
    except Exception as e:
        logger.error(f"Error in rule-based image generation: {e}")
    
    return img

# 入口点函数 - 处理请求并生成图像
def generate_image(prompt):
    # 处理空提示
    if not prompt or prompt.strip() == "":
        prompt = "a beautiful landscape"
        logger.info(f"Empty prompt, using default: {prompt}")
    
    logger.info(f"Received prompt: {prompt}")
    
    # 尝试使用 Hugging Face API 生成
    image = generate_image_with_api(prompt)
    if image is not None:
        return image
    
    # 如果 API 方法失败，尝试使用 Spaces
    logger.info("API method failed, trying Spaces method...")
    image = generate_image_with_spaces(prompt)
    if image is not None:
        return image
    
    # 如果所有 AI 方法都失败，使用规则生成
    logger.warning("All AI methods failed, using rule-based image generation")
    return generate_rule_based_image(prompt)

# 为旧版 gradio 创建界面
def create_demo():
    # 使用 Interface 替代 Blocks (兼容3.19.1)
    demo = gr.Interface(
        fn=generate_image,
        inputs=gr.Textbox(
            label="Prompt",
            placeholder="Describe the image you want, e.g.: a cute cat, sunset over mountains...",
            lines=2
        ),
        outputs=gr.Image(label="Generated Image", type="pil"),
        title="Text to Image Generator",
        description="Enter a text description to generate an image. This demo uses Hugging Face API for image generation.",
        examples=[
            "a cute cat sitting on a windowsill",
            "beautiful sunset over mountains",
            "an astronaut riding a horse in space",
            "a fantasy castle on a floating island"
        ],
        flagging_mode="never",  # 使用枚举值 'never' 而不是 False
        cache_examples=False
    )
    
    return demo

# 创建演示界面
demo = create_demo()

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