"""定位 NaN 的产生位置：对 dit / text_encoder 挂 forward hook，
找出第一个『输入有限、输出非有限』的模块。"""
import os, sys, time, torch
sys.argv = [sys.argv[0]]
BASE = "/data/models/minimax-h3"
os.environ.setdefault("DIFFSYNTH_ATTENTION_IMPLEMENTATION", "torch")
os.environ.setdefault("DIFFSYNTH_MODEL_BASE_PATH", f"{BASE}/models")
os.environ.setdefault("MODELSCOPE_CACHE", f"{BASE}/.mscache")

from diffsynth.pipelines.minimax_h3_audio_video import MiniMaxH3Pipeline, ModelConfig

DTYPE = torch.float16
free_gb = torch.cuda.mem_get_info("cuda")[0] / 1024**3
vram_limit = max(6.0, free_gb - 2.0)
print(f"[env] free={free_gb:.1f}G vram_limit={vram_limit:.1f}G", flush=True)

vc = 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")
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", **vc),
        ModelConfig(model_id=NF4, origin_file_pattern="minimax-h3-text-encoder-nf4.safetensors", **vc),
        ModelConfig(model_id=NF4, origin_file_pattern="video_vae_nf4.safetensors", **vc),
        ModelConfig(model_id=NF4, origin_file_pattern="audio_vae_nf4.safetensors", **vc),
    ],
    processor_config=ModelConfig(model_id="MiniMax/MiniMax-H3", origin_file_pattern="FL2VA/processor/"),
    vram_limit=vram_limit)
print(f"[load] {time.time()-t0:.0f}s", flush=True)

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

# 与 t2va.py 一致的修复：整条路径（含 LoRA）抬到 fp32
def _fp32_all(mod):
    mod.computation_dtype = torch.float32
    orig = mod.forward
    def fwd(x, *a, **kw):
        return orig(x.float(), *a, **kw).to(x.dtype)
    mod.forward = fwd
_fp32_all(pipe.dit.condition_proj)
print("[fix ] condition_proj 全路径 -> fp32", flush=True)

def finite(x):
    if torch.is_tensor(x):
        return bool(torch.isfinite(x).all()), (float(x.abs().max()) if x.numel() and x.is_floating_point() else None)
    return True, None

hits = []      # 第一个坏掉的模块
seen = set()
def mk(name):
    def hook(mod, inp, out):
        if name in seen: return
        ok_in = all(finite(t)[0] for t in inp if torch.is_tensor(t))
        outs = out if isinstance(out, (tuple, list)) else (out,)
        for o in outs:
            ok, mx = finite(o)
            if not ok:
                seen.add(name)
                ins=[float(t.abs().max()) for t in inp if torch.is_tensor(t) and t.is_floating_point() and t.numel()]
                hits.append((name, type(mod).__name__, ok_in, mx, max(ins) if ins else None))
                return
    return hook

handles = []
for tag, mod in (("dit", pipe.dit), ("txt", pipe.text_encoder)):
    for n, m in mod.named_modules():
        if n:  # 跳过根
            handles.append(m.register_forward_hook(mk(f"{tag}.{n}")))
print(f"[hook] {len(handles)} 个模块已挂钩", flush=True)

video, audio = pipe(prompt="一只橘猫在窗台上打哈欠，阳光洒进来",
                    height=480, width=832, num_frames=124,
                    num_inference_steps=1, seed=42, flow_shift=6.0)
for h in handles: h.remove()

print("\n===== 结果 =====", flush=True)
for t, nm in (("video", video), ("audio", audio)):
    if torch.is_tensor(nm):
        print(f"{t}: shape={tuple(nm.shape)} dtype={nm.dtype} finite={bool(torch.isfinite(nm).all())} "
              f"absmax={float(nm.abs().max()) if nm.is_floating_point() else 'n/a'}")
    else:
        print(f"{t}: type={type(nm)}")
print(f"\n第一批非有限输出的模块（共 {len(hits)}，只列前 15）：")
print(">>> 真源头（输入有限=True）：")
for h in hits:
    if h[2]: print(f"  ★ {h[0]}  [{h[1]}]  输入absmax={h[4]}  输出absmax={h[3]}")
print(">>> 下游受害者（前 5）：")
for h in [x for x in hits if not x[2]][:5]:
    print(f"    {h[0]}  [{h[1]}]")
if not hits:
    print("  无 —— dit/text_encoder 内部全程有限，问题在 VAE 解码或后处理")