"""Capture ComfyUI's effective dynamically cast AdaLN linear parameters."""

from pathlib import Path


model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
source = model.read_text(encoding="utf-8")
old = 'if capture_dir: torch.save({"norm": norm1.detach().cpu(), "shift": shift_msa.detach().cpu(), "scale": scale_msa.detach().cpu()}, os.path.join(capture_dir, "block0_norm1_adaln.pt"))'
new = 'if capture_dir:\n            effective_weight, effective_bias, effective_stream = comfy.ops.cast_bias_weight(self.adaln_proj.linear, t_emb, offloadable=True)\n            torch.save({"norm": norm1.detach().cpu(), "shift": shift_msa.detach().cpu(), "scale": scale_msa.detach().cpu(), "effective_weight": effective_weight.detach().cpu(), "effective_bias": effective_bias.detach().cpu()}, os.path.join(capture_dir, "block0_norm1_adaln.pt"))\n            comfy.ops.uncast_bias_weight(self.adaln_proj.linear, effective_weight, effective_bias, effective_stream)'
if source.count(old) == 1:
    source = source.replace(old, new)
elif new not in source:
    raise RuntimeError("Unable to locate AdaLN capture payload.")
model.write_text(source, encoding="utf-8")
print("Applied H3 AdaLN dispatch capture patch.")