import os
import gradio as gr
from diffusers import DiffusionPipeline
import torch
from huggingface_hub import login

# 初始化 Hugging Face 登录
if os.getenv('HF_TOKEN'):
    login(token=os.getenv('HF_TOKEN'))

# 模型加载函数
def load_model():
    model_id = "IndexTeam/Index-anisora"
    print("正在加载模型...")
    
    pipe = DiffusionPipeline.from_pretrained(
        model_id,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        variant="fp16" if torch.cuda.is_available() else None,
        use_safetensors=True
    )
    
    if torch.cuda.is_available():
        print("使用CUDA加速")
        pipe = pipe.to("cuda")
        pipe.enable_attention_slicing()
        if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
            pipe.enable_xformers_memory_efficient_attention()
    else:
        print("警告：未检测到CUDA，使用CPU模式运行")
    
    return pipe

# 加载模型
pipe = load_model()

# 生成图像函数
def generate_image(prompt, negative_prompt, steps=30, guidance_scale=7.5, seed=None):
    if seed is not None and seed >= 0:
        generator = torch.Generator("cuda" if torch.cuda.is_available() else "cpu").manual_seed(int(seed))
    else:
        generator = None
    
    try:
        with torch.inference_mode():
            print(f"生成图像: {prompt[:50]}...")
            image = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=steps,
                guidance_scale=guidance_scale,
                generator=generator
            ).images[0]
        return image
    except Exception as e:
        print(f"生成失败: {str(e)}")
        raise gr.Error(f"生成图像时出错: {str(e)}")

# 创建Gradio界面
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # Index-anisora 演示
    """)
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="提示词",
                value="1girl, beautiful anime style, detailed background",
                placeholder="描述你想要生成的图像..."
            )
            negative_prompt = gr.Textbox(
                label="负面提示词",
                value="low quality, blurry, bad anatomy",
                placeholder="描述你不希望在图像中出现的内容..."
            )
            
            with gr.Row():
                steps = gr.Slider(
                    label="推理步数", 
                    minimum=10, 
                    maximum=50, 
                    value=30, 
                    step=1
                )
                guidance_scale = gr.Slider(
                    label="引导强度", 
                    minimum=1.0, 
                    maximum=20.0, 
                    value=7.5, 
                    step=0.5
                )
                seed = gr.Number(
                    label="随机种子", 
                    value=-1,
                    precision=0
                )
            
            submit_btn = gr.Button("生成", variant="primary")
        
        with gr.Column():
            output_image = gr.Image(label="生成结果")
    
    # 提交按钮事件
    submit_btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, steps, guidance_scale, seed],
        outputs=output_image
    )

# 启动应用
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )