Upload export_decoder.py with huggingface_hub

export_decoder.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+"""
+OmniVoice Audio Decoder Export — Model Surgery + torch.export
+
+Extracts the HiggsAudioV2 decoder as a standalone nn.Module:
+  Input:  audio_codes [B, C, T] (C=8 codebooks, T=time steps)
+  Output: waveform [B, 1, samples] (24kHz)
+
+Each time step = 40ms = 960 samples at 24kHz
+"""
+
+import torch
+import torch.nn as nn
+import time
+
+
+class OmniVoiceDecoder(nn.Module):
+    """Standalone audio decoder wrapper for export."""
+
+    def __init__(self, audio_tokenizer):
+        super().__init__()
+        # Extract the actual decoder submodules
+        self.tokenizer = audio_tokenizer
+
+    def forward(self, audio_codes: torch.LongTensor) -> torch.Tensor:
+        """
+        Args:
+            audio_codes: [B, C, T] where C=8, T=time steps
+        Returns:
+            waveform: [B, 1, samples]
+        """
+        out = self.tokenizer.decode(audio_codes)
+        return out.audio_values
+
+
+def main():
+    print("Loading OmniVoice...", flush=True)
+    t0 = time.time()
+
+    from omnivoice import OmniVoice
+    model = OmniVoice.from_pretrained(
+        "k2-fsa/OmniVoice",
+        device_map="cpu",
+        dtype=torch.float32,
+    )
+    print(f"  Loaded in {time.time()-t0:.1f}s", flush=True)
+
+    # Extract decoder
+    print("Extracting decoder...", flush=True)
+    decoder = OmniVoiceDecoder(model.audio_tokenizer)
+    decoder.eval()
+
+    params = sum(p.numel() for p in decoder.parameters())
+    print(f"  Decoder params: {params/1e6:.1f}M ({params*4/1e9:.2f}GB fp32)", flush=True)
+
+    # Test with different sequence lengths
+    for T in [50, 100, 200]:
+        audio_codes = torch.randint(0, 1024, (1, 8, T), dtype=torch.long)
+        t0 = time.time()
+        with torch.no_grad():
+            waveform = decoder(audio_codes)
+        elapsed = time.time() - t0
+        print(f"  T={T}: output={waveform.shape}, duration={waveform.shape[-1]/24000:.2f}s, time={elapsed:.1f}s", flush=True)
+
+    # Try torch.export
+    print("\nAttempting torch.export...", flush=True)
+    test_codes = torch.randint(0, 1024, (1, 8, 100), dtype=torch.long)
+
+    try:
+        from torch.export import export
+        t0 = time.time()
+        ep = export(decoder, (test_codes,))
+        print(f"  torch.export succeeded in {time.time()-t0:.1f}s!", flush=True)
+        print(f"  Graph: {len(ep.graph.nodes)} nodes", flush=True)
+
+        torch.export.save(ep, "omnivoice_decoder_exported.pt2")
+        print("  Saved to omnivoice_decoder_exported.pt2", flush=True)
+
+    except Exception as e:
+        print(f"  torch.export failed: {type(e).__name__}: {e}", flush=True)
+        print("  Will need to dig into the decoder internals for fixes", flush=True)
+
+    # Save state dict
+    torch.save(decoder.state_dict(), "omnivoice_decoder_state.pt")
+    print(f"\nSaved decoder state dict to omnivoice_decoder_state.pt", flush=True)
+
+
+if __name__ == "__main__":
+    main()