#!/usr/bin/env python
"""MiniMax-H3 FL2VA (NF4) Text -> Video+Audio on RTX 2080 Ti (Turing / sm_75).

Turing deltas vs the official example (all deliberate):
  * FP16 everywhere, never BF16 -- sm_75 has no bf16 SDPA kernel at all
    (EFFICIENT and FLASH both report "No available kernel"), so bf16 silently
    falls back to the MATH backend and OOMs on a ~14k-token sequence.
  * DIFFSYNTH_ATTENTION_IMPLEMENTATION=torch -- FlashAttention 2/3/4 and
    SageAttention are sm_80+. xformers works here but would be auto-selected
    ahead of torch; pin the backend so runs are reproducible.
  * vram_limit from mem_get_info()[0] (FREE) not [1] (TOTAL) -- other processes
    share these GPUs.
  * offload to "cpu", never "disk" -- /data is a 7200rpm HDD (184 MB/s).
"""
import os, sys, time, argparse

BASE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
os.environ.setdefault("DIFFSYNTH_ATTENTION_IMPLEMENTATION", "torch")
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=4)
p.add_argument("--flow-shift", type=float, default=None, help="default 6.0 with turbo, 12.0 without")
p.add_argument("--seed",   type=int, default=0)
p.add_argument("--turbo",  action="store_true", default=True)
p.add_argument("--no-turbo", dest="turbo", action="store_false")
p.add_argument("--dtype", choices=["fp32","fp16"], default="fp32",
               help="fp32=唯一能在 sm_75 上出片的选择（默认）；fp16 只用于复现全黑故障")
p.add_argument("--no-fp32-residual", dest="fp32_residual", action="store_false",
               help="关闭 DiT 残差流的 fp32 兜底（默认开启；关掉在 sm_75 上会出全黑视频）")
p.add_argument("--no-fp32-cond", dest="fp32_cond", action="store_false",
               help="关闭 condition_proj 的 fp32 兜底（默认开启；关掉在 sm_75 上会出全黑视频）")
p.add_argument("--fp32-accum", action="store_true",
               help="disable fp16 reduced-precision reduction (~2x slower, use if output is NaN/artifacted)")
p.add_argument("--vram-reserve", type=float, default=None,
               help="留给激活的显存(GB)。默认按 dtype 自适应：fp32=14（实测全分辨率峰值 12.2G），fp16=5")
p.add_argument("--out", default=None)
args = p.parse_args()

import torch
if args.fp32_accum:
    torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False

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

# dtype 决定这条线能不能出片。实测(2026-08-18, sm_75, 448x256x28f, 唯一变量是 dtype)：
#   fp32 -> 视频 std=50.6、音频 finite ✅        fp16 -> 124 帧纯黑、音频 NaN ❌
# 原因是 fp16 的动态范围装不下这个 bf16 训练的模型：dit.blocks.0 输入 absmax 89
# 已放大到 38176，50 层残差累积后必然突破 fp16 上限 65504。
# bf16 在 sm_75 上不可用（无硬件支持），所以 fp32 是这张卡上唯一能出片的选择。
# 保留 --dtype fp16 只为复现/对照，不要用它出片。
DTYPE = {"fp32": torch.float32, "fp16": torch.float16}[args.dtype]
free_gb = torch.cuda.mem_get_info("cuda")[0] / 1024**3  # [0]=free, official example wrongly uses [1]=total
if args.vram_reserve is None:
    # fp32 的激活是 fp16 的两倍，权重要多让位。实测 832x480x124f fp32：
    # vram_limit 7.3G 时峰值 12.2G，跑通；vram_limit 19.3G 时 OOM。
    args.vram_reserve = 14.0 if DTYPE is torch.float32 else 5.0
vram_limit = max(4.0, free_gb - args.vram_reserve)
if vram_limit + 1.0 > free_gb:
    print(f"[warn] 空闲显存仅 {free_gb:.1f}G，vram_limit={vram_limit:.1f}G 几乎没有余量，"
          f"很可能 OOM。先停掉占显存的服务（见 qwen38-deploy/RESTORE-GPU.md）", flush=True)

prop = torch.cuda.get_device_properties(0)
print(f"[env] GPU={prop.name} sm_{prop.major}{prop.minor}  free={free_gb:.1f}G  vram_limit={vram_limit:.1f}G", flush=True)
print(f"[env] dtype={DTYPE}  attention={os.environ['DIFFSYNTH_ATTENTION_IMPLEMENTATION']}  "
      f"fp16_reduced_accum={torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction}", flush=True)
assert prop.major == 7, "this script is tuned for Turing; re-check dtype choices on other archs"

vram_config = dict(
    offload_dtype=DTYPE,   offload_device="cpu",     # not "disk": /data is a spinning HDD
    onload_dtype=DTYPE,    onload_device="cpu",
    preparing_dtype=DTYPE, preparing_device="cuda",
    computation_dtype=DTYPE, computation_device="cuda",
)
NF4 = "DiffSynth-Studio/MiniMax-H3-NF4"

t0 = time.time()
pipe = MiniMaxH3Pipeline.from_pretrained(
    torch_dtype=DTYPE, device="cuda",
    model_configs=[
        ModelConfig(model_id=NF4, origin_file_pattern="minimax-h3-fl2va-nf4.safetensors",        **vram_config),
        ModelConfig(model_id=NF4, origin_file_pattern="minimax-h3-text-encoder-nf4.safetensors", **vram_config),
        ModelConfig(model_id=NF4, origin_file_pattern="video_vae_nf4.safetensors",               **vram_config),
        ModelConfig(model_id=NF4, origin_file_pattern="audio_vae_nf4.safetensors",               **vram_config),
    ],
    processor_config=ModelConfig(model_id="MiniMax/MiniMax-H3", 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:
    # LoRA rides on the NF4 layers: AutoWrappedQuantizedModule inherits LoRAHotLoadMixin,
    # so lora_forward() adds x@A@B to the quantized output -- no base dequantisation needed.
    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"))
# --- sm_75 必需：condition_proj 在 fp16 下累加溢出 -> inf -> 全管线 NaN -> 全黑视频 ---
# 实测(2026-08-18)：该层输入 absmax≈17808、权重 absmax≈4.66、K=5120，
# fp16 累加器(上限 65504)有 0.04% 的输出溢出成 inf，随后污染下游 999 个模块。
# 它只是 (1,5120)x(5376,5120) 的一次投影，改走 fp32 的算力开销可忽略。
# 注意：--fp32-accum 对此无效——溢出发生在累加器量级，不是 reduced-precision reduction。
def _force_fp32_linear(mod, tag):
    """该层在 fp32 里算，且**输出保持 fp32**流入 token_refiner。

    实测(2026-08-18, sm_75)：
      文本编码器输出 absmax 15328（有限、健康，源头不在它）；
      condition_proj 的 fp32 真值 absmax 可达 79554，而 fp16 上限 65504。
      产生的 inf 会污染下游 2985 个模块，最终 124 帧纯黑（且原脚本退出码为 0）。
    根因是 Turing 无 bf16：模型按 bf16 训练（bf16 有 fp32 的动态范围），
    fp16 装不下这条条件路径。

    走过的四条弯路，都记下来免得后人重走：
      1) 替换整个 forward —— 绕过 preparing/onload 换入机制。
      2) 只换 linear_forward、最后仍 `.to(fp16)` —— 无效：fp32 算得对，
         溢出发生在那次类型转换上。自检若只查转换前的张量会永远显示"正常"。
      3) 钳位到 65504 再降回 fp16 —— **反而更糟**：token_refiner.blocks.0 是
         `h = x + attn(norm1(x))`，x 恰好等于 fp16 最大可表示值时，
         再加任何正数立刻变 inf（实测 3.45% 非有限）。钳位不留余量等于自造故障。
      4) 输出保持 fp32 但 vram_limit=19.3G —— 数值对了但 CUDA OOM。
         配 `--vram-reserve 5`(vram_limit 16.3G) + PYTORCH_ALLOC_CONF=expandable_segments:True
         后峰值 17.2G，有余量。

    库里 _embed() 在 token_refiner 之后本来就有 `text_embed.to(dtype)`，
    那时已归一化、量级回落，那次转换是安全的。
    下游 RMSNorm 会因 dtype 不一致发性能警告（非错误），可忽略。
    """
    mod.computation_dtype = torch.float32
    orig = mod.forward
    def fwd(x, *a, **kw):
        out = orig(x.float(), *a, **kw)
        if not torch.isfinite(out).all():
            print(f"[warn] {tag} 即使 fp32 仍非有限", flush=True)
        return out          # 保持 fp32：既不降回也不钳位
    mod.forward = fwd
    return mod

# --- sm_75 必需 #2：DiT 残差流保持 fp32 ---
# 实测(2026-08-18)：dit.blocks.0 输入 absmax 89.25 → 输出 38176，且有极少量元素溢出。
# 50 个 block × 每个 2 次残差相加 = 100 个溢出点，累积后必然爆 fp16 的 65504。
#
# 病灶与 condition_proj 同构，见 minimax_h3_dit.py:61-67：
#   def _modulate_gate(x, gate, other, indices):
#       return (x + gate.index_select(0, indices) * other).to(x.dtype)
#                                                           ^^^^^^^^^^ 提升到 fp32 算完又降回 fp16
# 库作者自己的注释也写了 "index_select on the AdaLN params can promote the expression"，
# 但把结果降回 fp16 恰恰丢掉了提升带来的动态范围。
#
# 做法是标准混合精度：**累加器(残差流)fp32，喂给子模块的张量仍 fp16**。
#   _modulate_gate        -> 返回 fp32（它的输出就是残差流，下游是 RMSNorm，容忍混合 dtype）
#   _modulate_scale_shift -> 显式返回 fp16（它的输出喂给 attn/mlp 的 Linear4bit，必须是 fp16）
# 额外显存约 300 MB（残差流 (seq,5376) 一份 fp32），实测峰值仍在 16 GB 以内。
def _patch_dit_residual_fp32():
    from diffsynth.models import minimax_h3_dit as _m

    def _mss(x, shift, scale, indices):
        return (x.float() * (1.0 + scale.index_select(0, indices).float())
                + shift.index_select(0, indices).float()).to(torch.float16)

    def _mg(x, gate, other, indices):
        # 不降回 fp16：这一份就是残差流
        return x.float() + gate.index_select(0, indices).float() * other.float()

    _m._modulate_scale_shift = _mss
    _m._modulate_gate = _mg
    print("[fix ] DiT 残差流已切到 fp32（_modulate_gate 不再降回 fp16）", flush=True)

# 这两个补丁是 fp16 专用的挽救措施（最终证明救不回来，见 README）。
# fp32 下必须失效，否则强制 .to(float16) 会与 fp32 权重冲突：
#   RuntimeError: mat1 and mat2 must have the same dtype, but got Half and Float
if args.fp32_residual and DTYPE is torch.float16:
    _patch_dit_residual_fp32()

if args.fp32_cond and DTYPE is torch.float16:
    _force_fp32_linear(pipe.dit.condition_proj, "dit.condition_proj")
    print("[fix ] dit.condition_proj 已切到 fp32 计算（规避 sm_75 fp16 累加溢出）", flush=True)

flow_shift = args.flow_shift if args.flow_shift is not None else (6.0 if args.turbo else 12.0)

print(f"[gen ] {args.width}x{args.height} {args.frames}f steps={args.steps} "
      f"turbo={args.turbo} flow_shift={flow_shift} seed={args.seed}", flush=True)
t1 = time.time()
video, audio = pipe(
    prompt=args.prompt,
    height=args.height, width=args.width, num_frames=args.frames,
    num_inference_steps=args.steps, seed=args.seed, flow_shift=flow_shift,
)
gen_s = time.time() - t1

import torch as _t
_bad = []
if not _t.isfinite(audio if _t.is_tensor(audio) else _t.tensor(0.)).all():
    _bad.append("audio 含非有限值")
# video 是 PIL 帧列表；原脚本从不检查它，导致全黑视频也返回 0 —— 实测踩过
def _video_stats(v):
    import numpy as _np
    if _t.is_tensor(v):
        a = v.float()
        return float(a.mean()), float(a.std()), bool(_t.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(args.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)   # 不要用退出码 0 报告一个坏产物