import os, torch, torch.nn.functional as F
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
DT=torch.float16
free=torch.cuda.mem_get_info("cuda")[0]/1024**3
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=max(6.0,free-2.0))
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"))

cp = pipe.dit.condition_proj
stats = {"calls":0, "inf_out":0}
orig_lf = cp.linear_forward           # 改 linear_forward 而不是 forward：
                                      # 这样 onload/preparing/LoRA 全套机制保持不变
def lf(x, weight, bias):
    stats["calls"] += 1
    out = F.linear(x.float(), weight.float(), bias.float() if bias is not None else None)
    if not torch.isfinite(out).all(): stats["inf_out"] += 1
    return out.to(x.dtype)
cp.linear_forward = lf
print("[patch] linear_forward 已替换", flush=True)

video, audio = pipe(prompt="一只橘猫在窗台上打哈欠", height=480, width=832,
                    num_frames=124, num_inference_steps=1, seed=42, flow_shift=6.0)
print(f"[stat] condition_proj 被调用 {stats['calls']} 次，其中 fp32 仍非有限 {stats['inf_out']} 次", flush=True)
import numpy as np
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} max={a.max():.1f}", flush=True)
print(f"[audio] finite={bool(torch.isfinite(audio).all())} absmax={float(audio.abs().max()):.4g}", flush=True)