"""SolarRenderer: few-step generative reference rendering on MiniMax-H3.

Structure video in, rendered video out: the released SolarRenderer LoRA pair (a four-step sampler
LoRA merged into the weights, a rank-384 control LoRA on top of the Ref2VA transformer) renders
what the structure video describes, guided by a six-part prompt and anchored on an optional first
frame. Sampling takes **four** model evaluations.

The 62 GiB Qwen3-VL text conditioner runs in a separate Space (`multimodalart/qwen3vl-conditioner`);
the wire format between the halves is a safetensors file of `prompt_embeds` + `text_token_tags`.
Everything else — the reference/target VAE encoding and the denoising itself — is this Space's, and
follows the authors' `scripts/render.py` and `tools/omni_preencode.py` step for step: the structure
video is encoded as an observation laid over every target frame, the first frame as a held still on
latent 0, and the target (the structure clip itself, for shape and length only) through the released
`ref2va` `latents_only` chain. The member safetensors the renderer reads is written exactly as
`omni_preencode.encode()` writes it.
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  (before torch: it patches torch.cuda)

import json
import math
import shutil
import sys
import tempfile
import time
from functools import cache
from pathlib import Path

import gradio as gr
import numpy as np
import torch

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE / "src"))
sys.path.insert(0, str(HERE / "tools"))

VERSION = "solarrenderer-v1"
FRAMES = 158  # the encoded length the release was trained at: 47 latents, 6.58 s at 24 fps
RANK = ALPHA = 384
FPS = 24.0
STILL_FRAMES = 13  # H3 takes no video shorter than this; a first frame is held still for it
CONDITIONER_SPACE = os.environ.get("CONDITIONER_SPACE", "multimodalart/qwen3vl-conditioner")
H3_ROOT = os.environ.get("H3_ROOT", "MiniMaxAI/MiniMax-H3")
LORA_REPO = os.environ.get("LORA_REPO", "Enderfga/SolarRenderer")
GPU_SIZE = os.environ.get("GPU_SIZE", "xlarge")  # 61.7 GiB DiT + fp32 VAE does not fit "large"
GPU_DURATION = int(os.environ.get("GPU_DURATION", "460"))  # measured 393 s (85 placement + 32 encode + 241 sample + 17 decode + moves) + 15%
CANVASES = {  # the conditioner's labels, one per target canvas
    (768, 1344): "1344x768 · 16:9 full",
    (1344, 768): "768x1344 · 9:16 full",
}
SILENCE = HERE / "assets" / "h3_silence_153_158_170.safetensors"


# -------------------------------------------------------------------------------------------- startup load

def _download_loras() -> tuple[Path, Path]:
    """The released non-causal LoRA pair, fetched once into the Space's cache."""

    from huggingface_hub import hf_hub_download

    return (
        Path(hf_hub_download(LORA_REPO, "solarrenderer_4step_lora.safetensors")),
        Path(hf_hub_download(LORA_REPO, "solarrenderer_control_lora.safetensors")),
    )


def load_model() -> dict:
    """Everything except the Qwen3-VL conditioner, on the CPU at startup.

    The transformer is 61.7 GiB and the two VAE copies another ~15 GiB — more than ZeroGPU's
    startup pack can hold — so placement happens on the card inside the first GPU call instead,
    exactly as the deployed H3 Spaces do. The 4-step LoRA is merged into the weights and the
    control LoRA injected as a PEFT adapter before placement, so the card only ever sees the final
    weights. The `ref2va` `latents_only` encoding chain is built here too, but **without** its
    text encoder: only the `before_encode` + `vae_encoder` blocks, whose single pretrained weight
    is the bf16 video VAE.
    """

    from diffusers import ComponentsManager, ModularPipeline
    from diffusers.modular_pipelines import SequentialPipelineBlocks
    from solarwm.backends.minimax_h3.artifacts import load_silence_latents
    from solarwm.backends.minimax_h3.lora import inject_h3_lora, merge_lora_delta
    from solarwm.backends.minimax_h3.observations import observation_profile_from_train_config
    from solarwm.backends.minimax_h3.optional import load_conditioners, load_transformer
    from solarwm.backends.minimax_h3.stage0p5 import H3Stage0p5Core

    began = time.perf_counter()
    model_cfg = {
        "checkpoint_path": H3_ROOT,
        "transformer_subfolder": "transformer_ref",
        "torch_dtype": "bfloat16",
        # flash-attn has no sm_120 (Blackwell) kernel; cuDNN's SDPA is what the deployed H3 Spaces
        # run on this hardware. Stamped first, and the VAEs are stamped "native" after it — the
        # registry's active backend is global.
        "attention_backend": "_native_cudnn",
        "adapter": {"rank": RANK, "alpha": ALPHA, "dropout": 0.0},
        "low_cpu_mem_usage": True,
    }
    modules = load_transformer(model_cfg, device=torch.device("cpu"))

    def lora_pairs(path: str) -> tuple[dict, float]:
        from safetensors import safe_open
        from safetensors.torch import load_file

        with safe_open(path, framework="pt") as handle:
            meta = handle.metadata() or {}
        return load_file(path), float(meta.get("alpha", ALPHA))

    fewstep, control = _download_loras()
    pairs, alpha = lora_pairs(str(fewstep))
    rank = int(next(v for k, v in pairs.items() if k.endswith("lora_A.weight")).shape[0])
    merged = merge_lora_delta(modules.transformer, pairs, alpha=alpha, rank=rank, scale=1.0)
    print(f"[{VERSION}] four-step LoRA merged into {merged} weights", flush=True)
    wrapped, lora = inject_h3_lora(modules.transformer, model_cfg["adapter"], base_identity={})
    pairs, _alpha = lora_pairs(str(control))
    lora.load_state_dict({f"base_model.model.{k}": v for k, v in pairs.items()}, broadcast=False)
    print(f"[{VERSION}] control LoRA loaded into the adapter", flush=True)

    conditioners = load_conditioners(
        model_cfg, device=torch.device("cpu"), qwen=False, video_vae=True, audio_vae=False, schedulers=False
    )
    conditioners.video_vae.set_attention_backend("native")  # the fp32 decoder has no cuDNN kernel

    # The authors' `latents_only` chain, built by hand so the 62 GiB Qwen text encoder that
    # `omni_preencode.build_pipelines` also loads never enters this Space.
    blocks = ModularPipeline.from_pretrained(H3_ROOT).blocks
    sub = blocks.get_workflow("ref2va").sub_blocks
    latents_only = SequentialPipelineBlocks.from_blocks_dict(
        {name: sub[name] for name in ("before_encode", "vae_encoder")}
    ).init_pipeline(H3_ROOT, components_manager=ComponentsManager())
    latents_only.load_components(dtype=torch.bfloat16)
    latents_only.vae.set_attention_backend("native")
    for name in latents_only.pretrained_component_names:
        component = getattr(latents_only, name, None)
        if isinstance(component, torch.nn.Module):
            component.eval().requires_grad_(False)

    silence, _ = load_silence_latents(str(SILENCE), pixel_frames=FRAMES)
    core = H3Stage0p5Core(
        wrapped, silence, torch.device("cpu"), geometry=None,
        observations=observation_profile_from_train_config({}), flow_map=False,
    )
    print(f"[{VERSION}] model loaded on CPU in {time.perf_counter() - began:.0f} s", flush=True)
    return {
        "core": core,
        "conditioners": conditioners,
        "latents_only": latents_only,
        "placed": False,
        "vae": None,
    }


MODEL = load_model()


# -------------------------------------------------------------------------------------------- media plumbing

def canvas_of(path: str) -> tuple[int, int]:
    """The canvas a clip is encoded on: landscape or portrait by its frame's shape.

    The same rule as `omni_preencode.canvas` — the part of the frame that is kept decides.
    """

    import omni_preencode as pre

    stream = pre.open_media(path).streams.video[0]
    width, height = int(stream.width), int(stream.height)
    return (768, 1344) if width >= height else (1344, 768)


def read_video(path: str, seconds: float = 0.0) -> tuple[np.ndarray, float]:
    """Decode a clip — only its first `seconds` when given — exactly as `omni_preencode.read_video`."""

    import omni_preencode as pre

    container = pre.open_media(path)
    stream = container.streams.video[0]
    stream.thread_type = "AUTO"
    fps = float(stream.average_rate)
    limit = int(math.ceil(seconds * fps)) + 2 if seconds else 0
    frames = []
    for frame in container.decode(stream):
        frames.append(frame.to_ndarray(format="rgb24"))
        if limit and len(frames) >= limit:
            break
    container.close()
    if len(frames) < 2:
        raise ValueError("the structure video is too short: SolarRenderer renders a 6.6 s clip")
    return np.stack(frames), fps


def still_video(first: str, canvas: tuple[int, int], out: str) -> str:
    """The first frame as the 13-frame 24 fps video H3's reference blocks take.

    `omni_preencode.reference` lays an aligned still over the canvas and repeats it `STILL_FRAMES`
    times; pre-encoding it here means the same still reaches both the conditioner (over its API)
    and the local VAE encoder pixel for pixel.
    """

    import av
    from PIL import Image

    import omni_preencode as pre

    image = Image.open(first).convert("RGB")
    frame = np.asarray(pre.aligned_image(image, canvas))[None]
    frames = np.repeat(frame, STILL_FRAMES, axis=0)
    height, width = canvas
    with av.open(out, "w") as container:
        stream = container.add_stream("libx264", rate=int(FPS))
        stream.width, stream.height = width, height
        stream.pix_fmt = "yuv444p"
        stream.options = {"crf": "0", "preset": "veryfast"}  # near-lossless; identical frames either way
        for pixels in frames:
            picture = av.VideoFrame.from_ndarray(pixels, format="rgb24")
            for packet in stream.encode(picture):
                container.mux(packet)
        for packet in stream.encode():
            container.mux(packet)
    return out


# -------------------------------------------------------------------------------------------- the conditioner half

@cache
def conditioner():
    """The Qwen3-VL half, running in its own Space so this one never loads the 62 GiB encoder."""

    from gradio_client import Client

    return Client(CONDITIONER_SPACE)


def encode_remote(prompt: str, structure: str, first_video: str | None, canvas: tuple[int, int]) -> tuple[Path, dict]:
    """`prompt_embeds` + `text_token_tags` for the request, from the conditioner Space.

    H3 numbers references per modality: both of this demo's references are videos, so the structure
    clip is `<Video 1>` and the first-frame still is `<Video 2>` — the labels the six-part prompt's
    `subject_definitions` section refers to. The returned plan carries the resolved canvas, which the
    generating half pins rather than re-deriving. The conditioner's own GPU booking is charged to
    this Space's visitors, which is why the call sits outside this Space's `@spaces.GPU` time.
    """

    from gradio_client import handle_file

    media = [handle_file(structure)] + ([handle_file(first_video)] if first_video else [])
    kinds = "video,video" if first_video else "video"
    path, plan = conditioner().predict(
        prompt=prompt,
        media=media,
        kinds=kinds,
        canvas=CANVASES[canvas],
        num_frames=FRAMES,
        rewrite_prompt=False,
        api_name="/encode_ref2va",
    )
    return Path(path), plan


# -------------------------------------------------------------------------------------------- the generating half

def vram() -> str:
    """The allocator's own numbers, for the log."""

    if not torch.cuda.is_available():
        return "no cuda"
    return (
        f"allocated {torch.cuda.memory_allocated() / 2**30:.1f} GiB, "
        f"peak {torch.cuda.max_memory_allocated() / 2**30:.1f} GiB, "
        f"reserved {torch.cuda.memory_reserved() / 2**30:.1f} GiB"
    )


def place_on_card() -> None:
    """Move the transformer onto the card once, inside the GPU call (the startup pack cannot hold it).

    Only the PEFT-wrapped DiT goes on here. **The two autoencoders are never resident at the same
    time**: the encode chain's bf16 VAE is needed only for `encode_member` (which runs first) and the
    fp32 decode VAE only for `decode_h3_latents` (which runs last), so each is placed for its own
    phase and released before the next. Both on the card at once is ~14.5 GiB of weights plus both
    phases' activations sitting under the sampling peak, which is what pushes the forward pass into
    the allocator's NVML assertion on a 96 GiB card.
    """

    core = MODEL["core"]
    if MODEL["placed"]:
        return
    began = time.perf_counter()
    cuda = torch.device("cuda")
    core.model.to(cuda)  # the PEFT wrapper holds the merged base + the control adapter
    core.silence = core.silence.to(cuda)
    core.device = cuda
    core._layouts.clear()  # layouts cached with tensors on the old device
    torch.cuda.empty_cache()
    MODEL["placed"] = True
    print(
        f"[{VERSION}] transformer placed on card in {time.perf_counter() - began:.0f} s ({vram()})",
        flush=True,
    )


def place_vae(which: str) -> None:
    """Put at most one of the two autoencoders on the card, moving the others off first.

    `which="encode"` is the bf16 VAE of the `latents_only` chain, `which="decode"` the fp32 decoder
    `load_conditioners` builds, `"none"` leaves both on the CPU. Requesting the state already in
    place is a no-op, so the (relatively slow) transfers only happen on an actual change.
    """

    if MODEL["vae"] == which:
        return
    began = time.perf_counter()
    cuda = torch.device("cuda")
    encode_vae = MODEL["latents_only"].vae
    decode_vae = MODEL["conditioners"].video_vae
    encode_vae.to(cuda if which == "encode" else "cpu")
    decode_vae.to(cuda if which == "decode" else "cpu")
    torch.cuda.empty_cache()
    MODEL["vae"] = which
    print(
        f"[{VERSION}] autoencoders now {which!r} in {time.perf_counter() - began:.0f} s ({vram()})",
        flush=True,
    )


def encode_member(prompt: str, structure: str, first: str | None, plan: dict, canvas: tuple[int, int], work: Path) -> Path:
    """The generating half of the authors' `encode`: target and observation latents, member file.

    The target goes through the `latents_only` chain on its own, as `omni_preencode.encode` does —
    it sets the canvas and the length, and sampling never reads its values. `height` / `width` are
    pinned from the conditioner's plan: the two halves resolve the same canvas, and pinning is what
    keeps them from disagreeing.
    """

    from safetensors import safe_open
    from safetensors.torch import save_file

    import omni_preencode as pre

    latents_only = MODEL["latents_only"]
    height, width = int(plan["height"]), int(plan["width"])
    size = {"height": height, "width": width, "num_frames": FRAMES}
    shape = (height, width)
    video = read_video(structure, seconds=FRAMES / FPS + 0.5)
    clip_frames, clip_fps = int(video[0].shape[0]), float(video[1])
    target_state = latents_only(
        references=[pre.reference(structure, shape, video, num_frames=FRAMES)],
        prompt="",
        **size,
    )
    del video
    target = target_state.get("condition_latents")[0][0]  # [24, F, h, w]
    frames = int(target.shape[1])
    # `omni_preencode`: the container's own count wins, since only the clip's start was decoded.
    whole_frames = max(pre.frame_count(structure), clip_frames)
    tensors = {"target_latents": target.to(torch.bfloat16).contiguous()}
    # The structure is laid over every target frame and the target is the structure clip itself, so
    # its observation latents are the target's own (omni_preencode: aligned="all" over the same clip).
    tensors["observation_0_latents"] = target.to(torch.bfloat16).contiguous()
    observations = [{"frames": list(range(frames)), "degree": 0.999, "boxes": None}]
    if first:
        still = still_video(first, shape, str(work / "first_still.mp4"))
        state = latents_only(
            references=[pre.reference(still, shape, video=None, num_frames=FRAMES)],
            prompt="",
            **size,
        )
        latents = state.get("condition_latents")[0][0]
        # An aligned still keeps its first latent alone (omni_preencode: frames=[0], latents[:, :1]).
        tensors["observation_1_latents"] = latents[:, :1].to(torch.bfloat16).contiguous()
        observations.append({"frames": [0], "degree": 0.999, "boxes": None})
    # prompt_embeds / text_token_tags cross from the conditioner half's safetensors.
    with safe_open(str(work / "prompt.safetensors"), framework="pt") as handle:
        tensors["prompt_embeds"] = (
            handle.get_tensor("prompt_embeds").reshape(-1, 5120).to(torch.bfloat16).contiguous()
        )
        tensors["text_token_tags"] = (
            handle.get_tensor("text_token_tags").reshape(-1).to(torch.int64).contiguous()
        )
    meta = {
        "observations": observations,
        "latent_frames": frames,
        "canvas": [height, width],
        # `omni_preencode`: clip_seconds covers the whole clip as the container counted it, and
        # window_seconds the 158-frame stretch the latents hold.
        "clip_seconds": round(whole_frames / clip_fps, 3),
        "start_seconds": 0.0,
        "window_seconds": round(FRAMES / FPS, 3),
        "source": "demo",
        "crop": None,
        "target": structure,
        "media": [structure] + ([first] if first else []),
        "known_mask": "",
        "text": prompt[:16000],
    }
    member = work / "member.safetensors"
    save_file(tensors, str(member), metadata={"omni": json.dumps(meta)})
    return member


@spaces.GPU(duration=GPU_DURATION, size=GPU_SIZE)
def _generate(prompt: str, structure: str, first: str | None, prompt_file: str, height: int, width: int, seed: int, forwards: int) -> str:
    """The only thing on GPU time: placement, the two reference encodes, sampling, decode, mux.

    The prompt was already encoded by the conditioner Space, outside this booking; it crosses as
    the safetensors path `prompt_file`, with the plan's resolved `height` / `width` as plain ints.
    """

    from solarwm.backends.minimax_h3.inference import _mp4_bytes, decode_h3_latents
    from solarwm.backends.minimax_h3.omni_data import H3OmniStream

    core, conditioners = MODEL["core"], MODEL["conditioners"]
    place_on_card()
    with tempfile.TemporaryDirectory() as scratch:
        work = Path(scratch)
        shutil.copy(prompt_file, work / "prompt.safetensors")
        plan = {"height": int(height), "width": int(width)}
        place_vae("encode")  # the bf16 encode VAE, and only it, while the references are encoded
        began = time.perf_counter()
        member = encode_member(prompt, structure, first, plan, (int(height), int(width)), work)
        print(f"[{VERSION}] references encoded in {time.perf_counter() - began:.0f} s ({vram()})", flush=True)
        stream = object.__new__(H3OmniStream)
        stream.root, stream.current_plan_fingerprint, stream.aligned_control = "/", "render", False
        batch = stream.read({"id": "render", "member": str(member)})
        # Both autoencoders leave the card before the sampling peak: the encode VAE has done its
        # work and the decode VAE is not needed until the latents exist.
        place_vae("none")
        torch.cuda.reset_peak_memory_stats()
        began = time.perf_counter()
        with torch.no_grad():
            # generate() counts schedule points: N points run N - 1 evaluations
            latents = core.generate(batch, noise_seed=seed, num_inference_steps=forwards + 1, renoise=True)
        print(f"[{VERSION}] sampled {forwards} steps in {time.perf_counter() - began:.0f} s ({vram()})", flush=True)
        place_vae("decode")  # the fp32 decoder, and only it, for the decode
        began = time.perf_counter()
        decoded = decode_h3_latents(conditioners.video_vae, latents, device=torch.device("cuda"))
        out = work / "render.mp4"
        out.write_bytes(_mp4_bytes(decoded))
        print(f"[{VERSION}] decoded + muxed in {time.perf_counter() - began:.0f} s ({vram()})", flush=True)
        final = Path(tempfile.gettempdir()) / f"solarrenderer_{seed}_{int(time.time())}.mp4"
        shutil.copy(out, final)
    return str(final)


# -------------------------------------------------------------------------------------------- the Gradio app

def render(prompt: str, structure, first=None, seed: int = 0, forwards: int = 4, progress=gr.Progress(track_tqdm=True)):
    """Render a video that follows the structure clip.

    Args:
        prompt: the six-part prompt (`subject_definitions` … `non_diegetic_music`), as the release
            captions with. `<Video 1>` is the structure video, `<Video 2>` the first frame when given.
        structure: the structure video — a semantic map, a drawing, a pose skeleton, a game-engine
            proxy — landscape or portrait, 24 fps, at least 158 frames (6.6 s).
        first: an optional first frame; the render begins from it, held still.
        seed: the sampling seed.
        forwards: model evaluations; the four-step LoRA was trained for 4.

    Returns:
        The rendered video: 158 frames, 24 fps, the structure clip's canvas.
    """

    path = structure.get("video") if isinstance(structure, dict) else structure
    if not path or not Path(str(path)).exists():
        raise gr.Error("a structure video is required")
    if not prompt or not prompt.strip():
        raise gr.Error("a prompt is required")
    structure = str(path)
    canvas = canvas_of(structure)
    first_path = None
    if first:
        candidate = first.get("video") if isinstance(first, dict) else first
        if candidate and Path(str(candidate)).exists():
            first_path = str(candidate)
    began = time.perf_counter()
    with tempfile.TemporaryDirectory() as scratch:
        work = Path(scratch)
        still = still_video(first_path, canvas, str(work / "first_for_conditioner.mp4")) if first_path else None
        try:
            prompt_path, plan = encode_remote(prompt, structure, still, canvas)
        except Exception as error:
            raise gr.Error(
                f"The conditioner ({CONDITIONER_SPACE}) failed with `{type(error).__name__}: {error}`. "
                "Try again in a moment."
            ) from error
        print(f"[{VERSION}] prompt encoded remotely in {time.perf_counter() - began:.0f} s", flush=True)
        # The plan's resolved height / width cross as plain ints: @spaces.GPU arguments pickle
        # across processes, and the safetensors itself crosses as the path alone.
        return _generate(
            prompt, structure, first_path, str(prompt_path), int(plan["height"]), int(plan["width"]), int(seed), int(forwards)
        )


CSS = """
.main.fillable { max-width: 1250px !important; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="SolarRenderer") as demo:
    gr.Markdown(
        """
        # SolarRenderer: few-step generative reference rendering
        Render a video from a **structure video** (semantic map, drawing, pose skeleton, engine
        proxy) and a six-part prompt, anchored on an optional first frame — in **four** model
        evaluations, on the MiniMax-H3 Ref2VA transformer with the released four-step + control LoRA
        pair. [Paper](https://huggingface.co/papers/2610.28630) ·
        [Code](https://github.com/Enderfga/SolarRenderer) ·
        [Weights](https://huggingface.co/Enderfga/SolarRenderer)
        """
    )
    with gr.Row():
        with gr.Column():
            structure = gr.Video(label="Structure video", sources=["upload"])
            gr.Markdown(
                "*Landscape clips render at 1344x768, portrait at 768x1344. Keep it at 24 fps and at "
                "least 158 frames (6.6 s) — the first 158 frames are used.*"
            )
            first = gr.Image(label="First frame (optional)", type="filepath", sources=["upload"])
            prompt = gr.Textbox(
                label="Prompt (six parts)",
                placeholder="subject_definitions:\n<Video 1> is …\n\nsummary:\n…",
                lines=14,
            )
            with gr.Accordion("Advanced", open=False):
                seed = gr.Slider(0, 2**31 - 1, value=0, step=1, label="Seed")
                forwards = gr.Slider(1, 8, value=4, step=1, label="Model evaluations",
                                     info="The four-step LoRA was trained for 4.")
        with gr.Column():
            video = gr.Video(label="Rendered video")
            button = gr.Button("Render", variant="primary")
    button.click(
        render, inputs=[prompt, structure, first, seed, forwards], outputs=[video],
        concurrency_limit=1, api_name="render",
    )
    gr.Examples(
        examples=[
            [
                Path("examples/sem-highway_synthwave/prompt.txt").read_text(),
                "examples/sem-highway_synthwave/structure.mp4",
                "examples/sem-highway_synthwave/first.png",
            ]
        ],
        inputs=[prompt, structure, first],
        fn=render,
        outputs=[video],
        cache_examples=True,
        cache_mode="lazy",
    )

demo.launch(show_error=True, theme=gr.themes.Citrus(), css=CSS, max_threads=1000)