# SPDX-License-Identifier: Apache-2.0
"""Viggle Remake: re-render a video with the people (swap) or the scene (edit) changed.

    python inference/sample.py --mode swap --video source.mp4 --prompt prompt.txt \
        --ref ref1.png [--ref ref2.png ...] [--pin 0 [--pin 96 ...]] \
        --frames 124 --short-edge 480 --seed 0 \
        --base MiniMax-H3 --lora-repo Viggle-Remake --out out.mp4

`--video` supplies motion, camera, framing and everything the prompt does not change. Each
`--ref` is an image the prompt refers to as <Picture N>. `--pin` is optional and, when given,
appears once per `--ref`: it pins that ref as a keyframe at that real frame of the output for
the first 90% of the denoise (`--keyframe-until`), which is how a painted frame of the source
becomes an anchor rather than only a reference.

The target audio rows hold the source soundtrack (`--audio`, defaults to `--video`) as a clean
latent for the whole denoise, so mouths follow the real speech and the output carries that
soundtrack back through the audio VAE.

Memory order: the 64-layer Qwen3-VL conditioner (~59 GiB) is loaded, encodes the prompt once,
and is freed before the pipeline loads, so the two never share the card.

Runs on upstream diffusers d6726f3. The keyframe pin (`ref2va` references *plus* keyframes inside
the generated span, the layout the LoRAs are run with) is not in
upstream; the three small patches under "keyframes on a ref2va layout" below add it.
"""

import argparse
import gc
import subprocess
import tempfile
import time
from pathlib import Path

import av
import numpy as np
import torch
from PIL import Image
from transformers import AutoProcessor, AutoTokenizer, Qwen3VLForConditionalGeneration
from diffusers import MiniMaxH3Transformer3DModel, ModularPipeline, attention_backend
from diffusers.modular_pipelines.minimax_h3 import MiniMaxH3ImageReference, MiniMaxH3VideoReference
from diffusers.modular_pipelines.minimax_h3.before_denoise import (_ROPE_FRAME_RESCALE, MiniMaxH3Ref2VAPrepareLayoutStep,
                                                                   MiniMaxH3SetTimestepsStep, _frame_position_grid)
from diffusers.modular_pipelines.minimax_h3.denoise import MiniMaxH3LoopSchedulerStep, MiniMaxH3Ref2VALoopDenoiser
from diffusers.modular_pipelines.minimax_h3.encoders import (MiniMaxH3Ref2VAReferenceEncoderStep,
                                                             MiniMaxH3Ref2VATextEncoderStep, encode_vae_condition,
                                                             get_qwen3vl_prompt_embeds)
from diffusers.modular_pipelines.minimax_h3.modular_pipeline import audio_latent_num_frames
from diffusers.utils.export_utils import encode_video

# 6-step grids, in the video stream's sigma; each stream's own shift is applied below
SIGMAS = {"swap": [1, .990476, .965243, .923077, .8, .631579, 0],
          "edit": [1, .983607, .96, .923077, .857143, .705882, 0]}
TEXT_TAG = 1

parser = argparse.ArgumentParser()
parser.add_argument("--mode", required=True, choices=["swap", "edit"])
parser.add_argument("--video", required=True, help="the source clip")
parser.add_argument("--audio", default=None, help="the soundtrack pinned into the output; defaults to --video")
parser.add_argument("--prompt", required=True, help="text file with the six-section prompt, read verbatim")
parser.add_argument("--ref", action="append", required=True, help="reference image, <Picture N> in the prompt")
parser.add_argument("--pin", action="append", type=int, default=[],
                    help="output frame each --ref is pinned to as a keyframe; absent, or once per --ref")
parser.add_argument("--frames", type=int, default=124, help="17n+5 in 90..345, at 24 fps")
parser.add_argument("--short-edge", type=int, default=480, help="480 (832x480 for 16:9) or 736 (1280x736)")
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--base", required=True, help="local copy of MiniMaxAI/MiniMax-H3")
parser.add_argument("--lora-repo", required=True, help="local copy of Viggle/Viggle-Remake")
parser.add_argument("--rank", type=int, default=256, choices=[256, 128, 64],
                    help="256 loads lora/<mode>; 128 or 64 load the smaller lora_lowrank/<mode>-r<rank>")
parser.add_argument("--keyframe-until", type=float, default=0.9, help="drop the pinned keyframes once 1-t reaches this")
parser.add_argument("--attn-backend", default="_native_cudnn")
parser.add_argument("--out", required=True)
args = parser.parse_args()

nf, se, refs, pin = args.frames, args.short_edge, args.ref, args.pin
assert 90 <= nf <= 345 and (nf - 5) % 17 == 0, f"--frames must be 17n+5 in 90..345, got {nf}"
assert se % 32 == 0, f"--short-edge must be a multiple of 32, got {se}"
assert len(pin) in (0, len(refs)), "give --pin once per --ref, or not at all"
assert all(0 <= k < nf for k in pin), f"--pin must be in [0, {nf})"
prompt = Path(args.prompt).read_text()
sigmas = SIGMAS[args.mode]

# The source on its canvas. A clip already at the canvas (short edge == --short-edge, both sides multiples of 32) is used
# as is; anything else is cover-cropped onto the 16:9-class canvas the way our renderer does (832x480, 1280x736).
with av.open(args.video) as c:
    w, h = c.streams.video[0].codec_context.width, c.streams.video[0].codec_context.height
video_path = args.video
if min(w, h) != se or w % 32 or h % 32:
    long = se * 16 // 9 // 32 * 32
    w, h = (long, se) if w >= h else (se, long)
    video_path = str(Path(tempfile.mkdtemp()) / "cond.mp4")
    subprocess.run(["ffmpeg", "-v", "error", "-i", args.video, "-vf",
                    f"fps=24,scale={w}:{h}:force_original_aspect_ratio=increase,crop={w}:{h}", "-frames:v", str(nf),
                    "-an", "-c:v", "libx264", "-crf", "12", "-pix_fmt", "yuv420p", video_path], check=True)
video = MiniMaxH3VideoReference.from_file(video_path)
video.audio, video.sample_rate = None, None  # the soundtrack goes in as the target audio pin, not as reference audio
height, width = video.frames.shape[1], video.frames.shape[2]

# --- prompt -> embeds. The presentation is text only: a label per 2 fps block of the source, one per ref, then the
# prompt, each piece tokenized on its own. H3 reads hidden_states[50] of the 64-layer conditioner.
t0 = time.time()
processor = AutoProcessor.from_pretrained(args.base, subfolder="processor")
tokenizer = AutoTokenizer.from_pretrained(args.base, subfolder="tokenizer")
text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(args.base, subfolder="text_encoder",
                                                               dtype=torch.bfloat16).to("cuda")
t1 = time.time()
_, stamps = MiniMaxH3Ref2VATextEncoderStep._sample_video_condition_frames(
    np.zeros((nf, 1, 1, 3), dtype=np.uint8), 24.0, 2.0, processor.video_processor.temporal_patch_size)
segments = ["<Video 1>: ", *[f"<{s:.1f} seconds>" for s in stamps], *[f"<Picture {i + 1}>: " for i in range(len(refs))],
            prompt]
token_ids = [i for s in segments for i in tokenizer(s, add_special_tokens=False)["input_ids"]]
with torch.no_grad():
    prompt_embeds = get_qwen3vl_prompt_embeds(text_encoder, processor, token_ids, device=torch.device("cuda"),
                                              dtype=torch.bfloat16).cpu()
text_token_tags = torch.full((len(token_ids),), TEXT_TAG, dtype=torch.long)
t2 = time.time()
print(f"text encoder: load {t1 - t0:.0f}s, encode {t2 - t1:.1f}s, {len(token_ids)} tokens, "
      f"peak {torch.cuda.max_memory_allocated() / 2**30:.1f} GiB", flush=True)
del text_encoder
gc.collect()
torch.cuda.empty_cache()

# --- pipeline: the LoRA sits on the fl2va `transformer`, loaded into the ref2va slot
pipe = ModularPipeline.from_pretrained(args.base, workflow="ref2va")
type(pipe).min_duration = property(lambda self: 90 / 24)  # the adapters run from 90 frames; the base floor is 5 s
pipe.load_components(names=["vae", "audio_vae", "scheduler", "audio_scheduler"], pretrained_model_name_or_path=args.base,
                     dtype=torch.bfloat16)
pipe.transformer_ref = MiniMaxH3Transformer3DModel.from_pretrained(args.base, subfolder="transformer",
                                                                    torch_dtype=torch.bfloat16)
# prefix=None: the saved keys are bare module names; the default "transformer." filter loads nothing and only warns
lora = Path(args.lora_repo) / (f"lora/{args.mode}" if args.rank == 256 else f"lora_lowrank/{args.mode}-r{args.rank}")
pipe.transformer_ref.load_lora_adapter(str(lora),
                                       weight_name="pytorch_lora_weights.safetensors", prefix=None)
pipe.transformer_ref.set_attention_backend(args.attn_backend)  # process-wide: the audio encode below opts out
pipe.to("cuda")
pipe.register_to_config(canvas_short_edge=se, canvas_max_pixels=se * max(height, width), reference_image_short_edge=se)
sh = lambda u, s: s * u / (1 + (s - 1) * u)  # noqa: E731
for sched, s in ((pipe.scheduler, 12.0), (pipe.audio_scheduler, 3.0)):
    grid = [sh(x / (12 - 11 * x), s) for x in sigmas]  # video sigma -> u -> this stream's sigma
    sched.set_timesteps = (lambda f, g: lambda *a, **k: f(sigmas=g, device=k.get("device")))(sched.set_timesteps, grid)
t3 = time.time()
print(f"pipeline on cuda in {t3 - t2:.0f}s", flush=True)


def use_fixed_embeds(self, components, state):
    block_state = self.get_block_state(state)
    block_state.prompt_embeds = prompt_embeds.to(components._execution_device, torch.bfloat16)
    block_state.text_token_tags = text_token_tags
    self.set_block_state(state, block_state)
    return components, state


MiniMaxH3Ref2VATextEncoderStep.__call__ = use_fixed_embeds

# audio pin: the target audio rows are the source soundtrack as a clean latent (positional arg 6 = audio_timestep, and
# H3's t = 1 is clean), and only the video rows are stepped, so they stay clean
_build_row_timesteps = MiniMaxH3SetTimestepsStep.build_row_timesteps
MiniMaxH3SetTimestepsStep.build_row_timesteps = staticmethod(lambda *a: _build_row_timesteps(*a[:6], 1.0, *a[7:]))


@torch.no_grad()
def video_only_step(self, components, block_state, i, t):
    n = block_state.num_condition_video_rows
    block_state.latents[n:] = components.scheduler.step(
        block_state.noise_pred[0, n:].float(), t, block_state.latents[n:], return_dict=False)[0]
    return components, block_state


MiniMaxH3LoopSchedulerStep.__call__ = video_only_step

# --- keyframes on a ref2va layout. A keyframe block is packed after the references, right before the target audio, and
# is pinned *inside* the target's rotary span at its real frame index rather than advancing the clock like a reference.
KF = {"images": [], "anchors": list(pin)}
_encode_references = MiniMaxH3Ref2VAReferenceEncoderStep.__call__


def encode_references_then_keyframes(self, components, state):
    components, state = _encode_references(self, components, state)
    with torch.no_grad():
        kf = [encode_vae_condition(components.vae, torch.from_numpy(np.array(im)).to(components._execution_device)
                                   .permute(2, 0, 1)[None, :, None], components.pixel_mean, components.pixel_std,
                                   components.keyframe_encode_seed) for im in KF["images"]]
    state.set("condition_latents", list(state.get("condition_latents") or []) + kf)
    return components, state


MiniMaxH3Ref2VAReferenceEncoderStep.__call__ = encode_references_then_keyframes
_layout = MiniMaxH3Ref2VAPrepareLayoutStep.build_ref2va_packed_sequence


def layout_with_keyframes(text_token_tags, references, condition_latents, audio_condition_latents, num_latent_frames,
                          latent_height, latent_width, num_audio_latents, patch_size, audio_channels, audio_tag,
                          video_tag):
    k = len(KF["anchors"])
    pos, tags, vi, ai, ti, ncv, nca = _layout(text_token_tags, references, condition_latents[:len(condition_latents) - k],
                                              audio_condition_latents, num_latent_frames, latent_height, latent_width,
                                              num_audio_latents, patch_size, audio_channels, audio_tag, video_tag)
    if not k:
        return pos, tags, vi, ai, ti, ncv, nca
    s = text_token_tags.shape[0] + ncv + nca  # the first target audio row; its time is the target's rotary origin
    origin = float(pos[s, 0])
    blocks = []
    for anchor, lat in zip(KF["anchors"], condition_latents[-k:]):
        grid, _ = _frame_position_grid(lat.shape[3], lat.shape[4], patch_size[1], patch_size[2])
        p = torch.zeros(grid.shape[0], 3, dtype=torch.float64)
        p[:, 0] = origin + _ROPE_FRAME_RESCALE * float(anchor)
        p[:, 1:] = grid
        blocks.append(p)
    n = sum(b.shape[0] for b in blocks)
    pos = torch.cat([pos[:s], *blocks, pos[s:]])
    tags = torch.cat([tags[:s], torch.full((n,), video_tag, dtype=tags.dtype), tags[s:]])
    vi = torch.cat([vi[vi < s], torch.arange(s, s + n), vi[vi >= s] + n])
    ai = torch.where(ai >= s, ai + n, ai)
    return pos, tags, vi, ai, ti, ncv + n, nca


MiniMaxH3Ref2VAPrepareLayoutStep.build_ref2va_packed_sequence = staticmethod(layout_with_keyframes)

# --keyframe-until: once 1 - t reaches it the keyframe blocks are cut out of every per-row input, leaving the plain
# ref2va sequence; their velocity comes back as zeros, which the scheduler never reads (conditioning rows)
r, (_, ph, pw) = pipe.vae_spatial_compression_ratio, pipe.patch_size
nkf = len(pin) * (height // r // ph) * (width // r // pw)
_denoise = MiniMaxH3Ref2VALoopDenoiser.__call__


def kf_release(self, components, block_state, i, t):
    if not nkf or 1 - float(t) >= args.keyframe_until:
        return _denoise(self, components, block_state, i, t)
    ncv = block_state.num_condition_video_rows
    f = block_state.denoiser_input_fields
    vi, ai = f["video_indices"], f["audio_indices"]
    lo, hi = int(vi[ncv - nkf]), int(vi[ncv - 1]) + 1
    assert hi - lo == nkf and hi == int(ai[block_state.num_condition_audio_rows]), (lo, hi, nkf)
    keep = torch.ones(f["position_ids"].shape[0], dtype=torch.bool, device=vi.device)
    keep[lo:hi] = False
    plan = list(block_state.row_timestep_plan)
    plan[i] = (plan[i][0], plan[i][1][keep])
    full = block_state.latents, f, block_state.row_timestep_plan
    block_state.latents = torch.cat([block_state.latents[:ncv - nkf], block_state.latents[ncv:]])
    block_state.denoiser_input_fields = dict(
        f, position_ids=f["position_ids"][keep], token_tags=f["token_tags"][keep],
        video_indices=torch.cat([vi[:ncv - nkf], vi[ncv:] - nkf]), audio_indices=ai - nkf * (ai >= hi))
    block_state.row_timestep_plan = plan
    components, block_state = _denoise(self, components, block_state, i, t)
    block_state.latents, block_state.denoiser_input_fields, block_state.row_timestep_plan = full
    p = block_state.noise_pred
    block_state.noise_pred = torch.cat([p[:, :ncv - nkf], p.new_zeros(1, nkf, p.shape[2]), p[:, ncv - nkf:]], 1)
    return components, block_state


MiniMaxH3Ref2VALoopDenoiser.__call__ = kf_release

# keyframes onto the canvas the way our renderer does: the first one stretched, the rest cover-cropped (rounded
# size, (src - dst) // 2 centring -- not VaeImageProcessor's "crop" arithmetic, which differs by a pixel on some ratios)
if pin:
    for i, im in enumerate(Image.open(x).convert("RGB") for x in refs):
        if im.size != (width, height) and i == 0:
            im = pipe.image_processor.resize(im, height=height, width=width)
        elif im.size != (width, height):
            scale = max(width / im.size[0], height / im.size[1])
            size = (max(width, round(im.size[0] * scale)), max(height, round(im.size[1] * scale)))
            left, top = max(0, (size[0] - width) // 2), max(0, (size[1] - height) // 2)
            im = im.resize(size, Image.Resampling.LANCZOS).crop((left, top, left + width, top + height))
        KF["images"].append(im)

# --- the soundtrack as clean target audio latents
asr = pipe.audio_vae.config.sampling_rate
pcm = subprocess.run(["ffmpeg", "-loglevel", "error", "-i", args.audio or args.video, "-t", f"{nf / 24:.3f}", "-vn",
                      "-f", "f32le", "-ac", "2", "-ar", str(asr), "-"], capture_output=True, check=True).stdout
wav = torch.frombuffer(bytearray(pcm), dtype=torch.float32).view(-1, 2).T
n_samp = round((nf + 1) / 24 * asr)
wav = torch.nn.functional.pad(wav, (0, max(0, n_samp - wav.shape[1])))[:, :n_samp]
with attention_backend(), torch.no_grad():
    posterior = pipe.audio_vae.encode(wav[:, None].to("cuda", torch.bfloat16), return_dict=False)[0]
mean = torch.tensor(pipe.audio_vae.config.latents_mean).view(1, 1, -1)
std = torch.tensor(pipe.audio_vae.config.latents_std).view(1, 1, -1)
rows = (posterior.mode().float().cpu().transpose(1, 2)[:, :audio_latent_num_frames(nf)] - mean) / std
audio_latents = rows.permute(0, 2, 1).contiguous()

torch.cuda.reset_peak_memory_stats()
t4 = time.time()
result = pipe(prompt="".join(segments), references=[video, *[MiniMaxH3ImageReference.from_file(x) for x in refs]],
              num_frames=nf, height=height, width=width, num_inference_steps=len(sigmas) - 1,
              generator=torch.Generator().manual_seed(args.seed), audio_latents=audio_latents,
              output=["videos", "audio", "sampling_rate"])
t5 = time.time()
Path(args.out).parent.mkdir(parents=True, exist_ok=True)
encode_video(result["videos"][0], fps=24, output_path=args.out, audio=result["audio"][0],
             audio_sample_rate=result["sampling_rate"])
print(f"render {t5 - t4:.1f}s at {width}x{height}x{nf}, peak {torch.cuda.max_memory_allocated() / 2**30:.1f} GiB "
      f"allocated / {torch.cuda.max_memory_reserved() / 2**30:.1f} GiB reserved -> {args.out}", flush=True)