"""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`.

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 MiniMaxH3ImageReference, MiniMaxH3VideoReference
from diffusers.modular_pipelines.minimax_h3.encoders import MiniMaxH3Ref2VATextEncoderStep
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("--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

# 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, as it is in training.
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)

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}")

t0 = time.time()
result = pipe(
    prompt=fixed["presentation"],
    references=[video, MiniMaxH3ImageReference.from_file(args.ref)],
    num_frames=args.num_frames,
    height=height,
    width=width,
    num_inference_steps=args.steps,
    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}")