import gradio as gr
import spaces
import torch
from diffusers import NewbiePipeline, EulerDiscreteScheduler
from transformers import AutoModel
import random
import warnings

warnings.filterwarnings("ignore")

# 优化：使用 float16 以获得更好的 GPU 性能和内存效率
torch_dtype = torch.float16

model_path = "Disty0/NewBie-image-Exp0.1-Diffusers"

print("正在加载 Text Encoder...")
text_encoder_2 = AutoModel.from_pretrained(
    model_path, 
    subfolder="text_encoder_2", 
    trust_remote_code=True, 
    torch_dtype=torch_dtype
)

print("正在加载 Pipeline...")
# 优化：移除显式调度器设置，让模型自动加载正确的调度器
pipe_newbie = NewbiePipeline.from_pretrained(
    model_path, 
    text_encoder_2=text_encoder_2, 
    torch_dtype=torch_dtype
)

# 打印调度器信息以调试
print(f"模型使用的调度器: {type(pipe_newbie.scheduler).__name__}")

pipe_newbie.to("cuda")
print("模型加载完成。")

@spaces.GPU()
def generate_image_newbie(prompt, negative_prompt, height, width, num_inference_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
    if seed < 0:
        seed = random.randint(0, 2**32 - 1)

    generator = torch.Generator("cuda").manual_seed(int(seed))

    try:
        # 优化：调整参数以避免黑色图像
        image = pipe_newbie(
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=int(height),
            width=int(width),
            num_inference_steps=num_inference_steps,  # 使用用户指定的步数
            guidance_scale=guidance_scale,
            generator=generator,
            output_type="pil"
        ).images[0]
        
        return image, seed
    except Exception as e:
        print(f"生成图像时出错: {e}")
        # 返回一个错误图像或占位符
        return None, seed

# 默认提示词 - 优化提示词以生成更好的图像
newbie_prompt = """masterpiece, best quality, high resolution, detailed, 1girl, solo, blonde hair, long hair, large hat, white hat, white blouse, puffy sleeves, shoulder cutout, black skirt, shirt tucked in, socks, shoes, sitting, reclining, arm support, from side, cowboy shot, wide shot, center, cheerful, dynamic angle, depth of field, high contrast, colorful, detailed light, light leaks, beautiful detailed glow, best shadow, shiny skin, cinematic lighting, ray tracing, from above, female focus, close-up, dutch angle, blue archive, detailed background, scenery, very aesthetic, masterpiece, no text, bag, 2024 year"""

with gr.Blocks(title="NewBie 图像生成器") as demo:
    gr.Markdown("# NewBie 图像生成器 (NewBie Image Generator)")
    gr.Markdown("使用 NewBie 模型生成高质量动漫风格图像。")
    
    with gr.Row(variant="panel"):
        with gr.Column(scale=2):
            prompt_newbie = gr.Textbox(
                label="提示词 (Prompt)",
                value=newbie_prompt,
                lines=10,
                placeholder="在此输入生成提示词..."
            )
            negative_prompt_newbie = gr.Textbox(
                label="负面提示词 (Negative Prompt)",
                value="low quality, bad quality, blurry, low resolution, deformed, ugly, bad anatomy, black image, dark, black background",
                lines=3,
                placeholder="在此输入不需要出现的元素..."
            )
            
            with gr.Row():
                height_newbie = gr.Slider(label="图片高度 (Height)", minimum=512, maximum=1024, step=64, value=768)
                width_newbie = gr.Slider(label="图片宽度 (Width)", minimum=512, maximum=1024, step=64, value=768)
            
            with gr.Row():
                steps_newbie = gr.Slider(label="推理步数 (Inference Steps)", minimum=10, maximum=50, step=1, value=25)  # 增加默认步数
                guidance_scale_newbie = gr.Slider(label="引导系数 (Guidance Scale)", minimum=1.0, maximum=20.0, step=0.5, value=7.5)  # 增加默认值
            
            seed_newbie = gr.Number(label="随机种子 (Seed, -1 为随机)", value=-1, precision=0)
            
            generate_btn_newbie = gr.Button("生成图片 (Generate)", variant="primary")

        with gr.Column(scale=1):
            image_output_newbie = gr.Image(label="生成结果 (Output)", format="png", interactive=False)
            used_seed_newbie = gr.Number(label="使用的种子 (Used Seed)", interactive=False)
            error_message = gr.Textbox(label="错误信息", interactive=False, visible=False)

    generate_btn_newbie.click(
        fn=generate_image_newbie,
        inputs=[
            prompt_newbie, 
            negative_prompt_newbie, 
            height_newbie, 
            width_newbie, 
            steps_newbie, 
            guidance_scale_newbie, 
            seed_newbie
        ],
        outputs=[image_output_newbie, used_seed_newbie]
    ).then(
        fn=lambda img, seed: (img, seed, gr.update(visible=True) if img is None else gr.update(visible=False)),
        inputs=[image_output_newbie, used_seed_newbie],
        outputs=[image_output_newbie, used_seed_newbie, error_message]
    )

if __name__ == "__main__":
    demo.launch(
        theme=gr.themes.Soft(primary_hue="blue"),
        footer_links=[{"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}],
        share=True
    )