Update export_all.py to include PEAFrame and CLAP exports

onnx_export/export_all.py

@@ -7,6 +7,10 @@ This script exports:
 2. T5 text encoder
 3. DiT transformer (single-step for ODE solving)
 4. Vision encoder (CLIP-based, for video-guided separation)
+5. PEAFrame span predictor (audio-text similarity for span detection)
+6. CLAP reranker (audio-text similarity for candidate selection)
+
+Usage:
     python -m onnx_export.export_all --output-dir onnx_models --verify
 """
 
@@ -66,7 +70,17 @@ def main():
         action="store_true",
         help="Skip Vision encoder export",
     )
-    
+    parser.add_argument(
+        "--skip-peaframe",
+        action="store_true",
+        help="Skip PEAFrame export",
+    )
+    parser.add_argument(
+        "--skip-clap",
+        action="store_true",
+        help="Skip CLAP export",
+    )
+
     args = parser.parse_args()
     
     os.makedirs(args.output_dir, exist_ok=True)
@@ -98,7 +112,21 @@ def main():
     if not args.skip_vision:
         export_args = ["--output", args.output_dir, "--model", args.model]
         results["Vision"] = run_export("onnx_export.export_vision", export_args)
-    
+
+    # Export PEAFrame
+    if not args.skip_peaframe:
+        export_args = ["--output-dir", args.output_dir]
+        if args.verify:
+            export_args.append("--verify")
+        results["PEAFrame"] = run_export("onnx_export.export_peaframe", export_args)
+
+    # Export CLAP
+    if not args.skip_clap:
+        export_args = ["--output-dir", args.output_dir]
+        if args.verify:
+            export_args.append("--verify")
+        results["CLAP"] = run_export("onnx_export.export_clap", export_args)
+
     # Print summary
     print(f"\n{'='*60}")
     print("Export Summary")