import random

import spaces  # MUST come before any CUDA-touching import (ZeroGPU rule 1)
import gradio as gr
import torch
from diffusers import DiffusionPipeline

MAX_SEED = 2**31 - 1
BASE_MODEL = "black-forest-labs/FLUX.2-klein-base-4b"
LORA_REPO = "aiconiccompany/jppy-logo-flux-lora"

# Module-scope load, eager .to("cuda") — ZeroGPU intercepts and packs these weights.
pipe = DiffusionPipeline.from_pretrained(BASE_MODEL, torch_dtype=torch.bfloat16).to("cuda")
pipe.load_lora_weights(LORA_REPO, adapter_name="jppy-logo")

DESCRIPTION = """
Demo for [aiconiccompany/jppy-logo-flux-lora](https://huggingface.co/aiconiccompany/jppy-logo-flux-lora) —
a logo-design LoRA for FLUX.2-klein trained on 15,815 clean vector-quality logos.
Just describe the mark. The `jppy-logo` trigger is appended automatically unless it's already in your prompt.
"""


def build_prompt(prompt: str, append_trigger: bool) -> str:
    p = prompt.strip()
    if append_trigger and "jppy-logo" not in p.lower():
        p = f"{p}, jppy-logo style"
    return p


@spaces.GPU(duration=90)
def generate(prompt, lora_scale, steps, guidance, width, height, seed, randomize_seed, append_trigger):
    if not prompt or not prompt.strip():
        raise gr.Error("Describe the logo first — e.g. 'a minimalist logo for a coffee shop'.")
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    pipe.set_adapters(["jppy-logo"], adapter_weights=[float(lora_scale)])
    image = pipe(
        prompt=build_prompt(prompt, append_trigger),
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        width=int(width),
        height=int(height),
        generator=torch.Generator("cuda").manual_seed(int(seed)),
    ).images[0]
    return image, int(seed)


with gr.Blocks(title="jppy-logo FLUX.2-klein LoRA demo") as demo:
    gr.Markdown(f"# jppy-logo — logo LoRA demo\n{DESCRIPTION}")
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="a minimalist logo for a coffee shop",
                lines=2,
            )
            append_trigger = gr.Checkbox(value=True, label="Append 'jppy-logo style' trigger")
            run = gr.Button("Generate logo", variant="primary")
            with gr.Accordion("Advanced", open=False):
                lora_scale = gr.Slider(0.0, 1.5, value=1.0, step=0.05, label="LoRA weight (card recommends 0.8–1.0)")
                steps = gr.Slider(8, 50, value=28, step=1, label="Inference steps")
                guidance = gr.Slider(1.0, 7.0, value=3.5, step=0.1, label="Guidance scale")
                width = gr.Slider(512, 1280, value=1024, step=128, label="Width (1024 = crisp vector-friendly)")
                height = gr.Slider(512, 1280, value=1024, step=128, label="Height")
                seed = gr.Number(value=0, precision=0, label="Seed")
                randomize_seed = gr.Checkbox(value=True, label="Randomize seed")
        with gr.Column():
            output = gr.Image(label="Generated logo")
            seed_out = gr.Number(label="Used seed", precision=0)

    run.click(
        generate,
        inputs=[prompt, lora_scale, steps, guidance, width, height, seed, randomize_seed, append_trigger],
        outputs=[output, seed_out],
        api_name="generate",
    )
    prompt.submit(generate,
        inputs=[prompt, lora_scale, steps, guidance, width, height, seed, randomize_seed, append_trigger],
        outputs=[output, seed_out],
        api_name="generate_submit",
    )

    gr.Examples(
        examples=[
            ["a minimalist logo for a coffee shop", 1.0, 28, 3.5, 1024, 1024, 0, True, True],
            ["navy and gold shield monogram logo for a law firm, with text A&L", 1.0, 28, 3.5, 1024, 1024, 0, True, True],
            ["monochrome geometric tech logo, circular badge silhouette", 1.0, 28, 3.5, 1024, 1024, 0, True, True],
            ["red and white minimalist food logo", 0.9, 28, 3.5, 1024, 1024, 0, True, True],
        ],
        inputs=[prompt, lora_scale, steps, guidance, width, height, seed, randomize_seed, append_trigger],
    )

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