import random

import spaces
import torch
import gradio as gr
from diffusers import DiffusionPipeline

BASE_MODEL = "krea/Krea-2-Turbo"
LORA_REPO = "ostris/krea2_turbo_style_reference"
LORA_WEIGHT = "krea2_style_reference.safetensors"
CUSTOM_PIPELINE = "ostris/Krea2OstrisEdit"

DTYPE = torch.bfloat16
MAX_SEED = 2**31 - 1

pipe = DiffusionPipeline.from_pretrained(
    BASE_MODEL,
    custom_pipeline=CUSTOM_PIPELINE,
    torch_dtype=DTYPE,
    trust_remote_code=True,
)
pipe.to("cuda")
pipe.load_lora_weights(LORA_REPO, weight_name=LORA_WEIGHT)


@spaces.GPU(duration=90, size="large")
def generate(
    prompt,
    style_ref_image,
    lora_scale=1.0,
    steps=8,
    guidance=0.0,
    width=1024,
    height=1024,
    seed=0,
    randomize_seed=True,
    progress=gr.Progress(track_tqdm=True),
):
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")
    if style_ref_image is None:
        raise gr.Error("Please upload a style reference image.")

    if randomize_seed or seed is None:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)

    # Defensive defaults: if a value ever arrives as None (e.g. an example
    # row that omits an input), fall back to the same defaults used elsewhere.
    if lora_scale is None:
        lora_scale = 1.0
    if steps is None:
        steps = 8
    if guidance is None:
        guidance = 0.0
    if width is None:
        width = 1024
    if height is None:
        height = 1024

    # Snap dimensions to multiples of 16 (vae_scale_factor * patch_size = 16)
    multiple = 16
    width = ((int(width) + multiple - 1) // multiple) * multiple
    height = ((int(height) + multiple - 1) // multiple) * multiple

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

    image = pipe(
        prompt=prompt,
        image=style_ref_image,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        width=width,
        height=height,
        generator=generator,
        attention_kwargs={"scale": float(lora_scale)},
    ).images[0]

    return image, seed


CSS = """
#page { max-width: 1100px; margin: 0 auto; padding: 4px 8px 32px; }
#header { padding: 24px 4px 18px; border-bottom: 1px solid #e5e5e5; margin-bottom: 20px; }
#header h1 { font-size: 32px; font-weight: 700; margin: 0 0 6px; letter-spacing: -0.02em; }
#header .subtitle { font-size: 15px; color: #666; margin: 0; max-width: 70ch; line-height: 1.5; }
#header .links { margin-top: 12px; display: flex; gap: 16px; }
#header .links a {
    font-size: 13px; color: #888; text-decoration: none;
    border: 1px solid #ddd; border-radius: 6px; padding: 4px 10px;
}
#header .links a:hover { color: #333; border-color: #aaa; }
footer { display: none !important; }
"""

HEADER = """
<div id="header">
  <h1>Krea 2 Style Reference</h1>
  <p class="subtitle">
    Generate an image from a text prompt, guided by a style reference image.
    Powered by Krea&nbsp;2&nbsp;Turbo with the
    <a href="https://huggingface.co/ostris/krea2_turbo_style_reference" target="_blank">Krea2 Style Reference LoRA</a>
    and the
    <a href="https://huggingface.co/ostris/Krea2OstrisEdit" target="_blank">Krea2OstrisEdit community pipeline</a>.
  </p>
  <div class="links">
    <a href="https://huggingface.co/krea/Krea-2-Turbo" target="_blank">Base model ↗</a>
    <a href="https://huggingface.co/ostris/krea2_turbo_style_reference" target="_blank">LoRA ↗</a>
    <a href="https://github.com/ostris/ComfyUI-Krea2-Ostris-Edit" target="_blank">ComfyUI nodes ↗</a>
  </div>
</div>
"""

with gr.Blocks(title="Krea 2 Style Reference") as demo:
    with gr.Column(elem_id="page"):
        gr.HTML(HEADER)
        with gr.Row(equal_height=False):
            with gr.Column(scale=1):
                prompt = gr.Textbox(
                    label="Prompt",
                    lines=3,
                    placeholder="Describe what you want to generate, e.g. 'a white yeti with horns reading a book'",
                )
                style_ref = gr.Image(
                    label="Style Reference Image",
                    type="pil",
                    height=300,
                )
                generate_btn = gr.Button("Generate", variant="primary", size="lg")
                with gr.Accordion("Advanced", open=False):
                    lora_scale = gr.Slider(
                        0.0, 2.0, value=1.0, step=0.01,
                        label="LoRA scale",
                        info="Strength of the style reference influence",
                    )
                    steps = gr.Slider(1, 30, value=8, step=1, label="Steps")
                    guidance = gr.Slider(
                        0.0, 10.0, value=0.0, step=0.1,
                        label="Guidance scale",
                        info="Krea 2 Turbo uses 0.0 (guidance disabled)",
                    )
                    with gr.Row():
                        width = gr.Slider(512, 1536, value=1024, step=16, label="Width")
                        height = gr.Slider(512, 1536, value=1024, step=16, label="Height")
                    with gr.Row():
                        seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
                        randomize_seed = gr.Checkbox(value=True, label="Randomize seed")
            with gr.Column(scale=1):
                result = gr.Image(label="Result", format="png", height=420)
                used_seed = gr.Number(label="Seed used", visible=True, interactive=False)

        inputs = [
            prompt, style_ref, lora_scale, steps, guidance,
            width, height, seed, randomize_seed,
        ]
        outputs = [result, used_seed]

        # Lean examples: only supply the two inputs a user actually varies
        # (prompt + style reference image). generate() provides defaults for
        # the remaining params, so the missing example columns fall back to
        # sane values instead of None.
        # Columns: prompt, style_ref
        gr.Examples(
            examples=[
                ["a white yeti with horns reading a book", "examples/style_ref_yeti.png"],
                ["a futuristic city skyline at sunset, cyberpunk aesthetic", "examples/style_ref_01.png"],
            ],
            inputs=[prompt, style_ref],
            outputs=outputs,
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

        gr.on([generate_btn.click, prompt.submit], generate, inputs, outputs)

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS)