import os
import random
import gradio as gr
import torch
import spaces
from diffusers import Flux2KleinPipeline

BASE = "black-forest-labs/FLUX.2-klein-4B"
LORA = "noema-art/impresstation-klein"
# Locked-in look (dialed in: strength 1.05 / steps 9 / guidance 3).
PROMPT = (
    "stationthis, low poly playstation screenshot style. A low-poly PlayStation 2 era 3D model, "
    "lightly upscaled like a PS2 emulator running at 2x internal resolution: keep the chunky "
    "low-polygon geometry and flat-shaded surfaces, with cleaner, moderately sharp edges and "
    "readable textures — but keep a gentle retro softness and a little texture warmth. Slightly "
    "crisp, not razor-sharp; clean but not clinical. Extrude any flat 2D elements into solid 3D "
    "polygonal forms with real depth and foreshortening. A nicely emulated PS2 screenshot — "
    "clearer than the console, softer than 4K."
)
STRENGTH = 1.05
STEPS = 9
GUIDANCE = 3.0
HF_TOKEN = os.environ.get("HF_TOKEN")
MAX_SEED = 2**31 - 1

dtype = torch.bfloat16
pipe = Flux2KleinPipeline.from_pretrained(BASE, torch_dtype=dtype, token=HF_TOKEN)
pipe.load_lora_weights(LORA, adapter_name="stationthis", token=HF_TOKEN)
pipe.set_adapters("stationthis", adapter_weights=STRENGTH)
pipe.to("cuda")


@spaces.GPU(duration=120)
def station_this(image, progress=gr.Progress(track_tqdm=True)):
    if image is None:
        raise gr.Error("Drop in an image first.")
    seed = random.randint(0, MAX_SEED)
    g = torch.Generator("cuda").manual_seed(seed)
    out = pipe(
        image=[image.convert("RGB")],
        prompt=PROMPT,
        guidance_scale=GUIDANCE,
        num_inference_steps=STEPS,
        generator=g,
    ).images[0]
    return out


CSS = """
:root { --accent:#5b8cff; }
body, .gradio-container { background:#08090A; color:#e7eaef; }
#title { text-align:center; font-family:ui-monospace,monospace; letter-spacing:.06em; }
#title b { color:var(--accent); }
.gr-button-primary { background:var(--accent) !important; border:0 !important; color:#06080f !important; }
footer { display:none !important; }
"""

with gr.Blocks(css=CSS, title="stationthis") as demo:
    gr.HTML(
        "<div id='title'><h1>⭐👌 stationthis</h1>"
        "<p>Drop an image — get it back as a low-poly PlayStation-era still. "
        "Powered by FLUX.2 klein + <b>NOEMA</b>. "
        "<a href='https://noema.art'>noema.art</a></p></div>"
    )
    with gr.Row():
        with gr.Column():
            inp = gr.Image(label="Your image", type="pil", height=380)
            go = gr.Button("Station this", variant="primary")
        with gr.Column():
            out = gr.Image(label="result", height=380)
    go.click(station_this, inp, out)

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