"""Simple text-to-video interface for krea/krea-realtime-video, tuned for Apple Silicon (MPS).

This is *not* the realtime/streaming server from the model card (that one assumes CUDA,
flash_attn, and a B200 GPU). Instead it drives the diffusers ModularPipeline directly,
running the standard (non-realtime) block-by-block denoising loop on MPS, and exposes
the generation knobs the pipeline actually supports through a small Gradio UI.

Run with:
    uv run main.py
"""

import os
import time

import gradio as gr
import torch
from diffusers import ModularPipeline
from diffusers.modular_pipelines import PipelineState
from diffusers.utils import export_to_video

REPO_ID = "krea/krea-realtime-video"
OUTPUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "outputs")

DTYPE_CHOICES = {
    "bfloat16 (recommended on M3/M4)": torch.bfloat16,
    "float16 (safer on M1/M2)": torch.float16,
    "float32 (slow, most compatible)": torch.float32,
}

_pipe = None


def get_device() -> str:
    if torch.backends.mps.is_available():
        return "mps"
    print("MPS not available, falling back to CPU (this will be very slow).")
    return "cpu"


def load_pipeline(dtype: torch.dtype, progress: gr.Progress | None = None):
    """Load (and cache) the modular pipeline. Downloads ~28GB of weights on first run."""
    global _pipe
    if _pipe is not None and _pipe._dtype == dtype:
        return _pipe

    if progress is not None:
        progress(0, desc="Loading pipeline (first run downloads ~28GB from Hugging Face)...")

    device = get_device()
    pipe = ModularPipeline.from_pretrained(REPO_ID, trust_remote_code=True)
    pipe.load_components(
        trust_remote_code=True,
        device_map=device,
        torch_dtype={"default": dtype, "vae": torch.float32},
    )
    for block in pipe.transformer.blocks:
        block.self_attn.fuse_projections()

    pipe._dtype = dtype
    _pipe = pipe
    return pipe


def generate_video(
    prompt: str,
    negative_prompt: str,
    height: int,
    width: int,
    num_inference_steps: int,
    num_blocks: int,
    num_frames_per_block: int,
    fps: int,
    seed: int,
    dtype_label: str,
    lora_repo: str,
    lora_weight_name: str,
    lora_adapter_name: str,
    progress=gr.Progress(),
):
    if not prompt or not prompt.strip():
        raise gr.Error("Prompt cannot be empty.")

    dtype = DTYPE_CHOICES[dtype_label]
    pipe = load_pipeline(dtype, progress)

    if lora_repo.strip():
        progress(0, desc=f"Loading LoRA from {lora_repo}...")
        pipe.transformer.load_lora_adapter(
            lora_repo.strip(),
            prefix="diffusion_model",
            weight_name=lora_weight_name.strip() or None,
            adapter_name=(lora_adapter_name.strip() or "user_lora"),
        )

    device = get_device()
    seed = int(seed)
    if seed < 0:
        seed = torch.seed() % (2**32)
    generator = torch.Generator(device=device if device != "mps" else "cpu").manual_seed(seed)

    frames = []
    state = PipelineState()
    start = time.time()
    for block_idx in range(num_blocks):
        progress(block_idx / num_blocks, desc=f"Denoising block {block_idx + 1}/{num_blocks}...")
        state = pipe(
            state,
            prompt=[prompt],
            negative_prompt=[negative_prompt] if negative_prompt.strip() else None,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            num_blocks=num_blocks,
            num_frames_per_block=num_frames_per_block,
            block_idx=block_idx,
            generator=generator,
        )
        frames.extend(state.values["videos"][0])

    elapsed = time.time() - start
    progress(1.0, desc="Exporting video...")

    os.makedirs(OUTPUT_DIR, exist_ok=True)
    out_path = os.path.join(OUTPUT_DIR, f"t2v_seed{seed}_{int(time.time())}.mp4")
    export_to_video(frames, out_path, fps=fps)

    info = f"Seed: {seed} | {len(frames)} frames @ {fps}fps | generated in {elapsed:.1f}s on {device}"
    return out_path, info


with gr.Blocks(title="Krea Realtime Video - Text to Video") as demo:
    gr.Markdown(
        "# Krea Realtime Video — Text to Video\n"
        "Local text-to-video with `krea/krea-realtime-video` (Wan2.1 14B distilled) via 🧨 diffusers, "
        "running on Apple Silicon (MPS). First generation downloads ~28GB of weights."
    )

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(label="Prompt", placeholder="a cat sitting on a boat", lines=3)
            negative_prompt = gr.Textbox(label="Negative prompt (optional)", lines=2)

            with gr.Row():
                height = gr.Slider(256, 720, value=480, step=16, label="Height")
                width = gr.Slider(256, 1280, value=832, step=16, label="Width")

            with gr.Row():
                num_inference_steps = gr.Slider(1, 12, value=6, step=1, label="Inference steps per block")
                fps = gr.Slider(8, 30, value=24, step=1, label="Output FPS")

            with gr.Row():
                num_blocks = gr.Slider(1, 30, value=9, step=1, label="Number of blocks (length)")
                num_frames_per_block = gr.Slider(1, 12, value=3, step=1, label="Frames per block")

            with gr.Row():
                seed = gr.Number(value=-1, label="Seed (-1 = random)", precision=0)
                dtype_label = gr.Dropdown(
                    choices=list(DTYPE_CHOICES.keys()),
                    value="bfloat16 (recommended on M3/M4)",
                    label="Compute dtype",
                )

            with gr.Accordion("LoRA (optional)", open=False):
                lora_repo = gr.Textbox(label="LoRA repo id", placeholder="shauray/Origami_WanLora")
                lora_weight_name = gr.Textbox(
                    label="LoRA weight filename", placeholder="origami_000000500.safetensors"
                )
                lora_adapter_name = gr.Textbox(label="Adapter name", placeholder="origami")

            generate_btn = gr.Button("Generate", variant="primary")

        with gr.Column(scale=1):
            video_out = gr.Video(label="Result")
            info_out = gr.Textbox(label="Run info", interactive=False)

    generate_btn.click(
        fn=generate_video,
        inputs=[
            prompt,
            negative_prompt,
            height,
            width,
            num_inference_steps,
            num_blocks,
            num_frames_per_block,
            fps,
            seed,
            dtype_label,
            lora_repo,
            lora_weight_name,
            lora_adapter_name,
        ],
        outputs=[video_out, info_out],
    )


def main():
    demo.queue().launch()


if __name__ == "__main__":
    main()