import spaces  # MUST come before torch / diffusers

import random

import gradio as gr
import numpy as np
import torch
from diffusers import ChromaPipeline

MODEL_ID = "lodestones/Chroma1-HD"
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024

DEFAULT_NEGATIVE = (
    "low quality, ugly, unfinished, out of focus, deformed, disfigure, "
    "blurry, smudged, restricted palette, flat colors"
)

pipe = ChromaPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
pipe = pipe.to("cuda")


def _gpu_duration(
    prompt,
    negative_prompt,
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,
    num_inference_steps,
    *args,
    **kwargs,
):
    """Estimate GPU seconds from steps and resolution."""
    pixels = max(int(width) * int(height), 256 * 256)
    scale = pixels / (1024 * 1024)
    return min(240, max(45, int(25 + int(num_inference_steps) * 2.2 * scale)))


@spaces.GPU(duration=_gpu_duration)
def generate(
    prompt: str,
    negative_prompt: str = DEFAULT_NEGATIVE,
    seed: int = 433,
    randomize_seed: bool = True,
    width: int = 1024,
    height: int = 1024,
    guidance_scale: float = 3.0,
    num_inference_steps: int = 40,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate an image with Chroma1-HD from a text prompt.

    Chroma1-HD is an 8.9B Apache-2.0 text-to-image model based on FLUX.1-schnell.
    """
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    if randomize_seed:
        seed = int(random.randint(0, MAX_SEED))
    else:
        seed = int(seed)

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

    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt or DEFAULT_NEGATIVE,
        guidance_scale=float(guidance_scale),
        num_inference_steps=int(num_inference_steps),
        width=int(width),
        height=int(height),
        generator=generator,
        num_images_per_prompt=1,
    ).images[0]

    return image, seed


EXAMPLES = [
    [
        "A high-fashion close-up portrait of a blonde woman in clear sunglasses. "
        "The image uses a bold teal and red color split for dramatic lighting. "
        "The background is a simple teal-green. The photo is sharp and well-composed."
    ],
    ["A golden retriever wearing a tiny crown, eating a slice of pepperoni pizza, studio photo"],
    ["The spirit of a tamagotchi wandering in foggy San Francisco at dusk, cinematic lighting"],
    [
        "An ancient library floating in the clouds, spiral staircases of books, "
        "golden hour light through stained glass, highly detailed"
    ],
]

CSS = """
#col-container {
    margin: 0 auto;
    max-width: 820px;
}
.generate-btn {
    align-self: stretch;
}
"""

with gr.Blocks(
    title="Chroma1-HD",
    css=CSS,
    theme=gr.themes.Soft(
        primary_hue="indigo",
        secondary_hue="violet",
        neutral_hue="slate",
    ),
) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
# Chroma1-HD

[**Chroma1-HD**](https://huggingface.co/lodestones/Chroma1-HD) is an **8.9B** parameter
Apache-2.0 text-to-image foundational model based on **FLUX.1-schnell**.
Built as a strong, neutral base for finetuning and creative work.
"""
        )

        with gr.Row():
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe the image you want to create…",
                lines=3,
                max_lines=8,
            )
        with gr.Row():
            negative_prompt = gr.Textbox(
                label="Negative prompt",
                value=DEFAULT_NEGATIVE,
                lines=2,
                max_lines=6,
            )

        run_button = gr.Button("Generate", variant="primary", elem_classes=["generate-btn"])

        result = gr.Image(label="Result", type="pil", show_label=True)

        with gr.Accordion("Advanced settings", open=False):
            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=1.0,
                    maximum=10.0,
                    step=0.1,
                    value=3.0,
                )
                num_inference_steps = gr.Slider(
                    label="Inference steps",
                    minimum=1,
                    maximum=80,
                    step=1,
                    value=40,
                )
            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
            with gr.Row():
                seed = gr.Slider(
                    label="Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=433,
                )
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

        gr.Examples(
            examples=EXAMPLES,
            inputs=[prompt],
            outputs=[result, seed],
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
            label="Example prompts",
        )

        gr.Markdown(
            """
---
**Model:** [lodestones/Chroma1-HD](https://huggingface.co/lodestones/Chroma1-HD) ·
**License:** Apache 2.0 ·
**Tips:** Start with guidance ~3.0 and 30–40 steps at 1024×1024.
"""
        )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=generate,
        inputs=[
            prompt,
            negative_prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, seed],
        api_name="generate",
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch(mcp_server=True)