import os
import time
import torch
from PIL import Image
import gradio as gr
from diffusers import DiffusionPipeline

# 设置环境变量，确保使用CPU
os.environ["CUDA_VISIBLE_DEVICES"] = ""

class SimpleImageGenerator:
    def __init__(self, model_path="./models"):
        self.model_path = model_path
        self.pipeline = None
        self.is_initialized = False
        
    def initialize(self):
        """初始化模型"""
        if self.is_initialized:
            return
        
        print("正在加载模型...")
        start_time = time.time()
        
        # 设置PyTorch优化选项
        torch.backends.cudnn.enabled = False
        torch.set_num_threads(4)
        torch.set_num_interop_threads(2)
        
        # 加载模型
        self.pipeline = DiffusionPipeline.from_pretrained(
            self.model_path,
            torch_dtype=torch.float32,
            device_map="auto",
            max_memory={"cpu": "16GB"},
            low_cpu_mem_usage=True,
            use_safetensors=True,
            trust_remote_code=True
        )
        
        # 优化管道
        self.pipeline.enable_attention_slicing()
        self.pipeline.enable_sequential_cpu_offload()
        self.pipeline.enable_model_cpu_offload()
        
        self.is_initialized = True
        end_time = time.time()
        print(f"模型加载完成，耗时: {end_time - start_time:.2f}秒")
    
    def generate_image(self, prompt, negative_prompt, width, height, num_inference_steps, guidance_scale, seed, num_images):
        """生成图像"""
        if not self.is_initialized:
            self.initialize()
        
        if not prompt:
            return [], "0.00秒", "0", "请输入生成提示"
        
        try:
            start_time = time.time()
            
            # 设置随机种子
            if seed is not None:
                generator = torch.Generator(device="cpu").manual_seed(seed)
            else:
                generator = None
            
            # 生成图像
            with torch.no_grad():
                results = self.pipeline(
                    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,
                    num_images_per_prompt=num_images,
                    output_type="pil"
                )
            
            end_time = time.time()
            execution_time = end_time - start_time
            
            return (
                results.images,
                f"{execution_time:.2f}秒",
                f"{len(results.images)}",
                "生成完成"
            )
            
        except Exception as e:
            print(f"生成失败: {str(e)}")
            return [], "0.00秒", "0", f"生成失败: {str(e)}"
    
    def clear_all(self):
        """清除所有输入和输出"""
        return (
            "",  # prompt
            "",  # negative_prompt
            1024,  # width
            1024,  # height
            50,  # num_inference_steps
            7.5,  # guidance_scale
            None,  # seed
            1,  # num_images
            [],  # gallery
            "",  # execution_time
            "",  # image_count
            "就绪，可以生成图像"  # status_text
        )

# 创建生成器实例
generator = SimpleImageGenerator()

# 创建Gradio界面
with gr.Blocks(
    title="Qwen-Image-2512 文本到图像生成",
    theme=gr.themes.Soft(),
    css="""
    .title { text-align: center; margin-bottom: 2rem; }
    .param-group { display: flex; flex-wrap: wrap; gap: 1rem; margin-bottom: 1rem; }
    .param-item { flex: 1 1 200px; }
    """
) as interface:
    
    # 标题
    gr.HTML("""
    <h1 class="title">Qwen-Image-2512 文本到图像生成</h1>
    <p class="title" style="font-size: 1.2rem; color: #666;">基于阿里通义千问的高性能图像生成模型</p>
    """)
    
    # 状态显示
    with gr.Row():
        status_text = gr.Textbox(
            label="状态",
            value="模型加载中...",
            interactive=False
        )
    
    # 主要内容区域
    with gr.Row():
        # 左侧：输入和参数
        with gr.Column(scale=1, min_width=300):
            # 文本提示输入
            with gr.Group():
                prompt = gr.Textbox(
                    label="生成提示",
                    placeholder="输入您想要生成的图像描述...",
                    lines=3,
                    max_lines=5
                )
                
                negative_prompt = gr.Textbox(
                    label="负面提示",
                    placeholder="输入您想要避免的内容...",
                    lines=2,
                    max_lines=3
                )
            
            # 参数控制面板
            with gr.Group():
                gr.Markdown("### 生成参数")
                
                with gr.Row():
                    # 图像尺寸
                    with gr.Column():
                        width = gr.Slider(
                            label="宽度",
                            minimum=256,
                            maximum=2512,
                            step=64,
                            value=1024
                        )
                        
                        height = gr.Slider(
                            label="高度",
                            minimum=256,
                            maximum=2512,
                            step=64,
                            value=1024
                        )
                    
                    # 推理参数
                    with gr.Column():
                        num_inference_steps = gr.Slider(
                            label="推理步数",
                            minimum=1,
                            maximum=100,
                            step=1,
                            value=50
                        )
                        
                        guidance_scale = gr.Slider(
                            label="引导尺度",
                            minimum=0.0,
                            maximum=20.0,
                            step=0.1,
                            value=7.5
                        )
                    
                    # 其他参数
                    with gr.Column():
                        seed = gr.Number(
                            label="随机种子",
                            value=None,
                            precision=0
                        )
                        
                        num_images = gr.Slider(
                            label="生成数量",
                            minimum=1,
                            maximum=4,
                            step=1,
                            value=1
                        )
            
            # 生成按钮
            with gr.Row():
                generate_btn = gr.Button(
                    "生成图像",
                    variant="primary",
                    size="lg"
                )
                
                clear_btn = gr.Button(
                    "清除",
                    variant="secondary",
                    size="lg"
                )
        
        # 右侧：结果展示
        with gr.Column(scale=2, min_width=500):
            with gr.Group():
                gr.Markdown("### 生成结果")
                
                # 图像输出区域
                gallery = gr.Gallery(
                    label="生成的图像",
                    show_label=False,
                    columns=2,
                    rows=2,
                    object_fit="contain",
                    height="auto"
                )
                
                # 生成信息
                with gr.Row():
                    execution_time = gr.Textbox(
                        label="生成时间",
                        interactive=False
                    )
                    
                    image_count = gr.Textbox(
                        label="图像数量",
                        interactive=False
                    )
    
    # 示例提示
    with gr.Row():
        gr.Markdown("### 示例提示")
        
        examples = gr.Examples(
            examples=[
                ["一只可爱的柯基犬在草地上奔跑，阳光明媚，高清细节", "模糊, 低质量, 变形", 1024, 1024, 50, 7.5, None, 1],
                ["一个未来主义城市的夜景，霓虹灯闪烁，飞行器穿梭", "模糊, 低质量, 变形", 1024, 1024, 50, 7.5, None, 1],
                ["一朵盛开的玫瑰花，特写镜头，超高清细节，自然光线", "模糊, 低质量, 变形", 1024, 1024, 50, 7.5, None, 1],
            ],
            inputs=[prompt, negative_prompt, width, height, num_inference_steps, guidance_scale, seed, num_images],
            outputs=[gallery, execution_time, image_count, status_text],
            cache_examples=False
        )
    
    # 事件监听
    generate_btn.click(
        fn=generator.generate_image,
        inputs=[prompt, negative_prompt, width, height, num_inference_steps, guidance_scale, seed, num_images],
        outputs=[gallery, execution_time, image_count, status_text],
        show_progress=True
    )
    
    clear_btn.click(
        fn=generator.clear_all,
        inputs=[],
        outputs=[prompt, negative_prompt, width, height, num_inference_steps, guidance_scale, seed, num_images, gallery, execution_time, image_count, status_text]
    )
    
    # 初始化状态
    status_text.value = "就绪，可以生成图像"

if __name__ == "__main__":
    # 初始化生成器
    generator = SimpleImageGenerator()
    generator.initialize()
    
    # 启动Gradio界面
    interface.launch(
        share=False,
        server_name="0.0.0.0",
        server_port=7860,
        show_api=True,
        quiet=True
    )