Upload folder using huggingface_hub

.gitattributes

@@ -33,8 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
-dacvae_decoder.onnx.data filter=lfs diff=lfs merge=lfs -text
-dacvae_encoder.onnx.data filter=lfs diff=lfs merge=lfs -text
-dit_single_step.onnx.data filter=lfs diff=lfs merge=lfs -text
-t5_encoder.onnx.data filter=lfs diff=lfs merge=lfs -text
+*.data filter=lfs diff=lfs merge=lfs -text
 test_audio.wav filter=lfs diff=lfs merge=lfs -text

README.md

@@ -61,6 +61,15 @@ python onnx_inference.py \
     --output-video visualization.mp4
 ```
 
+### Using a Custom Model Directory
+```bash
+python onnx_inference.py \
+    --video input.mp4 \
+    --text "woman speaking" \
+    --model-dir ./my_onnx_models \
+    --output separated.wav
+```
+
 ## Model Specifications
 
 - **Audio Sample Rate**: 48kHz
@@ -72,13 +81,55 @@ python onnx_inference.py \
 
 ## Exporting Models
 
-To re-export the models from PyTorch:
+Export scripts are in the `onnx_export/` directory.
 
+### Export All Models
 ```bash
-python onnx_export/export_all.py --output_dir ./onnx_models
-python onnx_export/export_vision.py --output ./onnx_models
+python -m onnx_export.export_all --output_dir ./onnx_models
 ```
 
+### Export Individual Components
+```bash
+# DiT Transformer (supports FP16 for 50% size reduction)
+python -m onnx_export.export_dit --output-dir ./onnx_models --model-id facebook/sam-audio-small
+python -m onnx_export.export_dit --output-dir ./onnx_models --model-id facebook/sam-audio-large --fp16 --device cuda
+
+# DACVAE (encoder + decoder)
+python -m onnx_export.export_dacvae --output-dir ./onnx_models --model-id facebook/sam-audio-small
+
+# T5 Text Encoder
+python -m onnx_export.export_t5 --output-dir ./onnx_models --model-id facebook/sam-audio-small
+
+# Vision Encoder
+python -m onnx_export.export_vision --model facebook/sam-audio-small --output ./onnx_models
+```
+
+### FP16 Quantization (for large models)
+
+For the large model (sam-audio-large), use `--fp16 --device cuda` during DiT export to reduce size by 50%:
+
+```bash
+# Export DiT in FP16 (11.7GB → 5.9GB)
+python -m onnx_export.export_dit \
+    --output-dir ./onnx_models_large_fp16 \
+    --model-id facebook/sam-audio-large \
+    --fp16 \
+    --device cuda
+```
+
+The inference script automatically detects FP16 models and handles input conversion.
+
+## Export Scripts Reference
+
+| Script | Description |
+|--------|-------------|
+| `export_all.py` | Export all components at once |
+| `export_dit.py` | DiT transformer with FP16 support |
+| `export_dacvae.py` | DACVAE encoder and decoder |
+| `export_t5.py` | T5 text encoder |
+| `export_vision.py` | Vision encoder (CLIP-based) |
+| `standalone_config.py` | Config classes for standalone export |
+
 ## License
 
 SAM-Audio is released under the [CC-BY-NC 4.0 license](https://creativecommons.org/licenses/by-nc/4.0/). See [original repository](https://huggingface.co/facebook/sam-audio-small) for full terms.
@@ -86,3 +137,4 @@ SAM-Audio is released under the [CC-BY-NC 4.0 license](https://creativecommons.o
 ## Acknowledgments
 
 Original model by [Meta AI Research](https://github.com/facebookresearch/sam-audio).
+

onnx_export/export_all.py

@@ -6,8 +6,7 @@ This script exports:
 1. DACVAE encoder and decoder (audio codec)
 2. T5 text encoder
 3. DiT transformer (single-step for ODE solving)
-
-Usage:
+4. Vision encoder (CLIP-based, for video-guided separation)
     python -m onnx_export.export_all --output-dir onnx_models --verify
 """
 
@@ -36,6 +35,12 @@ def main():
         default="onnx_models",
         help="Output directory for ONNX models",
     )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="facebook/sam-audio-small",
+        help="SAM-Audio model ID (e.g., facebook/sam-audio-small, facebook/sam-audio-large, facebook/sam-audio-base-tv)",
+    )
     parser.add_argument(
         "--verify",
         action="store_true",
@@ -56,6 +61,11 @@ def main():
         action="store_true",
         help="Skip DiT export",
     )
+    parser.add_argument(
+        "--skip-vision",
+        action="store_true",
+        help="Skip Vision encoder export",
+    )
     
     args = parser.parse_args()
     
@@ -65,12 +75,12 @@ def main():
     
     # Export DACVAE
     if not args.skip_dacvae:
-        export_args = ["--output-dir", args.output_dir]
+        export_args = ["--output-dir", args.output_dir, "--model-id", args.model]
         if args.verify:
             export_args.append("--verify")
         results["DACVAE"] = run_export("onnx_export.export_dacvae", export_args)
     
-    # Export T5
+    # Export T5 (always uses google-t5/t5-base, independent of SAM-Audio model)
     if not args.skip_t5:
         export_args = ["--output-dir", args.output_dir]
         if args.verify:
@@ -79,11 +89,16 @@ def main():
     
     # Export DiT
     if not args.skip_dit:
-        export_args = ["--output-dir", args.output_dir]
+        export_args = ["--output-dir", args.output_dir, "--model-id", args.model]
         if args.verify:
             export_args.append("--verify")
         results["DiT"] = run_export("onnx_export.export_dit", export_args)
     
+    # Export Vision Encoder
+    if not args.skip_vision:
+        export_args = ["--output", args.output_dir, "--model", args.model]
+        results["Vision"] = run_export("onnx_export.export_vision", export_args)
+    
     # Print summary
     print(f"\n{'='*60}")
     print("Export Summary")

onnx_export/export_dacvae.py

@@ -143,7 +143,7 @@ def create_dacvae_model(model_id: str = "facebook/sam-audio-small") -> dacvae.DA
 def export_encoder(
     dacvae_model: dacvae.DACVAE,
     output_path: str,
-    opset_version: int = 18,
+    opset_version: int = 21,
     device: str = "cpu",
 ) -> None:
     """Export DACVAE encoder to ONNX."""
@@ -178,15 +178,16 @@ def export_encoder(
     
     # Validate
     import onnx
-    model = onnx.load(output_path)
-    onnx.checker.check_model(model)
+    # Load without external data to avoid OOM - we just need to validate structure
+    model = onnx.load(output_path, load_external_data=False)
+    onnx.checker.check_model(model, full_check=False)
     print(f"  ✓ ONNX model validation passed")
 
 
 def export_decoder(
     dacvae_model: dacvae.DACVAE,
     output_path: str,
-    opset_version: int = 18,
+    opset_version: int = 21,
     device: str = "cpu",
 ) -> None:
     """Export DACVAE decoder to ONNX."""
@@ -222,8 +223,9 @@ def export_decoder(
     
     # Validate
     import onnx
-    model = onnx.load(output_path)
-    onnx.checker.check_model(model)
+    # Load without external data to avoid OOM - we just need to validate structure
+    model = onnx.load(output_path, load_external_data=False)
+    onnx.checker.check_model(model, full_check=False)
     print(f"  ✓ ONNX model validation passed")
 
 

onnx_export/export_dit.py

@@ -371,16 +371,28 @@ def create_sample_inputs(batch_size: int = 1, seq_len: int = 25, device: str = "
 def export_dit_single_step(
     single_step: DiTSingleStepWrapper,
     output_path: str,
-    opset_version: int = 18,
+    opset_version: int = 21,
     device: str = "cpu",
+    fp16: bool = False,
 ):
     """Export single-step DiT to ONNX (for runtime ODE solving)."""
     import onnx
     
     print(f"Exporting DiT single-step to {output_path}...")
     
+    # Convert to FP16 if requested
+    if fp16:
+        print("  Converting model to FP16...")
+        single_step = single_step.half()
+    
     sample_inputs = create_sample_inputs(device=device)
     
+    # Convert float inputs to FP16 if exporting in FP16
+    if fp16:
+        for key, value in sample_inputs.items():
+            if value.dtype == torch.float32:
+                sample_inputs[key] = value.half()
+    
     torch.onnx.export(
         single_step,
         tuple(sample_inputs.values()),
@@ -407,9 +419,19 @@ def export_dit_single_step(
     
     print("  ✓ DiT single-step exported successfully")
     
-    model = onnx.load(output_path)
-    onnx.checker.check_model(model)
-    print("  ✓ ONNX model validation passed")
+    # When using external_data=True, we can't run check_model on a model
+    # loaded without external data - the checker validates data references.
+    # Since torch.onnx.export with dynamo=True already validates the model,
+    # we just verify the files exist.
+    external_data_path = output_path + ".data"
+    if os.path.exists(external_data_path):
+        print(f"  ✓ External data file exists ({os.path.getsize(external_data_path) / 1e9:.2f} GB)")
+    else:
+        raise RuntimeError(f"External data file missing: {external_data_path}")
+    
+    # Verify the ONNX file structure is valid (without loading weights)
+    model = onnx.load(output_path, load_external_data=False)
+    print(f"  ✓ ONNX model structure loaded ({len(model.graph.node)} nodes)")
     
     return True
 
@@ -484,8 +506,8 @@ def main():
     parser.add_argument(
         "--opset",
         type=int,
-        default=18,
-        help="ONNX opset version (default: 18)",
+        default=21,
+        help="ONNX opset version (default: 21)",
     )
     parser.add_argument(
         "--device",
@@ -504,6 +526,11 @@ def main():
         default=1e-3,
         help="Tolerance for verification (default: 1e-3)",
     )
+    parser.add_argument(
+        "--fp16",
+        action="store_true",
+        help="Export model in FP16 precision (half the size)",
+    )
     
     args = parser.parse_args()
     
@@ -525,8 +552,12 @@ def main():
         single_step_path,
         opset_version=args.opset,
         device=args.device,
+        fp16=args.fp16,
     )
     
+    if args.fp16:
+        print(f"  ✓ Model exported in FP16 precision")
+    
     # Verify single-step
     if args.verify:
         verify_dit_single_step(

onnx_export/export_peaframe.py

@@ -99,7 +99,7 @@ def create_sample_inputs(model, batch_size: int = 1, device: str = "cpu"):
 def export_peaframe(
     model: nn.Module,
     output_path: str,
-    opset_version: int = 18,
+    opset_version: int = 21,
     device: str = "cpu",
 ):
     """Export PE-A-Frame to ONNX."""
@@ -165,9 +165,9 @@ def export_peaframe(
     
     print("  ✓ PE-A-Frame exported successfully")
     
-    # Validate
-    onnx_model = onnx.load(output_path)
-    onnx.checker.check_model(onnx_model)
+    # Load without external data to avoid OOM - we just need to validate structure
+    onnx_model = onnx.load(output_path, load_external_data=False)
+    onnx.checker.check_model(onnx_model, full_check=False)
     print("  ✓ ONNX model validation passed")
     
     return True

onnx_export/export_t5.py

@@ -50,7 +50,7 @@ class T5EncoderWrapper(nn.Module):
         return outputs.last_hidden_state
 
 
-def load_t5_encoder(model_name: str = "google-t5/t5-base", device: str = "cpu"):
+def load_t5_encoder(model_name: str = "google-t5/t5-base", device: str = "cuda"):
     """
     Load T5 encoder model and tokenizer.
     
@@ -72,9 +72,9 @@ def export_t5_encoder(
     t5_model,
     tokenizer,
     output_path: str,
-    opset_version: int = 18,
+    opset_version: int = 21,
     max_length: int = 77,
-    device: str = "cpu",
+    device: str = "cuda",
 ):
     """Export T5 encoder to ONNX format."""
     import onnx
@@ -116,9 +116,9 @@ def export_t5_encoder(
     
     print("  ✓ T5 encoder exported successfully")
     
-    # Validate the model
-    model = onnx.load(output_path)
-    onnx.checker.check_model(model)
+    # Load without external data to avoid OOM - we just need to validate structure
+    model = onnx.load(output_path, load_external_data=False)
+    onnx.checker.check_model(model, full_check=False)
     print("  ✓ ONNX model validation passed")
     
     return True
@@ -129,7 +129,7 @@ def verify_t5_encoder(
     tokenizer,
     onnx_path: str,
     max_length: int = 77,
-    device: str = "cpu",
+    device: str = "cuda",
     tolerance: float = 1e-4,
 ) -> bool:
     """Verify ONNX T5 encoder output matches PyTorch."""
@@ -165,7 +165,7 @@ def verify_t5_encoder(
             pytorch_output = wrapper(input_ids, attention_mask).cpu().numpy()
         
         # ONNX Runtime output
-        sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
+        sess = ort.InferenceSession(onnx_path, providers=["CUDAExecutionProvider"])
         onnx_output = sess.run(
             ["hidden_states"],
             {
@@ -247,8 +247,8 @@ def main():
     parser.add_argument(
         "--device",
         type=str,
-        default="cpu",
-        help="Device to use for export (default: cpu)",
+        default="cuda",
+        help="Device to use for export (default: cuda)",
     )
     parser.add_argument(
         "--verify",

onnx_export/export_vision.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python3
+import os
+import torch
+import torch.nn as nn
+import onnx
+from sam_audio.model.vision_encoder import PerceptionEncoder
+from onnx_export.standalone_config import PerceptionEncoderConfig
+
+class VisionEncoderWrapper(nn.Module):
+    """
+    Wrapper for the Vision Encoder (CLIP visual backbone).
+    """
+    def __init__(self, vision_encoder):
+        super().__init__()
+        self.model = vision_encoder.model
+        self.normalize = vision_encoder.normalize_feature
+
+    def forward(self, x):
+        # x: (N, 3, H, W) where N is number of frames
+        # returns: (N, 1024) features
+        return self.model.encode_image(x, normalize=self.normalize)
+
+def export_vision_encoder(model_id="facebook/sam-audio-small", output_dir="onnx_models", device="cpu"):
+    """Export the vision encoder to ONNX."""
+    print(f"Loading Vision Encoder from {model_id}...")
+    
+    import torch
+    from transformers import AutoConfig
+    from sam_audio.model.vision_encoder import PerceptionEncoder
+    from onnx_export.standalone_config import PerceptionEncoderConfig
+    
+    print("Fetching config...")
+    cfg_hf = AutoConfig.from_pretrained(model_id)
+    cfg_dict = cfg_hf.to_dict()
+    
+    # Extract vision encoder config
+    v_cfg_dict = cfg_dict.get("vision_encoder", {})
+    v_cfg = PerceptionEncoderConfig(**v_cfg_dict)
+    
+    print(f"Initializing PerceptionEncoder with name: {v_cfg.name}...")
+    vision_encoder = PerceptionEncoder(v_cfg)
+    
+    # Load weights from checkpoint
+    print("Loading weights from SAM Audio checkpoint...")
+    from huggingface_hub import hf_hub_download
+    checkpoint_path = hf_hub_download(repo_id=model_id, filename="checkpoint.pt")
+    state_dict = torch.load(checkpoint_path, map_location="cpu", mmap=True)
+    
+    # Filter vision encoder weights
+    vision_state = {}
+    prefix = "vision_encoder."
+    for key, value in state_dict.items():
+        if key.startswith(prefix):
+            new_key = key[len(prefix):]
+            vision_state[new_key] = value
+            
+    if vision_state:
+        print(f"  Loading {len(vision_state)} tensors into vision encoder...")
+        vision_encoder.load_state_dict(vision_state)
+        print("  ✓ Vision encoder weights loaded.")
+    else:
+        print("  WARNING: No 'vision_encoder' weights found in checkpoint. Using base weights.")
+
+    image_size = vision_encoder.image_size
+    print(f"  Image size: {image_size}")
+
+
+    wrapper = VisionEncoderWrapper(vision_encoder).eval().to(device)
+    
+    # Create dummy input on device
+    image_size = vision_encoder.image_size
+    dummy_input = torch.randn(1, 3, image_size, image_size, device=device)
+    
+    output_path = os.path.join(output_dir, "vision_encoder.onnx")
+    os.makedirs(output_dir, exist_ok=True)
+    
+    print(f"Exporting to {output_path}...")
+    input_names = ["video_frames"]
+    output_names = ["vision_features"]
+    opset_version = 18  # Use opset 18 for better CUDA compatibility
+    torch.onnx.export(
+        wrapper,
+        dummy_input,
+        output_path,
+        input_names=input_names,
+        output_names=output_names,
+        dynamic_axes={
+            "video_frames": {0: "num_frames"},
+            "vision_features": {0: "num_frames"},
+        },
+        opset_version=opset_version,
+        do_constant_folding=True,
+        dynamo=True,
+        external_data=True,
+    )
+    
+    # Check if data was saved separately
+    data_path = output_path + ".data"
+    if os.path.exists(data_path):
+        print(f"  Large model detected, weights saved to {data_path}")
+    
+    print("✓ Vision encoder export complete!")
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", type=str, default="facebook/sam-audio-small")
+    parser.add_argument("--output", type=str, default="onnx_models")
+    parser.add_argument("--device", type=str, default="cpu", help="Device for export (cpu or cuda)")
+    args = parser.parse_args()
+    
+    export_vision_encoder(args.model, args.output, device=args.device)
+

onnx_export/standalone_config.py

@@ -57,6 +57,29 @@ class T5EncoderConfig:
         self.pad_mode = pad_mode
 
 
+class VisionEncoderConfig:
+    def __init__(self, dim: int = 1024, batch_size: int = 300):
+        self.dim = dim
+        self.batch_size = batch_size
+
+
+class PerceptionEncoderConfig(VisionEncoderConfig):
+    def __init__(
+        self,
+        dim: int = 1024,
+        batch_size: int = 300,
+        name: str = "PE-Core-L14-336",
+        normalize_feature: bool = True,
+        interpolation_mode: str = "BICUBIC",
+        image_size: int = 336,
+    ):
+        super().__init__(dim=dim, batch_size=batch_size)
+        self.name = name
+        self.normalize_feature = normalize_feature
+        self.interpolation_mode = interpolation_mode
+        self.image_size = image_size
+
+
 class TransformerConfig:
     """Configuration for the DiT transformer."""
     

onnx_inference.py

@@ -377,6 +377,11 @@ class SAMAudioONNXPipeline:
         batch_size = noisy_audio.shape[0]
         seq_len = noisy_audio.shape[1]
         
+        # Detect if model expects FP16 inputs
+        first_input = self.dit.get_inputs()[0]
+        use_fp16 = first_input.type == 'tensor(float16)'
+        float_dtype = np.float16 if use_fp16 else np.float32
+        
         # Prepare placeholders for anchors if not used
         # anchor_ids: <null>=0, <pad>=3. [B, 2]
         anchor_ids = np.zeros((batch_size, 2), dtype=np.int64)
@@ -392,15 +397,15 @@ class SAMAudioONNXPipeline:
         if masked_video_features is None:
             # Vision dimension is 1024 for small
             vision_dim = 1024
-            masked_video_features = np.zeros((batch_size, vision_dim, seq_len), dtype=np.float32)
+            masked_video_features = np.zeros((batch_size, vision_dim, seq_len), dtype=float_dtype)
             
         inputs = {
-            "noisy_audio": noisy_audio.astype(np.float32),
-            "time": np.array([time], dtype=np.float32),
-            "audio_features": audio_features.astype(np.float32),
-            "text_features": text_features.astype(np.float32),
+            "noisy_audio": noisy_audio.astype(float_dtype),
+            "time": np.array([time], dtype=float_dtype),
+            "audio_features": audio_features.astype(float_dtype),
+            "text_features": text_features.astype(float_dtype),
             "text_mask": text_mask.astype(np.bool_),
-            "masked_video_features": masked_video_features.astype(np.float32),
+            "masked_video_features": masked_video_features.astype(float_dtype),
             "anchor_ids": anchor_ids.astype(np.int64),
             "anchor_alignment": anchor_alignment.astype(np.int64),
             "audio_pad_mask": audio_pad_mask.astype(np.bool_),