"""Bang Story image Space: draws clip pictures on Hugging Face ZeroGPU.

Upload this file together with requirements.txt and README.md into a Gradio Space, set the Space's
hardware to ZeroGPU, and Bang Story can use it as its image service — free, with no credits and no
API keys. The app calls the `generate` function through the Space's Gradio API.
"""

import os

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

# Tried in order, first one that loads wins. The official repository gives the best quality but is gated:
# without the licence accepted and a read token in the HF_TOKEN secret it answers 401. The open mirrors
# below are the same model, so the Space still comes up instead of dying on a missing token.
MODEL_CHAIN = [
    os.environ.get("MODEL_ID") or "black-forest-labs/FLUX.1-schnell",
    "Niansuh/FLUX.1-schnell",
    "playgroundai/playground-v2.5-1024px-aesthetic",
]

# A distilled model like schnell needs very few steps; a full pipeline needs many more.
MAX_STEPS = 8
device = "cuda" if torch.cuda.is_available() else "cpu"


def load():
    """Loads the first model of the chain that works, so one blocked repository cannot kill the Space."""
    last = None
    for model in MODEL_CHAIN:
        if not model:
            continue
        try:
            print(f"[bang-story] loading {model} ...", flush=True)
            pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.bfloat16).to(device)
            print(f"[bang-story] ready: {model} on {device}", flush=True)
            return pipe, model
        except Exception as e:
            last = e
            print(f"[bang-story] {model} failed: {type(e).__name__}: {e}", flush=True)
    raise RuntimeError(f"Tidak ada model yang bisa dimuat. Error terakhir: {last}")


pipe, MODEL = load()
IS_FLUX = "flux" in MODEL.lower()


@spaces.GPU()
def generate(prompt: str, width: int = 1024, height: int = 576, steps: int = 4):
    """Draws one picture. Bang Story calls this through the Gradio API."""
    kwargs = {"prompt": prompt, "width": int(width), "height": int(height)}
    if IS_FLUX:
        kwargs["num_inference_steps"] = max(1, min(MAX_STEPS, int(steps)))
        # schnell is distilled: it runs without classifier-free guidance.
        kwargs["guidance_scale"] = 0.0
    else:
        kwargs["num_inference_steps"] = max(20, min(35, int(steps) * 5))
        kwargs["guidance_scale"] = 7.0
    return pipe(**kwargs).images[0]


with gr.Blocks(title="Bang Story — Image") as demo:
    gr.Markdown(
        f"Draws the clip pictures for **Bang Story** with `{MODEL}`.\n\n"
        "Bang Story calls this Space from its Settings, so you normally never touch this page. "
        "It is left open so you can try a prompt yourself."
    )
    prompt = gr.Textbox(label="Prompt", lines=5, placeholder="Describe the picture...")
    with gr.Row():
        width = gr.Slider(256, 1536, value=1024, step=16, label="Width")
        height = gr.Slider(256, 1536, value=576, step=16, label="Height")
        steps = gr.Slider(1, MAX_STEPS, value=4, step=1, label="Steps")
    run = gr.Button("Generate", variant="primary")
    out = gr.Image(label="Result", type="pil")
    run.click(generate, [prompt, width, height, steps], out)
    prompt.submit(generate, [prompt, width, height, steps], out)


demo.launch()