# Patch for ComfyUI-WanVideoWrapper multitalk/wav2vec2.py with transformers >= 5:
# Wav2Vec2Encoder.forward no longer honours output_hidden_states, so hidden_states comes back None.
# We capture the encoder input and every layer output with hooks, matching transformers 4 ordering.
import sys
path = "/ComfyUI/custom_nodes/ComfyUI-WanVideoWrapper/multitalk/wav2vec2.py"
src = open(path).read()
MARK = "# PATCH transformers>=5 hidden_states"
if MARK in src:
    print("PATCH_ALREADY_APPLIED"); sys.exit(0)
old = """        encoder_outputs = self.encoder(
            hidden_states,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
"""
new = """        """ + MARK + """
        _captured = []
        _handles = []
        if output_hidden_states:
            def _pre(mod, args, kwargs):
                _captured.append(args[0] if args else kwargs.get("hidden_states"))
            _handles.append(self.encoder.layers[0].register_forward_pre_hook(_pre, with_kwargs=True))
            for _layer in self.encoder.layers:
                _handles.append(_layer.register_forward_hook(lambda mod, args, out: _captured.append(out[0] if isinstance(out, tuple) else out)))
        try:
            encoder_outputs = self.encoder(hidden_states, attention_mask=attention_mask)
        finally:
            for _h in _handles:
                _h.remove()
        if output_hidden_states:
            if len(_captured) != len(self.encoder.layers) + 1:
                raise RuntimeError(f"wav2vec2 patch: captured {len(_captured)} states, expected {len(self.encoder.layers) + 1}")
            encoder_outputs = BaseModelOutput(last_hidden_state=encoder_outputs[0], hidden_states=tuple(_captured), attentions=None)
"""
count = src.count(old)
if count < 1:
    print("PATCH_FAILED old block not found"); sys.exit(1)
open(path, "w").write(src.replace(old, new))
print(f"PATCH_APPLIED to {count} blocks")