import gradio as gr
import torch
from diffusers import DiffusionPipeline
import numpy as np
from PIL import Image
import time

# 全局变量存储模型
pipe = None

def load_model():
    """加载 Z-Image-Turbo 模型"""
    global pipe
    if pipe is None:
        try:
            pipe = DiffusionPipeline.from_pretrained(
                "Tongyi-MAI/Z-Image-Turbo",
                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                use_safetensors=True
            )
            
            if torch.cuda.is_available():
                pipe = pipe.to("cuda")
            
            # 启用内存优化
            if hasattr(pipe, 'enable_attention_slicing'):
                pipe.enable_attention_slicing()
            
            return "✅ 模型加载成功！"
        except Exception as e:
            return f"❌ 模型加载失败: {str(e)}"
    return "✅ 模型已加载"

def generate_image(
    prompt,
    negative_prompt,
    num_inference_steps,
    guidance_scale,
    width,
    height,
    seed,
    use_random_seed
):
    """使用 Z-Image-Turbo 生成图像"""
    global pipe
    
    # 确保模型已加载
    if pipe is None:
        load_model()
    
    if pipe is None:
        return None, "❌ 模型未加载，请先加载模型"
    
    try:
        # 设置随机种子
        if use_random_seed:
            seed = np.random.randint(0, 2**32 - 1)
        
        generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu")
        generator.manual_seed(int(seed))
        
        # 记录开始时间
        start_time = time.time()
        
        # 生成图像
        result = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt if negative_prompt else None,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            width=width,
            height=height,
            generator=generator
        )
        
        image = result.images[0]
        
        # 计算生成时间
        generation_time = time.time() - start_time
        
        info_text = f"""
🎨 **图像生成成功！**

⏱️ **生成时间**: {generation_time:.2f} 秒
🔢 **推理步数**: {num_inference_steps}
🌱 **随机种子**: {seed}
📐 **分辨率**: {width}x{height}
🎯 **引导强度**: {guidance_scale}
⚡ **Turbo 模式**: {'启用' if num_inference_steps <= 8 else '标准'}
        """
        
        return image, info_text
        
    except Exception as e:
        return None, f"❌ 生成失败: {str(e)}"

def batch_generate(prompt, negative_prompt, num_images, steps, guidance, width, height):
    """批量生成图像"""
    global pipe
    
    if pipe is None:
        load_model()
    
    if pipe is None:
        return [], "❌ 模型未加载"
    
    images = []
    try:
        for i in range(num_images):
            seed = np.random.randint(0, 2**32 - 1)
            generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu")
            generator.manual_seed(seed)
            
            result = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt if negative_prompt else None,
                num_inference_steps=steps,
                guidance_scale=guidance,
                width=width,
                height=height,
                generator=generator
            )
            
            images.append(result.images[0])
        
        return images, f"✅ 成功生成 {num_images} 张图像！"
    except Exception as e:
        return [], f"❌ 批量生成失败: {str(e)}"

# 创建自定义主题
z_turbo_theme = gr.themes.Soft(
    primary_hue="purple",
    secondary_hue="indigo",
    neutral_hue="slate",
    font=gr.themes.GoogleFont("Inter"),
    text_size="lg",
    spacing_size="md",
    radius_size="lg"
).set(
    button_primary_background_fill="linear-gradient(135deg, #667eea 0%, #764ba2 100%)",
    button_primary_background_fill_hover="linear-gradient(135deg, #764ba2 0%, #667eea 100%)",
    block_title_text_weight="700",
    block_label_text_weight="600"
)

# 示例提示词
example_prompts = [
    [
        "A majestic dragon flying over a mystical mountain landscape at sunset, highly detailed, 8k, fantasy art",
        "blurry, low quality, distorted, ugly",
        8,
        7.5,
        1024,
        1024,
        42,
        False
    ],
    [
        "Futuristic cyberpunk city with neon lights, flying cars, rain-soaked streets, cinematic lighting",
        "low quality, blurry, distorted",
        8,
        8.0,
        1024,
        768,
        123,
        False
    ],
    [
        "Beautiful anime girl with long flowing hair, cherry blossoms, spring scenery, studio ghibli style",
        "ugly, deformed, low quality",
        8,
        7.0,
        768,
        1024,
        456,
        False
    ],
    [
        "Photorealistic portrait of a wise old wizard with a long white beard, magical aura, fantasy art",
        "cartoon, anime, low quality",
        8,
        7.5,
        1024,
        1024,
        789,
        False
    ]
]

with gr.Blocks(fill_height=True) as demo:
    gr.HTML("""
        <div style='text-align: center; padding: 20px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 12px; margin-bottom: 20px;'>
            <h1 style='color: white; margin: 0; font-size: 2.5em;'>🚀 Z-Image Turbo Generator</h1>
            <p style='color: white; margin: 10px 0 0 0; font-size: 1.2em;'>高质量图像生成，仅需 8 步推理！</p>
        </div>
    """)
    
    with gr.Tabs() as tabs:
        # 单图生成标签页
        with gr.Tab("🎨 单图生成", id=0):
            with gr.Row():
                with gr.Column(scale=1):
                    prompt = gr.Textbox(
                        label="📝 提示词 (Prompt)",
                        placeholder="描述您想要生成的图像...",
                        lines=3,
                        value="A beautiful landscape with mountains and a lake at sunset, highly detailed, 8k"
                    )
                    
                    negative_prompt = gr.Textbox(
                        label="🚫 负面提示词 (Negative Prompt)",
                        placeholder="描述您不想要的元素...",
                        lines=2,
                        value="blurry, low quality, distorted, ugly"
                    )
                    
                    with gr.Row():
                        num_inference_steps = gr.Slider(
                            minimum=4,
                            maximum=20,
                            value=8,
                            step=1,
                            label="🔢 推理步数 (建议: 8 步 Turbo 模式)"
                        )
                        
                        guidance_scale = gr.Slider(
                            minimum=1.0,
                            maximum=15.0,
                            value=7.5,
                            step=0.5,
                            label="🎯 引导强度 (Guidance Scale)"
                        )
                    
                    with gr.Row():
                        width = gr.Slider(
                            minimum=512,
                            maximum=1024,
                            value=1024,
                            step=64,
                            label="📏 宽度"
                        )
                        
                        height = gr.Slider(
                            minimum=512,
                            maximum=1024,
                            value=1024,
                            step=64,
                            label="📐 高度"
                        )
                    
                    with gr.Row():
                        seed = gr.Number(
                            label="🌱 随机种子",
                            value=42,
                            precision=0
                        )
                        
                        use_random_seed = gr.Checkbox(
                            label="🎲 使用随机种子",
                            value=False
                        )
                    
                    generate_btn = gr.Button(
                        "🚀 生成图像",
                        variant="primary",
                        size="lg"
                    )
                
                with gr.Column(scale=1):
                    output_image = gr.Image(
                        label="✨ 生成结果",
                        type="pil",
                        height=600
                    )
                    
                    generation_info = gr.Markdown(
                        label="📊 生成信息",
                        value="等待生成..."
                    )
            
            # 示例
            gr.Examples(
                examples=example_prompts,
                inputs=[
                    prompt,
                    negative_prompt,
                    num_inference_steps,
                    guidance_scale,
                    width,
                    height,
                    seed,
                    use_random_seed
                ],
                outputs=[output_image, generation_info],
                fn=generate_image,
                cache_examples=False,
                label="💡 示例提示词"
            )
        
        # 批量生成标签页
        with gr.Tab("🎭 批量生成", id=1):
            with gr.Row():
                with gr.Column(scale=1):
                    batch_prompt = gr.Textbox(
                        label="📝 提示词",
                        placeholder="描述您想要生成的图像...",
                        lines=3,
                        value="A serene Japanese garden with cherry blossoms, koi pond, traditional architecture"
                    )
                    
                    batch_negative = gr.Textbox(
                        label="🚫 负面提示词",
                        placeholder="描述您不想要的元素...",
                        lines=2,
                        value="blurry, low quality, distorted"
                    )
                    
                    num_images = gr.Slider(
                        minimum=1,
                        maximum=4,
                        value=2,
                        step=1,
                        label="🔢 生成数量"
                    )
                    
                    with gr.Row():
                        batch_steps = gr.Slider(
                            minimum=4,
                            maximum=20,
                            value=8,
                            step=1,
                            label="🔢 推理步数"
                        )
                        
                        batch_guidance = gr.Slider(
                            minimum=1.0,
                            maximum=15.0,
                            value=7.5,
                            step=0.5,
                            label="🎯 引导强度"
                        )
                    
                    with gr.Row():
                        batch_width = gr.Slider(
                            minimum=512,
                            maximum=1024,
                            value=1024,
                            step=64,
                            label="📏 宽度"
                        )
                        
                        batch_height = gr.Slider(
                            minimum=512,
                            maximum=1024,
                            value=1024,
                            step=64,
                            label="📐 高度"
                        )
                    
                    batch_generate_btn = gr.Button(
                        "🎨 批量生成",
                        variant="primary",
                        size="lg"
                    )
                
                with gr.Column(scale=1):
                    batch_output = gr.Gallery(
                        label="✨ 批量生成结果",
                        columns=2,
                        height=600,
                        object_fit="contain"
                    )
                    
                    batch_info = gr.Markdown(
                        label="📊 批量生成信息",
                        value="等待批量生成..."
                    )
        
        # 设置标签页
        with gr.Tab("⚙️ 设置", id=2):
            gr.Markdown("### 模型设置")
            
            model_status = gr.Textbox(
                label="📡 模型状态",
                value="未加载",
                interactive=False
            )
            
            load_model_btn = gr.Button(
                "📥 加载模型",
                variant="primary",
                size="lg"
            )
            
            gr.Markdown("""
            ### 📖 使用说明
            
            **Z-Image-Turbo** 是一个高效的图像生成模型，特点：
            
            - ⚡ **超快速度**: 仅需 8 步推理即可生成高质量图像
            - 🎨 **高质量**: 生成细节丰富、色彩鲜艳的图像
            - 💪 **强大功能**: 支持多种风格和主题
            - 🔧 **灵活控制**: 可调节推理步数、引导强度等参数
            
            ### 💡 提示词技巧
            
            1. **详细描述**: 包含主题、风格、细节、光照等
            2. **使用关键词**: 如 "highly detailed", "8k", "cinematic"
            3. **负面提示**: 排除不想要的元素
            4. **推理步数**: 8 步为最佳平衡点（速度 vs 质量）
            
            ### 🎯 推荐参数
            
            - **推理步数**: 8 (Turbo 模式)
            - **引导强度**: 7.0-8.0
            - **分辨率**: 1024x1024 或 768x1024
            """)
            
            gr.Markdown("""
            ### 🖥️ 系统信息
            """)
            
            device_info = gr.Textbox(
                label="设备信息",
                value=f"CUDA 可用: {torch.cuda.is_available()}\n设备: {'GPU' if torch.cuda.is_available() else 'CPU'}",
                interactive=False,
                lines=2
            )
    
    # 事件绑定
    generate_btn.click(
        fn=generate_image,
        inputs=[
            prompt,
            negative_prompt,
            num_inference_steps,
            guidance_scale,
            width,
            height,
            seed,
            use_random_seed
        ],
        outputs=[output_image, generation_info],
        api_visibility="public"
    )
    
    batch_generate_btn.click(
        fn=batch_generate,
        inputs=[
            batch_prompt,
            batch_negative,
            num_images,
            batch_steps,
            batch_guidance,
            batch_width,
            batch_height
        ],
        outputs=[batch_output, batch_info],
        api_visibility="public"
    )
    
    load_model_btn.click(
        fn=load_model,
        inputs=[],
        outputs=[model_status],
        api_visibility="public"
    )
    
    # 页脚
    gr.HTML("""
        <div style='text-align: center; margin-top: 30px; padding: 20px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 12px;'>
            <p style='color: white; margin: 0; font-size: 1.1em;'>
                Built with <a href='https://huggingface.co/spaces/akhaliq/anycoder' style='color: #FFD700; text-decoration: none; font-weight: bold;'>anycoder</a>
            </p>
            <p style='color: rgba(255,255,255,0.8); margin: 10px 0 0 0; font-size: 0.9em;'>
                Powered by Tongyi-MAI/Z-Image-Turbo | Gradio 6
            </p>
        </div>
    """)

if __name__ == "__main__":
    demo.launch(
        theme=z_turbo_theme,
        footer_links=[
            {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"},
            {"label": "Model: Z-Image-Turbo", "url": "https://huggingface.co/Tongyi-MAI/Z-Image-Turbo"}
        ],
        css="""
        .gradio-container {
            max-width: 1400px !important;
        }
        .contain {
            max-width: 100% !important;
        }
        """,
        share=False
    )