import gc
import os
from typing import Optional

import gradio as gr
import torch
from diffusers import AutoPipelineForText2Image, EulerAncestralDiscreteScheduler

MODEL_ID = os.getenv("MODEL_ID", "stablediffusionapi/counterfeit-v30")
DEFAULT_NEGATIVE = os.getenv(
    "DEFAULT_NEGATIVE",
    "blurry, lowres, worst quality, low quality, bad anatomy, watermark, text, extra fingers, realistic, photo, 3d, child, loli, young-looking",
)

PRESETS = {
    "🎀 动漫头像": {
        "prompt": "masterpiece, best quality, 1girl, solo, anime style, upper body portrait, detailed eyes, clean lineart, soft cel shading",
        "negative": DEFAULT_NEGATIVE,
        "steps": 12,
        "cfg": 7.0,
        "width": 384,
        "height": 576,
        "seed": -1,
    },
    "💼 成熟御姐": {
        "prompt": "masterpiece, best quality, 1woman, mature female, solo, anime style, office lady, detailed eyes, upper body portrait, clean lineart",
        "negative": DEFAULT_NEGATIVE,
        "steps": 12,
        "cfg": 7.0,
        "width": 384,
        "height": 576,
        "seed": -1,
    },
    "🌸 樱花和服": {
        "prompt": "masterpiece, best quality, 1girl, anime style, kimono, cherry blossoms, elegant pose, detailed face, clean lineart",
        "negative": DEFAULT_NEGATIVE,
        "steps": 14,
        "cfg": 7.0,
        "width": 448,
        "height": 640,
        "seed": -1,
    },
    "🏙 赛博少女": {
        "prompt": "masterpiece, best quality, 1girl, anime style, cyberpunk city, neon lights, upper body, detailed eyes, clean lineart",
        "negative": DEFAULT_NEGATIVE,
        "steps": 14,
        "cfg": 7.0,
        "width": 448,
        "height": 640,
        "seed": -1,
    },
}

_pipe = None
_device = "cuda" if torch.cuda.is_available() else "cpu"
_dtype = torch.float16 if _device == "cuda" else torch.float32


def load_pipe():
    global _pipe
    if _pipe is None:
        pipe = AutoPipelineForText2Image.from_pretrained(
            MODEL_ID,
            torch_dtype=_dtype,
            safety_checker=None,
        )
        if hasattr(pipe, "safety_checker"):
            pipe.safety_checker = None
        if hasattr(pipe, "requires_safety_checker"):
            pipe.requires_safety_checker = False
        if hasattr(pipe, "scheduler") and pipe.scheduler is not None:
            pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
        if hasattr(pipe, "enable_attention_slicing"):
            pipe.enable_attention_slicing()
        if hasattr(pipe, "enable_vae_slicing"):
            pipe.enable_vae_slicing()
        if hasattr(pipe, "vae") and hasattr(pipe.vae, "enable_tiling"):
            pipe.vae.enable_tiling()
        pipe = pipe.to(_device)
        _pipe = pipe
    return _pipe


def unload_pipe():
    global _pipe
    if _pipe is not None:
        _pipe = None
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


def apply_preset(name: str):
    p = PRESETS[name]
    return (
        p["prompt"],
        p["negative"],
        p["steps"],
        p["cfg"],
        p["width"],
        p["height"],
        p["seed"],
    )


def generate(
    prompt: str,
    negative_prompt: str,
    steps: int,
    guidance_scale: float,
    width: int,
    height: int,
    seed: int,
):
    if not prompt or not prompt.strip():
        raise gr.Error("提示词不能为空")

    pipe = load_pipe()
    generator: Optional[torch.Generator] = None
    if seed >= 0:
        generator = torch.Generator(device=_device).manual_seed(int(seed))

    with torch.inference_mode():
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=int(steps),
            guidance_scale=float(guidance_scale),
            width=int(width),
            height=int(height),
            generator=generator,
        ).images[0]
    return image


with gr.Blocks(title="wode") as demo:
    gr.Markdown(
        f"# wode\n\n免费档二次元生图（Counterfeit 路线）\n\n当前模型：`{MODEL_ID}`\n\n推荐先用：`384x576 / 12 steps / CFG 7`。\n\n想更像动漫角色：正向词尽量写 `1girl/1woman + anime style + detailed eyes + clean lineart`。\n想偏成年风格：加 `mature female` 或 `adult woman`。默认反向词已压制 `realistic/photo` 和 `child-like`。\n\n⚠️ 免费 CPU 仍然会慢，通常要等 `1~5 分钟`。"
    )
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="正向提示词",
                lines=4,
                placeholder="例如：masterpiece, best quality, 1girl, anime style, detailed eyes, clean lineart",
            )
            negative_prompt = gr.Textbox(label="反向提示词", lines=3, value=DEFAULT_NEGATIVE)
            with gr.Row():
                steps = gr.Slider(6, 24, value=12, step=1, label="步数 Steps")
                guidance_scale = gr.Slider(1, 12, value=7.0, step=0.5, label="引导强度 CFG")
            with gr.Row():
                width = gr.Slider(256, 640, value=384, step=64, label="宽度")
                height = gr.Slider(256, 896, value=576, step=64, label="高度")
            seed = gr.Number(label="随机种子（-1 为随机）", value=-1, precision=0)
            with gr.Row():
                run = gr.Button("开始生成", variant="primary")
                unload = gr.Button("释放模型内存")
        with gr.Column():
            out = gr.Image(label="生成结果", type="pil")

    gr.Markdown("## 二次元预设按钮")
    with gr.Row():
        preset1 = gr.Button("🎀 动漫头像")
        preset2 = gr.Button("💼 成熟御姐")
        preset3 = gr.Button("🌸 樱花和服")
        preset4 = gr.Button("🏙 赛博少女")

    for btn, name in [
        (preset1, "🎀 动漫头像"),
        (preset2, "💼 成熟御姐"),
        (preset3, "🌸 樱花和服"),
        (preset4, "🏙 赛博少女"),
    ]:
        btn.click(
            fn=lambda n=name: apply_preset(n),
            inputs=[],
            outputs=[prompt, negative_prompt, steps, guidance_scale, width, height, seed],
        )

    gr.Examples(
        label="也可以直接点下面示例",
        examples=[
            [
                PRESETS["🎀 动漫头像"]["prompt"],
                DEFAULT_NEGATIVE,
                12,
                7.0,
                384,
                576,
                -1,
            ],
            [
                PRESETS["💼 成熟御姐"]["prompt"],
                DEFAULT_NEGATIVE,
                12,
                7.0,
                384,
                576,
                -1,
            ],
        ],
        inputs=[prompt, negative_prompt, steps, guidance_scale, width, height, seed],
    )

    run.click(
        fn=generate,
        inputs=[prompt, negative_prompt, steps, guidance_scale, width, height, seed],
        outputs=out,
    )
    unload.click(fn=unload_pipe, outputs=[])

if __name__ == "__main__":
    demo.queue(max_size=4).launch(show_error=True)