156 lines
8.2 KiB
Python
156 lines
8.2 KiB
Python
"""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
|
|
|
|
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.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.token_refiner import H3TokenRefiner
|
|
from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE, dtype_from_name
|
|
|
|
|
|
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("--attention", choices=("sage2", "sdpa", "sage3"), default="sage2")
|
|
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")))
|
|
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()
|
|
report_memory("h3_loaded")
|
|
text = H3TokenRefiner(checkpoint, attention_backend=args.attention)(conditioner(args.prompt))
|
|
report_memory("text_conditioned")
|
|
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
|
|
sampled = sample_video_res_multistep(model, H3PromptPacker(checkpoint), text, video, audio, steps=args.steps, model_timesteps=model_timesteps, return_audio=want_audio, progress=args.progress)
|
|
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)})
|