"""用与实际失败运行完全一致的配置定位 NaN：同 prompt、同修复、同显存参数。"""
import os, time, torch
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")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF","expandable_segments:True")
from diffsynth.pipelines.minimax_h3_audio_video import MiniMaxH3Pipeline, ModelConfig

PROMPT = ("A girl is very happy, she is speaking in english: "
          "“I enjoy working with Diffsynth-Studio, it's a perfect framework.”")
DT=torch.float16
free=torch.cuda.mem_get_info("cuda")[0]/1024**3
vram_limit=max(6.0, free-5.0)          # 与 --vram-reserve 5 一致
print(f"[env] free={free:.1f}G vram_limit={vram_limit:.1f}G", flush=True)
vc=dict(offload_dtype=DT,offload_device="cpu",onload_dtype=DT,onload_device="cpu",
        preparing_dtype=DT,preparing_device="cuda",computation_dtype=DT,computation_device="cuda")
NF4="DiffSynth-Studio/MiniMax-H3-NF4"
pipe=MiniMaxH3Pipeline.from_pretrained(torch_dtype=DT,device="cuda",
  model_configs=[ModelConfig(model_id=NF4,origin_file_pattern=p,**vc) for p in
    ["minimax-h3-fl2va-nf4.safetensors","minimax-h3-text-encoder-nf4.safetensors",
     "video_vae_nf4.safetensors","audio_vae_nf4.safetensors"]],
  processor_config=ModelConfig(model_id="MiniMax/MiniMax-H3",origin_file_pattern="FL2VA/processor/"),
  vram_limit=vram_limit)
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"))

FP16_MAX=65504.0
def _fix(mod, tag):
    mod.computation_dtype = torch.float32
    orig = mod.forward
    def fwd(x,*a,**kw):
        out = orig(x.float(),*a,**kw)
        print(f'[cp] fp32 out absmax={float(out.abs().max()):.1f} finite={bool(torch.isfinite(out).all())} dtype={out.dtype}', flush=True)
        return out
    mod.forward=fwd
_fix(pipe.dit.condition_proj, "condition_proj")
print("[fix ] condition_proj fp32+clamp 已装", flush=True)

# 先单独查文本编码器输出——它是整条条件链的源头
te_out={}
def te_hook(m,i,o):
    t = o[0] if isinstance(o,(tuple,list)) else o
    if torch.is_tensor(t) and t.is_floating_point():
        te_out["absmax"]=float(t.abs().max()); te_out["finite"]=bool(torch.isfinite(t).all())
        te_out["shape"]=tuple(t.shape); te_out["dtype"]=str(t.dtype)
h_te = pipe.text_encoder.register_forward_hook(te_hook)

hits=[]; seen=set()
def mk(name):
    def hook(mod, inp, out):
        if name in seen: return
        ok_in = all(bool(torch.isfinite(t).all()) for t in inp if torch.is_tensor(t) and t.is_floating_point())
        outs = out if isinstance(out,(tuple,list)) else (out,)
        for o in outs:
            if torch.is_tensor(o) and o.is_floating_point() and not torch.isfinite(o).all():
                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()]
                fin=torch.isfinite(o)
                hits.append((name, type(mod).__name__, ok_in,
                             float(o[fin].abs().max()) if fin.any() else float("nan"),
                             max(ins) if ins else None,
                             float((~fin).float().mean())*100))
                return
    return hook
handles=[]
for tag,mod in (("dit",pipe.dit),("txt",pipe.text_encoder),
                ("vvae",pipe.video_vae),("avae",pipe.audio_vae)):
    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=PROMPT, height=480,width=832,num_frames=124,
                    num_inference_steps=1, seed=0, flow_shift=6.0)
for h in handles: h.remove()
h_te.remove()

print("\n===== 结果 =====")
print("[文本编码器输出]", te_out)
import numpy as np
try:
    a=np.stack([np.asarray(f.convert("RGB"),dtype=np.float32) for f in video[::20]])
    print(f"[video] frames={len(video)} mean={a.mean():.4f} std={a.std():.4f}")
except Exception as e:
    print("[video]", type(video), e)
print(f"[audio] finite={bool(torch.isfinite(audio).all())}")
print(f"\n>>> 真源头（输入有限=True），共 {len([h for h in hits if h[2]])} 个：")
for h in hits:
    if h[2]: print(f"  ★ {h[0]}  [{h[1]}]  输入absmax={h[4]}  有限部分absmax={h[3]}  非有限占比={h[5]:.4f}%")
print(f">>> 受害者总数 {len([h for h in hits if not h[2]])}，前 5：")
for h in [x for x in hits if not x[2]][:5]: print(f"    {h[0]}  [{h[1]}]")