#!/usr/bin/env python
"""MiniMax-H3 FL2VA (int8-convrot) Text -> Video+Audio on RTX PRO 6000 (Blackwell / sm_120).

Blackwell has native bf16 attention kernels and native int8/fp8 tensor cores, so none of
the Turing-era workarounds in the 2080Ti script (t2va.py) apply here:
  * No fp32 residual-stream patch, no fp32 condition_proj patch -- those existed because
    sm_75 has zero bf16 attention kernel and fp16's 65504 max overflowed after 50 residual
    blocks (see 2080tiserver/minimax-h3/README.md). bf16 has fp32's dynamic range, so this
    pipeline runs bf16 end-to-end untouched, matching DiffSynth's own official example
    (examples/minimax_h3/model_inference/MiniMax-H3-Int8-ConvRot-FL2VA.py).
  * No CPU weight offload -- single 96GB GPU, ~71GB of int8-convrot weights fit resident.
    vram_config below still declares offload_device="cpu" (copied from the official example
    verbatim) but with this much headroom it should rarely trigger.
  * No sm_75 assert -- this script targets sm_80+ (native bf16 attention) only.

Precision: int8-convrot (Comfy-Org/MiniMax-H3), not NF4. Roughly 2x the bits of the NF4
build used on the 2080Ti; chosen specifically to test whether NF4's 4-bit quantization was
behind the quality complaints on that machine. See docs/todo/h3-rtxpro6000-setup.md for the
size/quality tradeoff table across NF4 / int8-convrot / fp8-pruned / bf16.

Needs `diffsynth[quant]` installed (comfy-kitchen backend for int8-convrot loading).
"""
import os, sys, time, argparse

BASE = os.path.dirname(os.path.abspath(__file__))
os.environ.setdefault("DIFFSYNTH_MODEL_BASE_PATH", os.path.join(BASE, "models"))

p = argparse.ArgumentParser()
p.add_argument("--prompt", default="A girl is very happy, she is speaking in english: "
                                   "“I enjoy working with Diffsynth-Studio, it's a perfect framework.”")
p.add_argument("--height", type=int, default=480)
p.add_argument("--width",  type=int, default=832)
p.add_argument("--frames", type=int, default=124)
p.add_argument("--steps",  type=int, default=None, help="默认: turbo=4 / 非turbo=50")
p.add_argument("--flow-shift", type=float, default=None,
               help="不传则用 pipeline 内置默认值——2080Ti 脚本硬编码过 6.0/12.0，"
                    "但官方 int8-convrot 示例根本没传这个参数，不要凭空套用未经验证的值")
p.add_argument("--seed",   type=int, default=0)
p.add_argument("--turbo",  action="store_true", default=False,
               help="4步 Turbo LoRA。未在 int8-convrot 权重上验证过 LoRA 挂载是否生效，"
                    "第一轮测画质建议先不开，跑官方默认的 50 步全量")
p.add_argument("--vram-reserve", type=float, default=2.0,
               help="留给激活的显存(GB)。单卡独占，照官方示例默认留 2G")
p.add_argument("--out", default=None)
args = p.parse_args()

import torch
from diffsynth.pipelines.minimax_h3_audio_video import MiniMaxH3Pipeline, ModelConfig
from diffsynth.utils.data.audio_video import write_video_audio

prop = torch.cuda.get_device_properties(0)
print(f"[env] GPU={prop.name} sm_{prop.major}{prop.minor}", flush=True)
assert prop.major >= 8, ("this script assumes native bf16 attention kernels (sm_80+); "
                          "on Turing (sm_75) use minimax-h3/scripts/t2va.py instead")

DTYPE = torch.bfloat16
vram_config = dict(
    offload_dtype=DTYPE, offload_device="cpu",
    onload_dtype=DTYPE,  onload_device="cpu",
    preparing_dtype=DTYPE, preparing_device="cuda",
    computation_dtype=DTYPE, computation_device="cuda",
)
COMFY = "Comfy-Org/MiniMax-H3"
MMH3  = "MiniMax/MiniMax-H3"

total_gb = torch.cuda.mem_get_info("cuda")[1] / 1024**3
vram_limit = total_gb - args.vram_reserve
print(f"[env] dtype={DTYPE}  total_vram={total_gb:.1f}G  vram_limit={vram_limit:.1f}G", flush=True)

t0 = time.time()
pipe = MiniMaxH3Pipeline.from_pretrained(
    torch_dtype=DTYPE, device="cuda",
    model_configs=[
        ModelConfig(model_id=COMFY, origin_file_pattern="diffusion_models/minimax_h3_fl2va_int8_convrot.safetensors", **vram_config),
        ModelConfig(model_id=COMFY, origin_file_pattern="text_encoders/qwen3vl_32b_minimax_h3_int8_convrot.safetensors", **vram_config),
        ModelConfig(model_id=MMH3, origin_file_pattern="FL2VA/video_vae/source/model.safetensors", **vram_config),
        ModelConfig(model_id=MMH3, origin_file_pattern="FL2VA/audio_vae/model.safetensors", **vram_config),
    ],
    processor_config=ModelConfig(model_id=MMH3, origin_file_pattern="FL2VA/processor/"),
    vram_limit=vram_limit,
)
print(f"[load] pipeline ready in {time.time()-t0:.0f}s", flush=True)

if args.turbo:
    pipe.load_lora(pipe.dit, ModelConfig(
        model_id="lightx2v/Minimax-h3-Turbo",
        origin_file_pattern="minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"))
steps = args.steps if args.steps is not None else (4 if args.turbo else 50)

gen_kwargs = dict(
    prompt=args.prompt,
    height=args.height, width=args.width, num_frames=args.frames,
    num_inference_steps=steps, seed=args.seed,
)
if args.flow_shift is not None:
    gen_kwargs["flow_shift"] = args.flow_shift

print(f"[gen ] {args.width}x{args.height} {args.frames}f steps={steps} "
      f"turbo={args.turbo} flow_shift={args.flow_shift if args.flow_shift is not None else 'pipeline默认'} "
      f"seed={args.seed}", flush=True)
t1 = time.time()
video, audio = pipe(**gen_kwargs)
gen_s = time.time() - t1

# 退出码不可信是 2080Ti 上踩过的坑：原脚本从不检查 video 是否非有限/纯色，全黑视频照样退出码 0。
# 照抄那次修的检查（见 minimax-h3/README.md「退出码不可信」一节）。
_bad = []
if not torch.isfinite(audio if torch.is_tensor(audio) else torch.tensor(0.)).all():
    _bad.append("audio 含非有限值")
def _video_stats(v):
    import numpy as _np
    if torch.is_tensor(v):
        a = v.float()
        return float(a.mean()), float(a.std()), bool(torch.isfinite(a).all())
    try:
        arrs = [_np.asarray(f.convert("RGB"), dtype=_np.float32) for f in v[:: max(1, len(v)//8)]]
        a = _np.stack(arrs)
        return float(a.mean()), float(a.std()), True
    except Exception:
        return None, None, True
_mean, _std, _fin = _video_stats(video)
if not _fin:
    _bad.append("video 含非有限值")
if _std is not None and _std < 1.0:
    _bad.append(f"video 近乎纯色 (std={_std:.3f}, mean={_mean:.3f})")
if _bad:
    print("[FAIL] " + "；".join(_bad), flush=True)

out = args.out or os.path.join(BASE, "outputs", f"t2va_{time.strftime('%Y%m%d-%H%M%S')}.mp4")
os.makedirs(os.path.dirname(out), exist_ok=True)
write_video_audio(video=video, audio=audio, output_path=out, fps=24, audio_sample_rate=32000)

print(f"[done] generation {gen_s:.0f}s ({gen_s/max(steps,1):.1f}s/step), "
      f"peak VRAM {torch.cuda.max_memory_allocated()/1024**3:.1f}G", flush=True)
print(f"[out ] {out}  ({os.path.getsize(out)/1024**2:.1f} MB)", flush=True)
if _bad:
    sys.exit(3)