"""
Z-Image Generation Demo
Clean UI for the non-distilled Z-Image model.
"""

import spaces
import random
import re
import torch
import gradio as gr
from diffusers import ZImagePipeline

# ==================== Configuration ====================
MODEL_PATH = "Tongyi-MAI/Z-Image"

# ==================== Resolution Choices ====================
RES_CHOICES = {
    "720": [
        "720x720 ( 1:1 )",
        "896x512 ( 16:9 )",
        "512x896 ( 9:16 )",
        "832x544 ( 3:2 )",
        "544x832 ( 2:3 )",
        "800x576 ( 4:3 )",
        "576x800 ( 3:4 )",
    ],
    "1024": [
        "1024x1024 ( 1:1 )",
        "1152x896 ( 9:7 )",
        "896x1152 ( 7:9 )",
        "1152x864 ( 4:3 )",
        "864x1152 ( 3:4 )",
        "1248x832 ( 3:2 )",
        "832x1248 ( 2:3 )",
        "1280x720 ( 16:9 )",
        "720x1280 ( 9:16 )",
        "1344x576 ( 21:9 )",
        "576x1344 ( 9:21 )",
    ],
    "1280": [
        "1280x1280 ( 1:1 )",
        "1440x1120 ( 9:7 )",
        "1120x1440 ( 7:9 )",
        "1472x1104 ( 4:3 )",
        "1104x1472 ( 3:4 )",
        "1536x1024 ( 3:2 )",
        "1024x1536 ( 2:3 )",
        "1536x864 ( 16:9 )",
        "864x1536 ( 9:16 )",
        "1680x720 ( 21:9 )",
        "720x1680 ( 9:21 )",
    ],
}

RESOLUTION_SET = []
for resolutions in RES_CHOICES.values():
    RESOLUTION_SET.extend(resolutions)

EXAMPLE_PROMPTS = [
    ["一位男士和他的贵宾犬穿着配套的服装参加狗狗秀，室内灯光，背景中有观众。"],
    ["极具氛围感的暗调人像，一位优雅的中国美女在黑暗的房间里。一束强光通过遮光板，在她的脸上投射出一个清晰的闪电形状的光影，正好照亮一只眼睛。高对比度，明暗交界清晰，神秘感，莱卡相机色调。"],
    ["Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. Soft-lit outdoor night background, silhouetted tiered pagoda, blurred colorful distant lights."],
    ["A serene mountain landscape at sunset with golden light reflecting off a calm lake, surrounded by pine trees"],
    ["A futuristic cityscape with flying cars and neon holographic advertisements, cyberpunk style"],
]


# ==================== Helper Functions ====================
def get_resolution(resolution: str) -> tuple[int, int]:
    """Parse resolution string to width and height."""
    match = re.search(r"(\d+)\s*[×x]\s*(\d+)", resolution)
    if match:
        return int(match.group(1)), int(match.group(2))
    return 1024, 1024


# ==================== Model Loading (Global Context) ====================
print(f"Loading Z-Image pipeline from {MODEL_PATH}...")
pipe = ZImagePipeline.from_pretrained(
    MODEL_PATH,
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=False,
)
pipe.to("cuda")
print("Pipeline loaded successfully!")

# pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
# spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")


# ==================== Generation Function ====================
@spaces.GPU
def generate(
    prompt: str,
    negative_prompt: str = "",
    resolution: str = "1024x1024 ( 1:1 )",
    seed: int = 42,
    num_inference_steps: int = 50,
    guidance_scale: float = 4.0,
    cfg_normalization: bool = False,
    random_seed: bool = True,
    gallery_images: list = None,
    progress=gr.Progress(track_tqdm=True),
):
    """
    Generate an image using the Z-Image diffusion model.

    Creates high-quality images from text prompts using the Z-Image
    single-stream diffusion transformer. Supports English and Chinese prompts.

    Args:
        prompt (str): Text description of the image to generate.
        negative_prompt (str): What to avoid in the generated image.
        resolution (str): Image resolution, e.g. "1024x1024 ( 1:1 )".
        seed (int): Random seed for reproducibility.
        num_inference_steps (int): Denoising steps (10-100). Higher = better quality.
        guidance_scale (float): CFG scale (1.0-20.0). Higher = more prompt adherence.
        cfg_normalization (bool): Whether to apply CFG normalization.
        random_seed (bool): Use a random seed instead of the provided one.
        gallery_images (list): Previous images to prepend new result to.

    Returns:
        tuple: Gallery with new image, seed string, and seed number.
    """
    
    if not prompt.strip():
        raise gr.Error("Please enter a prompt.")
    
    # Handle seed
    if random_seed:
        new_seed = random.randint(1, 1000000)
    else:
        new_seed = seed if seed != -1 else random.randint(1, 1000000)
    
    # Parse resolution
    width, height = get_resolution(resolution)
    
    # Generate
    generator = torch.Generator("cuda").manual_seed(new_seed)
    
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt if negative_prompt.strip() else None,
        height=height,
        width=width,
        cfg_normalization=cfg_normalization,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        generator=generator,
    ).images[0]
    
    # Update gallery
    if gallery_images is None:
        gallery_images = []
    gallery_images = [image] + gallery_images
    
    return gallery_images, str(new_seed), int(new_seed)


# ==================== Gradio Interface ====================
output_gallery = gr.Gallery(
    label="Generated Images",
    columns=2,
    rows=2,
    height=600,
    object_fit="contain",
    format="png",
    interactive=False,
)

used_seed = gr.Textbox(label="Seed Used", interactive=False)

with gr.Blocks(title="Z-Image Demo") as demo:
    gr.Markdown(
        """<div align="center">

# Z-Image Base
<div style="display: flex; justify-content: center; align-items: center; gap: 8px;">
  <a href="https://github.com/Tongyi-MAI/Z-Image">
    <img src="https://img.shields.io/badge/GitHub-Z--Image-181717?logo=github&logoColor=white" alt="GitHub">
  </a>
  <a href="https://huggingface.co/Tongyi-MAI/Z-Image">
    <img src="https://img.shields.io/badge/🤗%20Hugging%20Face-Model-yellow" alt="Hugging Face">
  </a>
</div>
*An Efficient Image Generation Foundation Model with Single-Stream Diffusion Transformer*

</div>"""
    )

    with gr.Row():
        with gr.Column(scale=1):
            prompt_input = gr.Textbox(
                label="Prompt",
                lines=3,
                placeholder="Enter your prompt here..."
            )
            negative_prompt_input = gr.Textbox(
                label="Negative Prompt (optional)",
                lines=2,
                placeholder="Enter what you want to avoid..."
            )

            with gr.Row():
                choices = [int(k) for k in RES_CHOICES.keys()]
                res_cat = gr.Dropdown(
                    value=1024,
                    choices=choices,
                    label="Resolution Category"
                )
                resolution = gr.Dropdown(
                    value=RES_CHOICES["1024"][0],
                    choices=RESOLUTION_SET,
                    label="Width x Height (Ratio)"
                )

            with gr.Row():
                seed = gr.Number(label="Seed", value=42, precision=0)
                random_seed = gr.Checkbox(label="Random Seed", value=True)

            with gr.Row():
                num_inference_steps = gr.Slider(
                    label="Inference Steps",
                    minimum=10,
                    maximum=100,
                    value=30,
                    step=1
                )
                guidance_scale = gr.Slider(
                    label="Guidance Scale (CFG)",
                    minimum=1.0,
                    maximum=20.0,
                    value=4.0,
                    step=0.5
                )

            cfg_normalization = gr.Checkbox(
                label="CFG Normalization",
                value=False
            )

            generate_btn = gr.Button("Generate", variant="primary")

            # Example prompts
            gr.Markdown("### 📝 Example Prompts")
            gr.Examples(
                fn=generate,
                examples=EXAMPLE_PROMPTS,
                inputs=prompt_input,
                outputs=[output_gallery, used_seed, seed],
                cache_examples=True,
                cache_mode="lazy",
                label=None
            )

        with gr.Column(scale=1):
            output_gallery.render()
            used_seed.render()

    # Event handlers
    def update_res_choices(res_cat_value):
        if str(res_cat_value) in RES_CHOICES:
            res_choices = RES_CHOICES[str(res_cat_value)]
        else:
            res_choices = RES_CHOICES["1024"]
        return gr.update(value=res_choices[0], choices=res_choices)

    res_cat.change(
        update_res_choices,
        inputs=res_cat,
        outputs=resolution
    )

    generate_btn.click(
        generate,
        inputs=[
            prompt_input,
            negative_prompt_input,
            resolution,
            seed,
            num_inference_steps,
            guidance_scale,
            cfg_normalization,
            random_seed,
            output_gallery,
        ],
        outputs=[output_gallery, used_seed, seed],
        api_name="generate",
    )


# ==================== Launch ====================
css = """
.fillable{max-width: 1230px !important}
"""
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        mcp_server=True,
        css=css
    )