Upload scripts/patch_unirig_restore_mha.py with huggingface_hub

scripts/patch_unirig_restore_mha.py

@@ -0,0 +1,38 @@
+# Restore unirig_skin.py to use flash_attn MHA (will be served by our shim in run.py)
+# Our earlier patch replaced flash_attn MHA with nn.MultiheadAttention which breaks
+# checkpoint loading because weight names don't match
+
+path = '/root/UniRig/src/model/unirig_skin.py'
+with open(path) as f:
+    src = f.read()
+
+# Restore flash_attn import (remove it if replaced, add it back)
+if 'from flash_attn.modules.mha import MHA' not in src:
+    # Add after the last import line before class definitions
+    src = src.replace(
+        'import torch_scatter\n',
+        'import torch_scatter\nfrom flash_attn.modules.mha import MHA\n'
+    )
+
+# Restore original MHA init (undo nn.MultiheadAttention replacement)
+src = src.replace(
+    '        self.attention = nn.MultiheadAttention(embed_dim=feat_dim, num_heads=num_heads, batch_first=True)',
+    '        self.attention = MHA(embed_dim=feat_dim, num_heads=num_heads, cross_attn=True)'
+)
+
+# Restore original MHA forward call
+src = src.replace(
+    '        attn_output, _ = self.attention(q, kv, kv)',
+    '        attn_output = self.attention(q, x_kv=kv)'
+)
+
+with open(path, 'w') as f:
+    f.write(src)
+
+print('unirig_skin.py restored to use flash_attn MHA (shim will serve it)')
+
+# Verify
+import subprocess
+r = subprocess.run(['grep', '-n', 'MHA\|flash_attn\|Wq\|in_proj', path],
+                   capture_output=True, text=True)
+print(r.stdout)