h3-blackwell-runtime/tools/direct_t2v_preview.py

"""Generate a minimal direct, video-only H3 T2V preview without ComfyUI."""

import argparse
import os
from pathlib import Path
import subprocess
import time
from datetime import datetime, timezone
import warnings
import numpy as np

warnings.filterwarnings("ignore", message="Found GPU0 NVIDIA GB10 which is of cuda capability 12.1.*", category=UserWarning)

import torch

from h3_blackwell_runtime.checkpoint import H3Checkpoint
from h3_blackwell_runtime.audio_vae_decoder import MiniMaxH3AudioVAE
from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, DEFAULT_ATTENTION_BACKEND
from h3_blackwell_runtime.block import configure_mlp_chunking
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
from h3_blackwell_runtime.packing import H3PromptPacker
from h3_blackwell_runtime.qwen3vl_text import Qwen3VLPromptConditioner
from h3_blackwell_runtime.sampler import sample_video_res_multistep
from h3_blackwell_runtime.t2v import random_av_latents
from h3_blackwell_runtime.qwen3vl_vision import build_fl2va_presentation, resize_keyframe
from h3_blackwell_runtime.token_refiner import H3TokenRefiner
from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE, dtype_from_name
from h3_blackwell_runtime.vae_encoder import MiniMaxH3VideoVAEEncoder


parser = argparse.ArgumentParser()
parser.add_argument("--prompt", default="A brass-and-paper dragon flies above a rain-washed old city at blue hour.")
parser.add_argument("--output", type=Path, default=Path("/output/direct-h3-preview.mp4"))
parser.add_argument("--width", type=int, default=320)
parser.add_argument("--height", type=int, default=192)
parser.add_argument("--frames", type=int, default=22)
parser.add_argument("--steps", type=int, default=12)
parser.add_argument("--seed", type=int, default=440204)
parser.add_argument("--first-frame", type=Path, help="First keyframe image (fl2va), PNG/JPG path.")
parser.add_argument("--last-frame", type=Path, help="Last keyframe image (fl2va), PNG/JPG path.")
parser.add_argument("--attention", choices=AVAILABLE_BACKENDS, default=DEFAULT_ATTENTION_BACKEND)
parser.add_argument("--model-timesteps-capture", type=Path, help="Directory containing captured input_XX.pt H3 timesteps for strict parity checks.")
parser.add_argument("--progress", action="store_true", help="Print per-step sampler progress.")
parser.add_argument("--profile-memory", action="store_true")
parser.add_argument("--ffmpeg-loglevel", default="error", help="ffmpeg loglevel, e.g. error, warning, info.")
parser.add_argument("--quiet", action="store_true", help="Suppress JSON summary lines.")
parser.add_argument("--save-latent", type=Path)
parser.add_argument("--save-audio-latent", type=Path)
parser.add_argument("--audio-output", type=Path)
parser.add_argument("--mux-audio", action="store_true")
parser.add_argument("--skip-decode", action="store_true")
parser.add_argument("--vae-dtype", choices=("float32", "float16", "bfloat16"), default=os.getenv("H3_VAE_DTYPE", "float16"))
parser.add_argument("--vae-tile-size", type=int, default=int(os.getenv("H3_VAE_TILE_SIZE", "256")))
parser.add_argument("--mlp-chunks", type=int, default=int(os.getenv("H3_MLP_CHUNKS", "1")))
parser.add_argument("--mlp-chunk-threshold", type=int, default=int(os.getenv("H3_MLP_CHUNK_THRESHOLD", "4096")))
parser.add_argument("--cache-mode", choices=("disabled", "easycache", "h3_cache"), default="disabled")
parser.add_argument("--cache-threshold", type=float, default=0.0)
parser.add_argument("--cache-start-percent", type=float, default=0.0)
parser.add_argument("--cache-end-percent", type=float, default=1.0)
parser.add_argument("--cache-subsample-factor", type=int, default=2)
args = parser.parse_args()
started = time.perf_counter()
last_report = started


def report_memory(stage: str) -> None:
    global last_report
    if not args.profile_memory:
        return
    now = time.perf_counter()
    rss_kb = 0
    try:
        with open("/proc/self/status", encoding="utf-8") as file:
            for line in file:
                if line.startswith("VmRSS:"):
                    rss_kb = int(line.split()[1])
                    break
    except OSError:
        pass
    cuda_alloc = torch.cuda.memory_allocated() / 1024**3 if torch.cuda.is_available() else 0.0
    cuda_reserved = torch.cuda.memory_reserved() / 1024**3 if torch.cuda.is_available() else 0.0
    print({"stage": stage, "ts": datetime.now(timezone.utc).isoformat(), "epoch_s": round(time.time(), 3), "elapsed_s": round(now - started, 3), "delta_s": round(now - last_report, 3), "rss_gb": round(rss_kb / 1024**2, 3), "cuda_alloc_gb": round(cuda_alloc, 3), "cuda_reserved_gb": round(cuda_reserved, 3), "fast_safetensors": os.getenv("H3_FAST_SAFETENSORS", ""), "disable_mmap": os.getenv("H3_DISABLE_MMAP", "")}, flush=True)
    last_report = now


def report(payload: dict) -> None:
    if not args.quiet:
        print(payload, flush=True)


def ffmpeg_command(*parts: str) -> list[str]:
    return ["ffmpeg", "-hide_banner", "-loglevel", args.ffmpeg_loglevel, *parts]

report_memory("start")
checkpoint = H3Checkpoint("/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
conditioner = Qwen3VLPromptConditioner(
    "/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors",
    "/opt/h3-blackwell-runtime/src/h3_blackwell_runtime/qwen25_tokenizer",
)
report_memory("qwen_loaded")
video, audio, frames = random_av_latents(args.width, args.height, args.frames, args.seed)
model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend=args.attention).eval()
configure_mlp_chunking(model, args.mlp_chunks, args.mlp_chunk_threshold)
report_memory("h3_loaded")
model_timesteps = None
if args.model_timesteps_capture is not None:
    model_timesteps = [
        torch.load(args.model_timesteps_capture / f"input_{index:02d}.pt", map_location="cuda", weights_only=False)["timesteps"]
        for index in range(args.steps)
    ]
want_audio = args.save_audio_latent is not None or args.audio_output is not None or args.mux_audio
cache_stats = {}
refiner = H3TokenRefiner(checkpoint, attention_backend=args.attention)
packer = H3PromptPacker(checkpoint)

if args.first_frame is not None or args.last_frame is not None:
    from PIL import Image

    def load_image(path: Path) -> torch.Tensor:
        img = Image.open(path).convert("RGB")
        return torch.from_numpy(np.array(img)).permute(2, 0, 1).unsqueeze(0).float() / 255.0

    first = load_image(args.first_frame) if args.first_frame is not None else None
    last = load_image(args.last_frame) if args.last_frame is not None else None
    vae_encoder = MiniMaxH3VideoVAEEncoder.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda").eval()
    report_memory("vae_encoder_loaded")
    from h3_blackwell_runtime.qwen3vl_vision import Qwen3VL32BVision

    vision_tower = Qwen3VL32BVision(
        "/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", device="cuda", dtype=torch.bfloat16
    )
    report_memory("vision_tower_loaded")
    presentation = build_fl2va_presentation(
        args.prompt,
        first,
        last,
        width=args.width,
        height=args.height,
        frame_count=frames,
        tokenizer=conditioner.tokenizer,
        vision=vision_tower,
        text_encoder=conditioner.encoder,
        device="cuda",
    )
    cond_latents = []
    for kf in presentation.keyframes:
        resized = resize_keyframe(kf["image"].cuda(), args.width, args.height, crop="disabled" if kf["resolved_frame_index"] == 0 else "center")
        cond_latents.append(vae_encoder.encode(resized.movedim(-1, 1).cuda().float()))
    report_memory("fl2va_conditioned")
    text = refiner(presentation.text_states)
    report_memory("text_conditioned")
    seed = args.seed
    sampled = sample_video_res_multistep(
        model,
        packer,
        text,
        video,
        audio,
        steps=args.steps,
        model_timesteps=model_timesteps,
        return_audio=want_audio,
        progress=args.progress,
        seed=seed,
        text_token_tags=presentation.text_token_tags,
        cond_latents=cond_latents,
        frame_count=frames,
        cache_mode=args.cache_mode,
        cache_threshold=args.cache_threshold,
        cache_start_percent=args.cache_start_percent,
        cache_end_percent=args.cache_end_percent,
        cache_subsample_factor=args.cache_subsample_factor,
        cache_stats=cache_stats,
    )
else:
    text = refiner(conditioner(args.prompt))
    report_memory("text_conditioned")
    sampled = sample_video_res_multistep(
        model,
        packer,
        text,
        video,
        audio,
        steps=args.steps,
        model_timesteps=model_timesteps,
        return_audio=want_audio,
        progress=args.progress,
        seed=args.seed,
        cache_mode=args.cache_mode,
        cache_threshold=args.cache_threshold,
        cache_start_percent=args.cache_start_percent,
        cache_end_percent=args.cache_end_percent,
        cache_subsample_factor=args.cache_subsample_factor,
        cache_stats=cache_stats,
    )
if cache_stats:
    report({"cache": cache_stats})
if want_audio:
    latent, audio_latent = sampled
else:
    latent, audio_latent = sampled, None
report_memory("sampled")
if args.save_latent is not None:
    args.save_latent.parent.mkdir(parents=True, exist_ok=True)
    state = {"latent": latent.detach().cpu(), "frames": frames, "width": args.width, "height": args.height, "prompt": args.prompt, "seed": args.seed}
    if audio_latent is not None:
        state["audio_latent"] = audio_latent.detach().cpu()
    torch.save(state, args.save_latent)
    report({"latent": str(args.save_latent)})
if args.save_audio_latent is not None:
    if audio_latent is None:
        raise RuntimeError("audio latent was not sampled")
    args.save_audio_latent.parent.mkdir(parents=True, exist_ok=True)
    torch.save({"audio_latent": audio_latent.detach().cpu(), "frames": frames, "prompt": args.prompt, "seed": args.seed}, args.save_audio_latent)
    report({"audio_latent": str(args.save_audio_latent)})
if args.skip_decode:
    raise SystemExit(0)
vae = MiniMaxH3VideoVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda", dtype=dtype_from_name(args.vae_dtype)).eval()
vae.tile_size = args.vae_tile_size
report_memory("vae_loaded")
with torch.inference_mode():
    pixels = vae.decode(latent.to(next(vae.parameters()).dtype))[:, :, :frames]
report_memory("vae_decoded")
del vae
torch.cuda.empty_cache()
pixels = ((pixels[0].permute(1, 2, 3, 0).clamp(-1, 1) + 1) * 127.5).to(torch.uint8).cpu()
report_memory("pixels_cpu")
args.output.parent.mkdir(parents=True, exist_ok=True)
raw = args.output.with_suffix(".rgb")
video_output = args.output.with_name(args.output.stem + ".video.mp4") if args.mux_audio else args.output
pixels.numpy().tofile(raw)
subprocess.run(ffmpeg_command("-y", "-f", "rawvideo", "-pixel_format", "rgb24", "-video_size", f"{pixels.shape[2]}x{pixels.shape[1]}", "-framerate", "24", "-i", str(raw), "-an", "-c:v", "libx264", "-pix_fmt", "yuv420p", str(video_output)), check=True)
raw.unlink()
audio_path = args.audio_output
if args.mux_audio and audio_path is None:
    audio_path = args.output.with_suffix(".wav")
if audio_path is not None:
    if audio_latent is None:
        raise RuntimeError("audio latent was not sampled")
    audio_vae = MiniMaxH3AudioVAE.from_safetensors("/vae/minimax_h3_audio_vae_fp32.safetensors", device="cuda").eval()
    report_memory("audio_vae_loaded")
    with torch.inference_mode():
        waveform = audio_vae.decode(audio_latent.to(next(audio_vae.parameters()).dtype)).clamp(-1, 1).cpu()[0]
    report_memory("audio_decoded")
    audio_path.parent.mkdir(parents=True, exist_ok=True)
    audio_raw = audio_path.with_suffix(".f32le")
    waveform.transpose(0, 1).contiguous().numpy().tofile(audio_raw)
    subprocess.run(ffmpeg_command("-y", "-f", "f32le", "-ar", "32000", "-ac", "2", "-i", str(audio_raw), str(audio_path)), check=True)
    audio_raw.unlink()
    report({"audio_output": str(audio_path), "sample_rate": 32000, "audio_shape": tuple(waveform.shape)})
if args.mux_audio:
    subprocess.run(ffmpeg_command("-y", "-i", str(video_output), "-i", str(audio_path), "-c:v", "copy", "-c:a", "aac", "-shortest", str(args.output)), check=True)
    video_output.unlink()
report({"output": str(args.output), "frames": frames, "shape": tuple(pixels.shape)})