# SPDX-License-Identifier: Apache-2.0
"""Viggle-Animate: replace the performers in a video with the person in a still.

    python inference/sample.py --cond driving.mp4 --ref character.png --out swapped.mp4

`--cond` supplies the motion, camera framing, background and lighting; `--ref` supplies
who is in it. Everything about the render except the people is copied from `--cond`.

By default the model generates its own soundtrack, and the fixed prompt asks for silence --
so there is nothing for the mouth to sync to. `--audio` pins a real one instead:

    python inference/sample.py --cond driving.mp4 --ref character.png \
        --audio driving.mp4 --out swapped.mp4

The soundtrack is encoded once and held in the target audio rows as a *clean* latent for the
whole denoise, so the model conditions on it rather than predicting it, and the mouth tracks
that speech. The track written to `--out` is then that same audio, back through the audio VAE.

The text encoder is never loaded. Conditioning comes from `assets/fixed_embed_fwd_anyframe.pt`,
a frozen 362 x 5120 tensor computed once from the fixed prompt in `assets/fixed_prompt.txt`,
so Qwen3-VL (63 GB of the base repo) stays on disk and the text block of the packed sequence
is 362 rows instead of several thousand. There is no per-clip prompt and no caption: nothing
in the output comes from text you write.

Needs `--model-dir` pointing at a local copy of MiniMaxAI/MiniMax-H3 for the VAE, the audio
VAE and the schedulers. This repository ships only the transformer and the LoRA.
"""

import argparse
import os
import time

import torch
from diffusers import MiniMaxH3Transformer3DModel, ModularPipeline
from diffusers.modular_pipelines.minimax_h3 import (MiniMaxH3AudioReference, MiniMaxH3ImageReference,
                                                    MiniMaxH3VideoReference)
from diffusers.modular_pipelines.minimax_h3.before_encoder import MiniMaxH3Ref2VASetupStep
from diffusers.modular_pipelines.minimax_h3.encoders import MiniMaxH3Ref2VATextEncoderStep
from diffusers.modular_pipelines.minimax_h3.modular_pipeline import (align_num_frames,
                                                                    audio_latent_num_frames)
from diffusers.utils.export_utils import encode_video

HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

parser = argparse.ArgumentParser()
parser.add_argument("--cond", required=True, help="the video whose motion, framing and background are kept")
parser.add_argument("--ref", required=True, help="a single still of the person to put in it")
parser.add_argument("--out", required=True)
parser.add_argument("--model-dir", required=True,
                    help="a local copy of MiniMaxAI/MiniMax-H3, for the VAE / audio VAE / schedulers")
parser.add_argument("--transformer", default=os.path.join(HERE, "transformer"))
parser.add_argument("--lora", default=os.path.join(HERE, "lora"))
parser.add_argument("--embed", default=os.path.join(HERE, "assets", "fixed_embed_fwd_anyframe.pt"))
parser.add_argument("--num-frames", type=int, default=124, help="at 24 fps; 124 frames is ~5.2 s")
parser.add_argument("--steps", type=int, default=4,
                    help="the distilled student's operating point. More is not monotonically better: "
                         "s4 is not a degraded s12")
parser.add_argument("--flow-shift", type=float, default=3.0,
                    help="the base model's released default is 12; the few-step student wants 3")
parser.add_argument("--height", type=int, default=None, help="defaults to the conditioning clip's own height")
parser.add_argument("--width", type=int, default=None, help="defaults to the conditioning clip's own width")
parser.add_argument("--short-edge", type=int, default=None,
                    help="the canvas both references are laid out on. Defaults to the conditioning clip's own "
                         "short edge, which is what this model was evaluated at")
parser.add_argument("--offload", action="store_true",
                    help="stream the transformer from CPU in groups of 5 blocks: ~12 GB resident instead of 62")
parser.add_argument("--audio", default=None,
                    help="pin the generated soundtrack to this file's audio -- usually the driving clip "
                         "itself -- so the mouth tracks real speech instead of the silence the fixed "
                         "prompt asks for. Any file PyAV can decode; a video's soundtrack is taken")
parser.add_argument("--seed", type=int, default=42)
args = parser.parse_args()

fixed = torch.load(args.embed, weights_only=False)


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


MiniMaxH3Ref2VATextEncoderStep.__call__ = use_fixed_embeds

# `--audio` holds the *target* audio rows at a real soundtrack, clean, for the whole denoise, rather
# than letting the model generate them. Two of the three things that takes have no argument on the
# pipeline, so they are patched here; the third is the `audio_latents=` passed to the call below.
if args.audio:
    from diffusers.modular_pipelines.minimax_h3.before_denoise import MiniMaxH3SetTimestepsStep
    from diffusers.modular_pipelines.minimax_h3.denoise import MiniMaxH3LoopSchedulerStep

    # (1) Those rows carry finished audio, so they have to be told they are clean. H3's flow
    # convention is reversed -- t = 1 is clean, not 0 -- and the library itself passes a literal 1.0
    # for a *reference* soundtrack. `audio_timestep` is positional argument 6.
    _build_row_timesteps = MiniMaxH3SetTimestepsStep.build_row_timesteps
    MiniMaxH3SetTimestepsStep.build_row_timesteps = staticmethod(
        lambda *a: _build_row_timesteps(*a[:6], 1.0, *a[7:]))

    # (2) ...and the scheduler must never write them, or the first step would walk them off the
    # soundtrack. Only the video rows are stepped. `num_condition_audio_rows` deliberately stays 0:
    # raising it empties the decoder's `audio_latents[num_condition_audio_rows:]` slice and trips the
    # reference-count check, and these rows are a pinned target, not a reference.
    @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

# The reference order is frozen: the presentation names `<Video 1>` then `<Picture 1>`, and that order
# advances the shared rotary clock, so it is part of the layout rather than a detail of the prompt. The
# driving clip's own soundtrack is dropped here, as it is in training: `--audio` puts one back, but as
# the *target* to be matched rather than as a reference to be imitated.
video = MiniMaxH3VideoReference.from_file(args.cond)
video.audio, video.sample_rate = None, None

# Passing an orientation that disagrees with the clip silently generates a transposed video, so the output
# geometry is derived from the clip rather than typed.
height = args.height or video.frames.shape[1]
width = args.width or video.frames.shape[2]
short_edge = args.short_edge or min(height, width)

pipe = ModularPipeline.from_pretrained(args.model_dir, workflow="ref2va")
# Both references are pinned to the target's own short edge. The base model's released defaults (768 for the
# video reference, 2048 for the image) put the references on a grid the target never shares; this model was
# finetuned and evaluated with them nested, and changing it changes the take.
pipe.register_to_config(canvas_short_edge=short_edge,
                        canvas_max_pixels=short_edge * max(height, width),
                        reference_image_short_edge=short_edge)

t0 = time.time()
# `transformer_ref` is deliberately absent: our finetune replaces it outright, so loading the base copy
# first would read 62 GB off disk only to drop it.
pipe.load_components(names=["vae", "audio_vae", "scheduler", "audio_scheduler"],
                     pretrained_model_name_or_path=args.model_dir, dtype=torch.bfloat16)
pipe.transformer_ref = MiniMaxH3Transformer3DModel.from_pretrained(args.transformer, torch_dtype=torch.bfloat16)
# `prefix=None` and the explicit `weight_name` are both required. The loader defaults to looking for a `.bin`
# (raises) and to filtering keys for a `transformer.` prefix, which these bare keys do not have -- that
# mismatch loads *nothing* and only warns, so the default would silently render the un-distilled model.
pipe.transformer_ref.load_lora_adapter(args.lora, weight_name="pytorch_lora_weights.safetensors", prefix=None)
pipe.scheduler.set_shift(args.flow_shift)

# The pipeline snaps `num_frames` up to the next `17 * n + 5` the video VAE can encode. Doing it here
# too means the pinned audio is cut to the length that is really rendered rather than the one asked
# for -- off by one grid step and the rows no longer line up with the video.
num_frames = align_num_frames(args.num_frames, pipe.vae_frames_per_chunk, pipe.vae_latents_per_chunk)

if args.offload:
    pipe.transformer_ref.enable_group_offload(
        onload_device=torch.device("cuda"), offload_type="block_level", num_blocks_per_group=5,
        non_blocking=True, use_stream=True, record_stream=True)
    pipe.vae.to("cuda")
    pipe.audio_vae.to("cuda")
else:
    pipe.to("cuda")
print(f"loaded in {time.time() - t0:.0f}s; canvas {height}x{width}, references on short edge {short_edge}")

# The soundtrack becomes target rows the same way the pipeline turns a *reference* soundtrack into
# reference rows: truncate at the source rate, resample once, take the posterior mean, normalize.
# Reusing its own helper is what keeps the two paths from drifting apart.
audio_latents = None
if args.audio:
    # One video frame more than the render needs, so the encoder cannot come up short. A source
    # shorter than the grid it renders on is padded, and that tail is real silence.
    n_samp = round((num_frames + 1) / pipe.fps * pipe.audio_sampling_rate)
    track = MiniMaxH3AudioReference.from_file(args.audio)
    wav = MiniMaxH3Ref2VASetupStep._normalize_audio_condition(
        track.audio, track.sample_rate or pipe.audio_sampling_rate, pipe.audio_sampling_rate,
        max_duration=(num_frames + 1) / pipe.fps)
    have = wav.shape[1]
    wav = torch.nn.functional.pad(wav, (0, max(0, n_samp - have)))[:, :n_samp]
    with torch.no_grad():
        # `encode` casts to the encoder's own dtype, so a float32 waveform is fine against a bf16 VAE.
        posterior = pipe.audio_vae.encode(wav[:, None].to(pipe.audio_vae.device), 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)
    n_lat = audio_latent_num_frames(num_frames, pipe.fps)
    # Channel-major rows: the two stereo channels are two batch items of the mono audio VAE.
    rows = (posterior.mode().float().cpu().transpose(1, 2)[:, :n_lat] - mean) / std
    if rows.shape[1] != n_lat:
        raise RuntimeError(f"the soundtrack encoded to {rows.shape[1]} latents, short of the {n_lat} "
                           f"that {num_frames} frames need")
    audio_latents = rows.permute(0, 2, 1).contiguous()
    print(f"pinned {min(have, n_samp) / pipe.audio_sampling_rate:.2f}s of audio -> "
          f"{tuple(audio_latents.shape)}, over {n_samp / pipe.audio_sampling_rate:.2f}s of video")

t0 = time.time()
result = pipe(
    prompt=fixed["presentation"],
    references=[video, MiniMaxH3ImageReference.from_file(args.ref)],
    num_frames=num_frames,
    height=height,
    width=width,
    num_inference_steps=args.steps,
    audio_latents=audio_latents,
    generator=torch.Generator().manual_seed(args.seed),
    output=["videos", "audio", "sampling_rate"],
)
encode_video(result["videos"][0], fps=24, output_path=args.out,
             audio=result["audio"][0], audio_sample_rate=result["sampling_rate"])
print(f"{time.time() - t0:.0f}s, peak {torch.cuda.max_memory_allocated() / 2**30:.1f} GiB -> {args.out}")