import random

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

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
BASE_MODEL = "krea/Krea-2-Turbo"
LORA_REPO = "ilkerzgi/krea-2-moody-golden-hour-editorial-lora"
LORA_WEIGHT_NAME = "moody-golden-hour-editorial.safetensors"
TRIGGER_WORD = "moody golden hour editorial style"
MAX_SEED = 2**31 - 1

# Resolution presets — all multiples of 16 (Krea 2 requirement).
RESOLUTIONS = {
    "Square · 1024×1024": (1024, 1024),
    "Portrait · 832×1216": (832, 1216),
    "Landscape · 1216×832": (1216, 832),
    "Tall · 1024×1280": (1024, 1280),
}

# Example prompts — drawn from the LoRA's own model-card example and
# extended with subjects that showcase the moody golden-hour editorial style.
EXAMPLES = [
    "a lighthouse on a rocky cliff",
    "a lone figure walking through a misty forest path",
    "a vintage car parked on a desert highway at dusk",
    "a fashion model leaning against a brick wall in soft warm light",
]

# ---------------------------------------------------------------------------
# Model loading — at module scope, on CUDA (ZeroGPU CUDA emulation).
# Krea 2 LoRAs load via the transformer's adapter API, per the official cards.
# ---------------------------------------------------------------------------
pipe = Krea2Pipeline.from_pretrained(BASE_MODEL, torch_dtype=torch.bfloat16)
pipe.to("cuda")
pipe.transformer.load_lora_adapter(LORA_REPO, weight_name=LORA_WEIGHT_NAME)
pipe.transformer.set_adapters("default", weights=1.0)


def _build_prompt(user_prompt: str) -> str:
    """Prepend the trigger phrase so the user doesn't have to."""
    p = user_prompt.strip()
    if not p:
        return p
    if TRIGGER_WORD.lower() in p.lower():
        return p
    return f"{p}. {TRIGGER_WORD}"


@spaces.GPU(duration=60)
def generate(
    prompt,
    resolution_label,
    lora_scale,
    seed,
    randomize_seed,
    progress=gr.Progress(track_tqdm=True),
):
    if not prompt or not prompt.strip():
        raise gr.Error("Enter a prompt to generate an image.")

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

    width, height = RESOLUTIONS[resolution_label]

    # Apply the LoRA scale dynamically.
    pipe.transformer.set_adapters("default", weights=float(lora_scale))

    full_prompt = _build_prompt(prompt)

    generator = torch.Generator("cuda").manual_seed(seed)
    image = pipe(
        prompt=full_prompt,
        num_inference_steps=8,
        guidance_scale=0.0,
        width=width,
        height=height,
        generator=generator,
    ).images[0]

    return image, seed


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
with gr.Blocks(title="Krea 2 · Moody Golden Hour Editorial") as demo:
    gr.Markdown(
        """
        # Krea 2 · Moody Golden Hour Editorial

        A style LoRA for [**Krea-2-Turbo**](https://huggingface.co/krea/Krea-2-Turbo) by
        [**ilkerzgi**](https://huggingface.co/ilkerzgi). Generates moody, warm, golden-hour
        editorial imagery. The trigger phrase *\"moody golden hour editorial style\"* is
        auto-appended to your prompt.
        """
    )

    with gr.Row(equal_height=True):
        with gr.Column(scale=5):
            prompt = gr.Textbox(
                label="Prompt",
                lines=3,
                placeholder="Describe your scene… e.g. a lighthouse on a rocky cliff",
                autofocus=True,
            )
            resolution = gr.Radio(
                choices=list(RESOLUTIONS.keys()),
                value="Square · 1024×1024",
                label="Aspect ratio",
            )
            run_button = gr.Button("Generate", variant="primary", size="lg")

            with gr.Accordion("Advanced", open=False):
                lora_scale = gr.Slider(
                    label="LoRA scale (style strength)",
                    minimum=0.0,
                    maximum=1.5,
                    step=0.05,
                    value=1.0,
                    info="1.0 = default, up to 1.25 for a stronger look.",
                )
                with gr.Row():
                    seed = gr.Slider(
                        label="Seed",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=0,
                    )
                    randomize_seed = gr.Checkbox(
                        label="Randomize seed", value=True
                    )

        with gr.Column(scale=5):
            result = gr.Image(label="Result", format="png", show_label=False)

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

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=generate,
        inputs=[prompt, resolution, lora_scale, seed, randomize_seed],
        outputs=[result, seed],
    )

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