Deploy hf-space @ c63b38a: batch export + ZeroGPU dynamic duration

.gitignore

@@ -8,7 +8,6 @@ ENV/
 env/
 .venv/
 .ENV/
-
 # Python IDEs
 .idea/
 .vscode/
@@ -29,6 +28,17 @@ tensorboard_logs/
 .DS_Store
 Thumbs.db
 
+# PyCharm files
+*.iml
+.idea/
+
+# Coverage and testing tools
+.coverage
+nosetests.xml
+coverage.xml
+*.cover
+*.log
+
 # Compiled extension modules
 *.so
 *.dylib
@@ -40,3 +50,9 @@ cython_debug/
 # Other custom ignore rules
 *.bak
 *.swp
+
+.ruff_cache/
+
+# Model checkpoints (downloaded at runtime)
+ckpts/
+.gradio/

LICENSE

@@ -0,0 +1,25 @@
+Creative Commons Attribution 4.0 International Public License
+
+By exercising the Licensed Rights (defined below), You accept and agree 
+to be bound by the terms and conditions of this Creative Commons 
+Attribution 4.0 International Public License ("Public License"). 
+To the extent this Public License may be interpreted as a contract, 
+You are granted the Licensed Rights in consideration of Your acceptance 
+of these terms and conditions, and the Licensor grants You such rights 
+in consideration of benefits the Licensor receives from making 
+the Licensed Material available under these terms and conditions.
+
+You are free to:
+- Share — copy and redistribute the material in any medium or format
+- Adapt — remix, transform, and build upon the material for any purpose, even commercially.
+
+Under the following terms:
+- Attribution — You must give appropriate credit, provide a link to the license, 
+  and indicate if changes were made. You may do so in any reasonable manner, 
+  but not in any way that suggests the licensor endorses you or your use.
+
+No additional restrictions — You may not apply legal terms or 
+technological measures that legally restrict others from doing 
+anything the license permits.
+
+Full license text: https://creativecommons.org/licenses/by/4.0/legalcode

README.md

@@ -1,6 +1,6 @@
 ---
 title: SongFormer
-emoji: "\U0001F3B5"
+emoji: "🎵"
 colorFrom: blue
 colorTo: indigo
 sdk: gradio
@@ -40,6 +40,8 @@ Chunbo Hao<sup>&ast;</sup>, Ruibin Yuan<sup>&ast;</sup>, Jixun Yao, Qixin Deng,
 
 SongFormer is a music structure analysis framework that leverages multi-resolution self-supervised representations and heterogeneous supervision, accompanied by the large-scale multilingual dataset SongFormDB and the high-quality benchmark SongFormBench to foster fair and reproducible research.
 
+**This Space offers two modes:** analyze a *Single File* with downloadable results (JSON / MSA / CSV / plot / ZIP), or use the *Batch* tab to process multiple files with live per-file status, a combined ZIP (downloadable mid-run), and a per-file inspector with audio playback. Runs on ZeroGPU — each analyzed file consumes daily GPU quota.
+
 ![](https://github.com/ASLP-lab/SongFormer/blob/main/figs/songformer.png?raw=true)
 
 ## Citation

app.py

@@ -1,30 +1,9 @@
 import os
 import sys
 
-current_file = os.path.abspath(__file__)
-current_dir = os.path.dirname(current_file)
-
-songformer_path = os.path.join(current_dir, "src", "SongFormer")
-if os.path.exists(songformer_path):
-    os.chdir(songformer_path)
-else:
-    print(f"The target working directory does not exist: {songformer_path}")
-
-working_dir = os.getcwd()
-
-third_party_path = os.path.join(current_dir, "src", "third_party")
-if os.path.exists(third_party_path):
-    sys.path.insert(0, third_party_path)
-sys.path.insert(0, working_dir)
-
-musicfm_paths = [
-    os.path.join(current_dir, "src"),
-    os.path.join(current_dir, "third_party"),
-    os.path.join(current_dir, "src", "SongFormer"),
-]
-for path in musicfm_paths:
-    if os.path.exists(path):
-        sys.path.insert(0, path)
+os.chdir(os.path.join("src", "SongFormer"))
+sys.path.append(os.path.join("..", "third_party"))
+sys.path.append(".")
 
 # monkey patch to fix issues in msaf
 import scipy
@@ -39,7 +18,8 @@ import json
 import math
 import importlib
 import matplotlib
-matplotlib.use('Agg')
+
+matplotlib.use("Agg")  # non-interactive backend: safe for rendering plots off the main thread
 import matplotlib.pyplot as plt
 import matplotlib.ticker as ticker
 from pathlib import Path
@@ -51,6 +31,11 @@ from musicfm.model.musicfm_25hz import MusicFM25Hz
 from postprocessing.functional import postprocess_functional_structure
 from dataset.label2id import DATASET_ID_ALLOWED_LABEL_IDS, DATASET_LABEL_TO_DATASET_ID
 from utils.fetch_pretrained import download_all
+
+import export_utils
+
+# ZeroGPU (Hugging Face Spaces). Preinstalled on the Space; this branch
+# is Space-only and never runs locally.
 import spaces
 
 # Constants
@@ -69,6 +54,23 @@ msa_model = None
 device = None
 
 
+def get_device():
+    """Select the best available device: MPS (Apple Silicon), CUDA, or CPU."""
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def clear_device_cache(device):
+    """Clear GPU memory cache for the given device type."""
+    if device.type == "cuda":
+        torch.cuda.empty_cache()
+    elif device.type == "mps":
+        torch.mps.empty_cache()
+
+
 def load_checkpoint(checkpoint_path, device=None):
     """Load checkpoint from path"""
     if device is None:
@@ -90,7 +92,7 @@ def initialize_models(model_name: str, checkpoint: str, config_path: str):
     global muq_model, musicfm_model, msa_model, device
 
     # Set device
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = get_device()
 
     # Load MuQ
     muq_model = MuQ.from_pretrained("OpenMuQ/MuQ-large-msd-iter")
@@ -123,17 +125,26 @@ def initialize_models(model_name: str, checkpoint: str, config_path: str):
     return hp
 
 
-@spaces.GPU(duration=300)
+def _gpu_duration(audio_path, win_size=420, hop_size=420, num_classes=128):
+    """Estimate GPU seconds for one file (ZeroGPU dynamic duration).
+
+    Conservative: 30s base + 0.2s per audio second, clamped to [60, 300].
+    Tune the constants from observed Space timings.
+    """
+    try:
+        audio_secs = librosa.get_duration(path=audio_path)
+    except Exception:
+        return 120
+    return int(min(300, max(60, 30 + 0.2 * audio_secs)))
+
+
+@spaces.GPU(duration=_gpu_duration)
 def process_audio(audio_path, win_size=420, hop_size=420, num_classes=128):
     """Process audio file and return structure analysis results"""
     global muq_model, musicfm_model, msa_model, device
 
     if muq_model is None:
-        hp = initialize_models(
-            model_name="SongFormer",
-            checkpoint="SongFormer.safetensors",
-            config_path="SongFormer.yaml",
-        )
+        hp = initialize_models()
     else:
         hp = OmegaConf.load(os.path.join("configs", "SongFormer.yaml"))
 
@@ -180,14 +191,14 @@ def process_audio(audio_path, win_size=420, hop_size=420, num_classes=128):
             muq_output = muq_model(audio_seg.unsqueeze(0), output_hidden_states=True)
             muq_embd_420s = muq_output["hidden_states"][10]
             del muq_output
-            torch.cuda.empty_cache()
+            clear_device_cache(device)
 
             _, musicfm_hidden_states = musicfm_model.get_predictions(
                 audio_seg.unsqueeze(0)
             )
             musicfm_embd_420s = musicfm_hidden_states[10]
             del musicfm_hidden_states
-            torch.cuda.empty_cache()
+            clear_device_cache(device)
 
             # Process 30-second segments
             wraped_muq_embd_30s = []
@@ -211,14 +222,14 @@ def process_audio(audio_path, win_size=420, hop_size=420, num_classes=128):
                         output_hidden_states=True,
                     )["hidden_states"][10]
                 )
-                torch.cuda.empty_cache()
+                clear_device_cache(device)
 
                 wraped_musicfm_embd_30s.append(
                     musicfm_model.get_predictions(
                         audio[start_idx_30s:end_idx_30s].unsqueeze(0)
                     )[1][10]
                 )
-                torch.cuda.empty_cache()
+                clear_device_cache(device)
 
             if wraped_muq_embd_30s:
                 wraped_muq_embd_30s = torch.concatenate(wraped_muq_embd_30s, dim=1)
@@ -325,6 +336,7 @@ def create_visualization(
     logits, msa_output, label_num=8, frame_rates=AFTER_DOWNSAMPLING_FRAME_RATES
 ):
     """Create visualization plot"""
+    # Assume ID_TO_LABEL mapping exists
     try:
         from dataset.label2id import ID_TO_LABEL
     except:
@@ -420,48 +432,241 @@ def rule_post_processing(msa_list):
     return result
 
 
+def analyze_one(audio_file, out_dir, stem=None):
+    """Run the full per-file analysis pipeline and write export files.
+
+    Shared by the single-file and batch handlers so the two paths cannot
+    drift. Returns (segments, json_str, msa_str, fig, export_paths). The
+    caller owns the returned figure (single-file displays it via gr.Plot;
+    batch saves+closes it); on a write failure the figure is closed here
+    before re-raising so it never leaks.
+    """
+    logits, msa_output = process_audio(audio_file)
+    # Apply rule-based post-processing, if not needed, use in cli infer
+    msa_output = rule_post_processing(msa_output)
+    segments = format_as_segments(msa_output)
+    msa_str = format_as_msa(msa_output)
+    json_str = format_as_json(segments)
+    fig = create_visualization(logits, msa_output)
+    try:
+        export_paths = export_utils.write_exports(
+            audio_file, segments, json_str, msa_str, fig, out_dir, stem=stem
+        )
+    except Exception:
+        plt.close(fig)
+        raise
+    return segments, json_str, msa_str, fig, export_paths
+
+
 def process_and_analyze(audio_file):
     """Main processing function"""
 
-    def format_time(t: float) -> str:
-        minutes = int(t // 60)
-        seconds = t % 60
-        return f"{minutes:02d}:{seconds:06.3f}"
-
     if audio_file is None:
-        return None, "", "", None
+        return None, "", "", None, None, None, None, None, None
 
     try:
-        # Process audio
-        logits, msa_output = process_audio(audio_file)
-        # Apply rule-based post-processing
-        msa_output = rule_post_processing(msa_output)
-        # Format outputs
-        segments = format_as_segments(msa_output)
-        msa_format = format_as_msa(msa_output)
-        json_format = format_as_json(segments)
+        # Shared pipeline; exports land in a fresh per-run temp directory
+        # (stale runs are swept automatically by the bootstrap).
+        out_dir = export_utils.new_run_dir()
+        segments, json_format, msa_format, fig, export_paths = analyze_one(
+            audio_file, out_dir
+        )
 
         # Create table data
-        table_data = [
-            [
-                f"{float(seg['start']):.2f} ({format_time(float(seg['start']))})",
-                f"{float(seg['end']):.2f} ({format_time(float(seg['end']))})",
-                seg["label"],
-            ]
-            for seg in segments
-        ]
+        table_data = export_utils.segments_to_table(segments)
 
-        # Create visualization
-        fig = create_visualization(logits, msa_output)
-
-        return table_data, json_format, msa_format, fig
+        zip_path = os.path.join(
+            out_dir, export_utils.stem_of(audio_file) + "_songformer.zip"
+        )
+        export_utils.make_zip(list(export_paths.values()), zip_path)
+
+        return (
+            table_data,
+            json_format,
+            msa_format,
+            fig,
+            export_paths["json"],
+            export_paths["msa"],
+            export_paths["csv"],
+            export_paths["png"],
+            zip_path,
+        )
 
     except Exception as e:
         import traceback
 
         error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
-        print(error_msg)
-        return None, "", error_msg, None
+        print(error_msg)  # 在命令行输出完整错误
+        return None, "", error_msg, None, None, None, None, None, None
+
+
+def process_batch(files):
+    """Analyze multiple files sequentially, yielding live status.
+
+    The status table itself is the progress display: every file is listed
+    as queued upfront, flips to processing, then to done/failed. Dropdown
+    choices update as files finish so completed results can be inspected
+    while the rest of the batch is still running.
+
+    Outputs (per yield): status rows, ZIP download update, file-selector
+    update, per-file results dict (for the detail viewer).
+    """
+    if not files:
+        yield (
+            [["(no files uploaded)", "", "", ""]],
+            gr.update(value=None),
+            gr.update(choices=[], value=None),
+            {},
+        )
+        return
+
+    run_dir = export_utils.new_run_dir()
+    bundle = os.path.join(run_dir, "bundle")
+    os.makedirs(bundle, exist_ok=True)
+
+    # De-duplicate stems upfront so same-named uploads don't overwrite each
+    # other and the queued list shows the final names.
+    used_stems = set()
+    queue = []
+    for audio_file in files:
+        base = export_utils.stem_of(audio_file)
+        stem = base
+        n = 2
+        while stem in used_stems:
+            stem = f"{base}_{n}"
+            n += 1
+        used_stems.add(stem)
+        queue.append((audio_file, stem))
+
+    status_rows = [[stem, "⏳ queued", "", ""] for _, stem in queue]
+    results = {}
+    zipped_count = 0  # how many files the on-disk ZIP actually contains
+    zip_path = os.path.join(run_dir, "songformer_batch.zip")
+
+    def _rebuild_bundle_zip():
+        """Rewrite manifests and atomically swap in an updated ZIP.
+
+        Called after each completed file so the download button always
+        serves "everything so far". os.replace is atomic, so a click can
+        never observe a half-written archive. The (stem, segments) pairs
+        are derived from `results` — the single source of truth.
+        """
+        named = [(s, r["segments"]) for s, r in results.items()]
+        with open(
+            os.path.join(bundle, "summary.csv"), "w", encoding="utf-8", newline=""
+        ) as f:
+            f.write(export_utils.segments_to_combined_csv(named))
+        with open(
+            os.path.join(bundle, "combined.json"), "w", encoding="utf-8"
+        ) as f:
+            f.write(export_utils.combined_json(named))
+        part = zip_path + ".part"
+        export_utils.zip_dir(bundle, part)
+        os.replace(part, zip_path)
+
+    # List every file as queued; clear any previous run's results
+    yield (
+        status_rows,
+        gr.update(value=None, interactive=False, label="⬇️ Download all (ZIP)"),
+        gr.update(choices=[], value=None),
+        {},
+    )
+
+    for idx, (audio_file, stem) in enumerate(queue):
+        status_rows[idx] = [stem, "🔄 processing…", "", ""]
+        yield status_rows, gr.update(), gr.update(), results
+        try:
+            file_dir = os.path.join(bundle, stem)
+            os.makedirs(file_dir, exist_ok=True)
+            segments, json_str, msa_str, fig, paths = analyze_one(
+                audio_file, file_dir, stem=stem
+            )
+            plt.close(fig)
+            duration = (
+                export_utils.format_time(float(segments[-1]["end"]))
+                if segments
+                else ""
+            )
+            status_rows[idx] = [stem, "✅", len(segments), duration]
+            results[stem] = {
+                "segments": segments,
+                "json": json_str,
+                "msa": msa_str,
+                "png": paths["png"],
+                "audio": audio_file,
+            }
+        except Exception as e:
+            import traceback
+
+            print(f"Batch error for {stem}:\n{traceback.format_exc()}")
+            status_rows[idx] = [stem, "❌ " + str(e)[:80], 0, ""]
+            # ZeroGPU quota exhausted: every remaining file would fail the
+            # same way, so skip them. (Message heuristic — ZeroGPU does not
+            # document a stable exception class.)
+            if "quota" in str(e).lower():
+                for j in range(idx + 1, len(queue)):
+                    status_rows[j] = [queue[j][1], "⏭️ skipped (GPU quota)", "", ""]
+                yield (
+                    status_rows,
+                    gr.update(),
+                    gr.update(choices=list(results.keys())),
+                    results,
+                )
+                break
+        else:
+            # A ZIP rebuild failure must NOT mark the analyzed file as
+            # failed: its exports exist and the next successful rebuild
+            # will include it (pairs derive from `results`).
+            try:
+                # Keep the ZIP downloadable mid-run with everything so far
+                _rebuild_bundle_zip()
+                zipped_count = len(results)
+            except Exception:
+                import traceback
+
+                print(f"ZIP rebuild error after {stem}:\n{traceback.format_exc()}")
+        if zipped_count:
+            zip_update = gr.update(
+                value=zip_path,
+                interactive=True,
+                label=f"⬇️ Download all (ZIP) — {zipped_count}/{len(queue)} files",
+            )
+        else:
+            zip_update = gr.update()
+        # Completed files become inspectable while the batch continues
+        yield status_rows, zip_update, gr.update(choices=list(results.keys())), results
+
+    # Manifests + ZIP were rebuilt incrementally per file; just normalize
+    # the button label now that the batch is complete. The button is only
+    # active if at least one rebuild actually produced a ZIP on disk.
+    yield (
+        status_rows,
+        gr.update(
+            value=zip_path if zipped_count else None,
+            interactive=bool(zipped_count),
+            label="⬇️ Download all (ZIP)",
+        ),
+        gr.update(choices=list(results.keys())),
+        results,
+    )
+
+
+def on_select_file(stem, results):
+    """Render a previously-computed file's result in the batch detail viewer."""
+    # A selection can race an in-flight batch iteration under rare scheduler
+    # timings (choices reach the browser just before the state lands); the
+    # guard degrades to an empty view, recoverable by re-selecting.
+    results = results or {}
+    if not stem or stem not in results:
+        return None, "", "", None, None
+    r = results[stem]
+    return (
+        export_utils.segments_to_table(r["segments"]),
+        r["json"],
+        r["msa"],
+        r["png"],
+        r.get("audio"),
+    )
 
 
 # Create Gradio interface
@@ -503,78 +708,173 @@ with gr.Blocks(
     # Links
     gr.HTML("""
         <div class="links-container">
-            <a href="https://img.shields.io/badge/Python-3.10-brightgreen"><img src="https://img.shields.io/badge/Python-3.10-brightgreen" alt="Python 3.10"></a>
-            <a href="https://img.shields.io/badge/License-CC%20BY%204.0-lightblue"><img src="https://img.shields.io/badge/License-CC%20BY%204.0-lightblue" alt="License CC BY 4.0"></a>
-            <a href="https://arxiv.org/abs/2510.02797"><img src="https://img.shields.io/badge/arXiv-2510.02797-blue" alt="arXiv Paper"></a>
-            <a href="https://github.com/ASLP-lab/SongFormer"><img src="https://img.shields.io/badge/GitHub-SongFormer-black" alt="GitHub"></a>
-            <a href="https://huggingface.co/spaces/SidSaxena/SongFormer"><img src="https://img.shields.io/badge/HuggingFace-space-yellow" alt="HuggingFace Space"></a>
-            <a href="https://huggingface.co/ASLP-lab/SongFormer"><img src="https://img.shields.io/badge/HuggingFace-model-blue" alt="HuggingFace Model"></a>
-            <a href="https://huggingface.co/datasets/ASLP-lab/SongFormDB"><img src="https://img.shields.io/badge/HF%20Dataset-SongFormDB-green" alt="Dataset SongFormDB"></a>
-            <a href="https://huggingface.co/datasets/ASLP-lab/SongFormBench"><img src="https://img.shields.io/badge/HF%20Dataset-SongFormBench-orange" alt="Dataset SongFormBench"></a>
-            <a href="https://discord.gg/p5uBryC4Zs"><img src="https://img.shields.io/badge/Discord-join%20us-purple?logo=discord&logoColor=white" alt="Discord"></a>
-            <a href="http://www.npu-aslp.org/"><img src="https://img.shields.io/badge/%F0%9F%8F%AB-ASLP-grey?labelColor=lightgrey" alt="ASLP lab"></a>
+            <img src="https://img.shields.io/badge/Python-3.10-brightgreen" alt="Python">
+            <img src="https://img.shields.io/badge/License-CC%20BY%204.0-lightblue" alt="License">
+            <a href="https://arxiv.org/abs/2510.02797">
+            <img src="https://img.shields.io/badge/arXiv-2510.02797-blue" alt="arXiv">
+            </a>
+            <a href="https://github.com/ASLP-lab/SongFormer">
+            <img src="https://img.shields.io/badge/GitHub-SongFormer-black" alt="GitHub">
+            </a>
+            <a href="https://huggingface.co/spaces/SidSaxena/SongFormer">
+            <img src="https://img.shields.io/badge/HuggingFace-space-yellow" alt="HuggingFace Space">
+            </a>
+            <a href="https://huggingface.co/ASLP-lab/SongFormer">
+            <img src="https://img.shields.io/badge/HuggingFace-model-blue" alt="HuggingFace Model">
+            </a>
+            <a href="https://huggingface.co/datasets/ASLP-lab/SongFormDB">
+            <img src="https://img.shields.io/badge/HF%20Dataset-SongFormDB-green" alt="Dataset SongFormDB">
+            </a>
+            <a href="https://huggingface.co/datasets/ASLP-lab/SongFormBench">
+            <img src="https://img.shields.io/badge/HF%20Dataset-SongFormBench-orange" alt="Dataset SongFormBench">
+            </a>
+            <a href="https://discord.gg/p5uBryC4Zs">
+            <img src="https://img.shields.io/badge/Discord-join%20us-purple?logo=discord&logoColor=white" alt="Discord">
+            </a>
+            <a href="http://www.npu-aslp.org/">
+            <img src="https://img.shields.io/badge/🏫-ASLP-grey?labelColor=lightgrey" alt="ASLP">
+            </a>
         </div>
     """)
 
+    with gr.Tabs():
+        with gr.Tab("Single File"):
+            # Main input area
+            with gr.Row():
+                with gr.Column(scale=3):
+                    audio_input = gr.Audio(
+                        label="Upload Audio File", type="filepath", elem_id="audio-input"
+                    )
 
-    # Main input area
-    with gr.Row():
-        with gr.Column(scale=3):
-            audio_input = gr.Audio(
-                label="Upload Audio File", type="filepath", elem_id="audio-input"
-            )
+                with gr.Column(scale=1):
+                    gr.Markdown("### 📌 Examples")
+                    gr.Examples(
+                        examples=[
+                            ["examples/BC_5cd6a6.mp3"],
+                            ["examples/BC_282ece.mp3"],
+                            ["examples/BHX_0158_letitrock.wav"],
+                            ["examples/BHX_0374_drunkonyou.wav"],
+                        ],
+                        inputs=[audio_input],
+                        label="Click to load example",
+                    )
 
-        with gr.Column(scale=1):
-            gr.Markdown("### Examples")
-            gr.Examples(
-                examples=[
-                    ["examples/BC_5cd6a6.mp3"],
-                    ["examples/BC_282ece.mp3"],
-                    ["examples/BHX_0158_letitrock.wav"],
-                    ["examples/BHX_0374_drunkonyou.wav"],
-                ],
-                inputs=[audio_input],
-                label="Click to load example",
-            )
+            # Analyze button
+            with gr.Row():
+                analyze_btn = gr.Button(
+                    "🚀 Analyze Music Structure", variant="primary", scale=1
+                )
 
-    # Analyze button
-    with gr.Row():
-        analyze_btn = gr.Button(
-            "Analyze Music Structure", variant="primary", scale=1
-        )
+            # Results display area
+            with gr.Row():
+                with gr.Column(scale=13):
+                    segments_table = gr.Dataframe(
+                        headers=["Start / s (m:s.ms)", "End / s (m:s.ms)", "Label"],
+                        label="Detected Music Segments",
+                        interactive=False,
+                        elem_id="result-table",
+                    )
+                with gr.Column(scale=8):
+                    with gr.Row():
+                        with gr.Accordion("📄 JSON Output", open=False):
+                            json_output = gr.Textbox(
+                                label="JSON Format",
+                                lines=15,
+                                max_lines=20,
+                                interactive=False,
+                                show_copy_button=True,
+                            )
+                    with gr.Row():
+                        with gr.Accordion("📋 MSA Text Output", open=False):
+                            msa_output = gr.Textbox(
+                                label="MSA Format",
+                                lines=15,
+                                max_lines=20,
+                                interactive=False,
+                                show_copy_button=True,
+                            )
+
+            # Visualization plot
+            with gr.Row():
+                plot_output = gr.Plot(label="Activation Curves Visualization")
 
-    # Results display area
-    with gr.Row():
-        with gr.Column(scale=13):
-            segments_table = gr.Dataframe(
-                headers=["Start / s (m:s.ms)", "End / s (m:s.ms)", "Label"],
-                label="Detected Music Segments",
-                interactive=False,
-                elem_id="result-table",
+            # Export / download buttons (populated after analysis)
+            with gr.Row():
+                download_json_btn = gr.DownloadButton("⬇️ JSON")
+                download_msa_btn = gr.DownloadButton("⬇️ MSA (.txt)")
+                download_csv_btn = gr.DownloadButton("⬇️ CSV")
+                download_png_btn = gr.DownloadButton("⬇️ Plot (.png)")
+                download_zip_btn = gr.DownloadButton(
+                    "⬇️ Download all (ZIP)", variant="primary"
+                )
+
+        with gr.Tab("Batch"):
+            gr.Markdown(
+                "Upload multiple audio files, analyze them sequentially, "
+                "and download all results as a single ZIP.\n\n"
+                "*This Space runs on ZeroGPU: each file consumes your daily "
+                "GPU quota (2–40 min depending on account tier). The ZIP "
+                "below always contains everything analyzed so far.*"
             )
-        with gr.Column(scale=8):
             with gr.Row():
-                with gr.Accordion("JSON Output", open=False):
-                    json_output = gr.Textbox(
-                        label="JSON Format",
-                        lines=15,
-                        max_lines=20,
-                        interactive=False,
-                        show_copy_button=True,
+                with gr.Column(scale=3):
+                    batch_files = gr.File(
+                        label="Upload Audio Files",
+                        file_count="multiple",
+                        type="filepath",
                     )
+                with gr.Column(scale=1):
+                    batch_analyze_btn = gr.Button(
+                        "🚀 Analyze Batch", variant="primary"
+                    )
+                    batch_zip_btn = gr.DownloadButton(
+                        "⬇️ Download all (ZIP)", variant="primary", interactive=False
+                    )
+            with gr.Row():
+                batch_status = gr.Dataframe(
+                    headers=["File", "Status", "Segments", "Duration"],
+                    label="Batch Status",
+                    interactive=False,
+                )
+            batch_results_state = gr.State({})
+            gr.Markdown("### Inspect a file")
             with gr.Row():
-                with gr.Accordion("MSA Text Output", open=False):
-                    msa_output = gr.Textbox(
-                        label="MSA Format",
-                        lines=15,
-                        max_lines=20,
+                with gr.Column(scale=1):
+                    batch_file_selector = gr.Dropdown(
+                        label="Processed File", choices=[], interactive=True
+                    )
+                with gr.Column(scale=2):
+                    batch_detail_audio = gr.Audio(
+                        label="Listen", type="filepath", interactive=False
+                    )
+            with gr.Row():
+                with gr.Column(scale=13):
+                    batch_detail_table = gr.Dataframe(
+                        headers=["Start / s (m:s.ms)", "End / s (m:s.ms)", "Label"],
+                        label="Detected Music Segments",
                         interactive=False,
-                        show_copy_button=True,
                     )
-
-    # Visualization plot
-    with gr.Row():
-        plot_output = gr.Plot(label="Activation Curves Visualization")
+                with gr.Column(scale=8):
+                    with gr.Row():
+                        with gr.Accordion("📄 JSON Output", open=False):
+                            batch_detail_json = gr.Textbox(
+                                label="JSON Format",
+                                lines=15,
+                                max_lines=20,
+                                interactive=False,
+                                show_copy_button=True,
+                            )
+                    with gr.Row():
+                        with gr.Accordion("📋 MSA Text Output", open=False):
+                            batch_detail_msa = gr.Textbox(
+                                label="MSA Format",
+                                lines=15,
+                                max_lines=20,
+                                interactive=False,
+                                show_copy_button=True,
+                            )
+            with gr.Row():
+                batch_detail_plot = gr.Image(label="Activation Curves Visualization")
 
     gr.HTML("""
         <div style="display: flex; justify-content: center; align-items: center;">
@@ -586,7 +886,39 @@ with gr.Blocks(
     analyze_btn.click(
         fn=process_and_analyze,
         inputs=[audio_input],
-        outputs=[segments_table, json_output, msa_output, plot_output],
+        outputs=[
+            segments_table,
+            json_output,
+            msa_output,
+            plot_output,
+            download_json_btn,
+            download_msa_btn,
+            download_csv_btn,
+            download_png_btn,
+            download_zip_btn,
+        ],
+    )
+    batch_analyze_btn.click(
+        fn=process_batch,
+        inputs=[batch_files],
+        outputs=[
+            batch_status,
+            batch_zip_btn,
+            batch_file_selector,
+            batch_results_state,
+        ],
+        show_progress="minimal",
+    )
+    batch_file_selector.change(
+        fn=on_select_file,
+        inputs=[batch_file_selector, batch_results_state],
+        outputs=[
+            batch_detail_table,
+            batch_detail_json,
+            batch_detail_msa,
+            batch_detail_plot,
+            batch_detail_audio,
+        ],
     )
 
 if __name__ == "__main__":
@@ -601,5 +933,6 @@ if __name__ == "__main__":
     )
     print("Models loaded successfully!")
 
-    # Launch interface
+    # Launch interface (Spaces injects its own server settings; an explicit
+    # port would break the platform health check)
     demo.launch()

export_utils.py

@@ -0,0 +1,206 @@
+"""Serialize SongFormer analysis results to downloadable files.
+
+Pure, UI-agnostic helpers used by app.py. No model or Gradio imports, so
+these can be unit-tested without loading any checkpoint.
+"""
+
+import csv
+import io
+import json
+import os
+import shutil
+import tempfile
+import time
+import zipfile
+
+# Per-run export directories older than this (seconds) are swept at the start
+# of each analysis. Recent runs are kept so their download files stay servable.
+DEFAULT_EXPORT_TTL_SECONDS = 3600
+
+
+def format_time(t: float) -> str:
+    """Render seconds as mm:ss.mmm (e.g. 61.5 -> '01:01.500')."""
+    minutes = int(t // 60)
+    seconds = t % 60
+    return f"{minutes:02d}:{seconds:06.3f}"
+
+
+def stem_of(audio_path: str) -> str:
+    """Return the audio filename without directory or extension."""
+    return os.path.splitext(os.path.basename(audio_path))[0]
+
+
+def segments_to_table(segments) -> list:
+    """Build display table rows: [start "(mm:ss.mmm)", end "(mm:ss.mmm)", label]."""
+    rows = []
+    for seg in segments:
+        start = float(seg["start"])
+        end = float(seg["end"])
+        rows.append(
+            [
+                f"{start:.2f} ({format_time(start)})",
+                f"{end:.2f} ({format_time(end)})",
+                seg["label"],
+            ]
+        )
+    return rows
+
+
+def segments_to_csv(segments) -> str:
+    """Build CSV text from segment dicts.
+
+    Each segment is {"start": str|float, "end": str|float, "label": str}.
+    Columns: start_sec, start_mmss, end_sec, end_mmss, label.
+    """
+    buf = io.StringIO()
+    writer = csv.writer(buf, lineterminator="\n")
+    writer.writerow(["start_sec", "start_mmss", "end_sec", "end_mmss", "label"])
+    for seg in segments:
+        start = float(seg["start"])
+        end = float(seg["end"])
+        writer.writerow(
+            [
+                f"{start:.2f}",
+                format_time(start),
+                f"{end:.2f}",
+                format_time(end),
+                seg["label"],
+            ]
+        )
+    return buf.getvalue()
+
+
+def segments_to_combined_csv(named) -> str:
+    """Build a combined CSV across files.
+
+    `named` is a list of (filename, segments). Columns:
+    filename, start_sec, start_mmss, end_sec, end_mmss, label.
+    """
+    buf = io.StringIO()
+    writer = csv.writer(buf, lineterminator="\n")
+    writer.writerow(
+        ["filename", "start_sec", "start_mmss", "end_sec", "end_mmss", "label"]
+    )
+    for filename, segments in named:
+        for seg in segments:
+            start = float(seg["start"])
+            end = float(seg["end"])
+            writer.writerow(
+                [
+                    filename,
+                    f"{start:.2f}",
+                    format_time(start),
+                    f"{end:.2f}",
+                    format_time(end),
+                    seg["label"],
+                ]
+            )
+    return buf.getvalue()
+
+
+def combined_json(named) -> str:
+    """Build a combined JSON mapping {filename: segments} across files."""
+    return json.dumps(
+        {filename: segments for filename, segments in named},
+        indent=2,
+        ensure_ascii=False,
+    )
+
+
+def write_exports(audio_path, segments, json_str, msa_str, fig, out_dir, stem=None) -> dict:
+    """Write json/msa/csv/png into out_dir; return {format: path}.
+
+    Reuses the already-built json_str/msa_str from app.py rather than
+    re-serializing. Saves the matplotlib figure as PNG. `stem` overrides the
+    filename stem (used by batch to keep de-duplicated folder and file names
+    consistent); defaults to the audio filename's stem.
+    """
+    if stem is None:
+        stem = stem_of(audio_path)
+    paths = {
+        "json": os.path.join(out_dir, f"{stem}.json"),
+        "msa": os.path.join(out_dir, f"{stem}.msa.txt"),
+        "csv": os.path.join(out_dir, f"{stem}.csv"),
+        "png": os.path.join(out_dir, f"{stem}.png"),
+    }
+    with open(paths["json"], "w", encoding="utf-8") as f:
+        f.write(json_str)
+    with open(paths["msa"], "w", encoding="utf-8") as f:
+        f.write(msa_str)
+    with open(paths["csv"], "w", encoding="utf-8", newline="") as f:
+        f.write(segments_to_csv(segments))
+    fig.savefig(paths["png"], dpi=150, bbox_inches="tight")
+    return paths
+
+
+def make_zip(paths, zip_path) -> str:
+    """Bundle the given files into zip_path using their basenames.
+
+    Returns zip_path.
+    """
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
+        for p in paths:
+            zf.write(p, arcname=os.path.basename(p))
+    return zip_path
+
+
+# File types that are already compressed: deflating them again wastes CPU
+# (which matters because the batch ZIP is rebuilt incrementally per file).
+_STORED_EXTENSIONS = {".png", ".jpg", ".jpeg", ".zip", ".mp3", ".flac", ".ogg"}
+
+
+def zip_dir(src_dir, zip_path) -> str:
+    """Zip the contents of src_dir into zip_path.
+
+    Arcnames are relative to src_dir, preserving subfolders. Files that are
+    already compressed (see _STORED_EXTENSIONS) are stored uncompressed.
+    Returns zip_path.
+    """
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
+        for root, _dirs, files in os.walk(src_dir):
+            for name in files:
+                full = os.path.join(root, name)
+                arcname = os.path.relpath(full, src_dir)
+                compress = (
+                    zipfile.ZIP_STORED
+                    if os.path.splitext(name)[1].lower() in _STORED_EXTENSIONS
+                    else zipfile.ZIP_DEFLATED
+                )
+                zf.write(full, arcname=arcname, compress_type=compress)
+    return zip_path
+
+
+def cleanup_old_exports(parent_dir, max_age_seconds, now=None) -> list:
+    """Remove run subdirectories of parent_dir older than max_age_seconds.
+
+    Only directories are swept (stray files are left alone). A missing
+    parent_dir is a no-op. Recent runs are preserved so their download files
+    remain servable. Returns the list of removed directory paths.
+    """
+    if now is None:
+        now = time.time()
+    removed = []
+    if not os.path.isdir(parent_dir):
+        return removed
+    cutoff = now - max_age_seconds
+    for name in sorted(os.listdir(parent_dir)):
+        path = os.path.join(parent_dir, name)
+        if not os.path.isdir(path):
+            continue
+        if os.path.getmtime(path) < cutoff:
+            shutil.rmtree(path, ignore_errors=True)
+            removed.append(path)
+    return removed
+
+
+def new_run_dir(parent_dir=None, ttl_seconds=DEFAULT_EXPORT_TTL_SECONDS) -> str:
+    """Create a fresh run directory for export files, sweeping stale runs.
+
+    Shared bootstrap for the single-file and batch handlers. parent_dir
+    defaults to <system tempdir>/songformer_exports.
+    """
+    if parent_dir is None:
+        parent_dir = os.path.join(tempfile.gettempdir(), "songformer_exports")
+    os.makedirs(parent_dir, exist_ok=True)
+    cleanup_old_exports(parent_dir, ttl_seconds)
+    return tempfile.mkdtemp(prefix="run_", dir=parent_dir)

src/SongFormer/infer.sh

@@ -16,6 +16,8 @@ python infer/infer.py \
 --checkpoint SongFormer.safetensors \
 --config_path SongFormer.yaml \
 -gn 1 \
--tn 1
+-tn 1 \
+--save_plots \
+--device mps
 # --debug
 # --no_rule_post_processing

src/SongFormer/infer/infer.py

@@ -22,9 +22,33 @@ from muq import MuQ
 from musicfm.model.musicfm_25hz import MusicFM25Hz
 from omegaconf import OmegaConf
 from tqdm import tqdm
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import matplotlib.ticker as ticker
 
 mp.set_start_method("spawn", force=True)
 
+
+def get_device(device_override=None, rank=0):
+    """Select the best available device: MPS (Apple Silicon), CUDA, or CPU."""
+    if device_override:
+        return torch.device(device_override)
+    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        return torch.device("mps")
+    if torch.cuda.is_available():
+        return torch.device(f"cuda:{rank}")
+    return torch.device("cpu")
+
+
+def clear_device_cache(device):
+    """Clear GPU memory cache for the given device type."""
+    if device.type == "cuda":
+        torch.cuda.empty_cache()
+    elif device.type == "mps":
+        torch.mps.empty_cache()
+
+
 MUSICFM_HOME_PATH = os.path.join("ckpts", "MusicFM")
 
 BEFORE_DOWNSAMPLING_FRAME_RATES = 25
@@ -36,7 +60,7 @@ DATASET_IDS = [5]
 TIME_DUR = 420
 INPUT_SAMPLING_RATE = 24000
 
-from dataset.label2id import DATASET_ID_ALLOWED_LABEL_IDS, DATASET_LABEL_TO_DATASET_ID
+from dataset.label2id import DATASET_ID_ALLOWED_LABEL_IDS, DATASET_LABEL_TO_DATASET_ID, ID_TO_LABEL
 from postprocessing.functional import postprocess_functional_structure
 
 
@@ -60,6 +84,112 @@ def get_processing_ids(input_path, processed_ids_set):
     return ret
 
 
+def save_visualizations(logits, msa_infer_output, data_id, output_dir, label_num=8):
+    """Save activation visualizations (line plot PDF + heatmap PNG) and raw logits (.npy)."""
+    try:
+        function_vals = logits["function_logits"].squeeze().cpu().numpy()  # [T, 128]
+        boundary_vals = logits["boundary_logits"].squeeze().cpu().numpy()  # [T]
+
+        # --- Raw logits ---
+        logits_dir = os.path.join(output_dir, "logits")
+        os.makedirs(logits_dir, exist_ok=True)
+        np.save(os.path.join(logits_dir, f"{data_id}_function.npy"), function_vals)
+        np.save(os.path.join(logits_dir, f"{data_id}_boundary.npy"), boundary_vals)
+
+        # --- Plots ---
+        plots_dir = os.path.join(output_dir, "plots")
+        os.makedirs(plots_dir, exist_ok=True)
+
+        T = function_vals.shape[0]
+        time_axis = np.arange(T) / AFTER_DOWNSAMPLING_FRAME_RATES
+        top_classes = np.argsort(function_vals.mean(axis=0))[-label_num:]
+
+        # A. Line plot (PDF)
+        fig, ax = plt.subplots(2, 1, figsize=(15, 8), sharex=True)
+
+        ax[0].plot(time_axis, boundary_vals, label="Boundary logit", color="orange")
+        ax[0].set_title("Boundary logits")
+        ax[0].set_ylabel("Logit")
+        ax[0].legend()
+        ax[0].grid(True)
+
+        for cls in top_classes:
+            ax[1].plot(
+                time_axis,
+                function_vals[:, cls],
+                label=f"{ID_TO_LABEL.get(cls, f'Class_{cls}')}",
+            )
+        ax[1].set_title(f"Top {label_num} Function logits by mean activation")
+        ax[1].set_xlabel("Time (seconds)")
+        ax[1].set_ylabel("Logit")
+        ax[1].xaxis.set_major_locator(ticker.MultipleLocator(20))
+        ax[1].xaxis.set_minor_locator(ticker.MultipleLocator(5))
+        ax[1].xaxis.set_major_formatter(ticker.FormatStrFormatter("%.1f"))
+        ax[1].legend()
+        ax[1].grid(True)
+
+        for t_sec, label in msa_infer_output:
+            for a in ax:
+                a.axvline(x=t_sec, color="red", linestyle="--", linewidth=0.8)
+            if label != "end":
+                ax[1].text(
+                    t_sec + 0.3,
+                    ax[1].get_ylim()[1] * 0.85,
+                    label,
+                    rotation=90,
+                    fontsize=8,
+                    color="red",
+                )
+
+        plt.suptitle(f"{data_id} — MSA Logits Overview", fontsize=16)
+        plt.tight_layout()
+        plt.savefig(os.path.join(plots_dir, f"{data_id}.pdf"), bbox_inches="tight")
+        plt.close(fig)
+
+        # B. Heatmap (PNG)
+        fig, ax = plt.subplots(figsize=(15, 6))
+        class_labels = [ID_TO_LABEL.get(cls, f"Class_{cls}") for cls in top_classes]
+        heatmap_data = function_vals[:, top_classes]  # [T, label_num]
+
+        im = ax.imshow(
+            heatmap_data.T,
+            aspect="auto",
+            origin="lower",
+            cmap="viridis",
+            extent=[time_axis[0], time_axis[-1], -0.5, label_num - 0.5],
+        )
+        ax.set_yticks(range(label_num))
+        ax.set_yticklabels(class_labels)
+        ax.set_xlabel("Time (seconds)")
+        ax.set_ylabel("Class")
+        ax.set_title(f"{data_id} — Function Logits Heatmap (Top {label_num})")
+
+        for t_sec, label in msa_infer_output:
+            ax.axvline(x=t_sec, color="white", linestyle="--", linewidth=0.8, alpha=0.7)
+            if label != "end":
+                ax.text(
+                    t_sec + 0.3,
+                    label_num - 0.7,
+                    label,
+                    rotation=90,
+                    fontsize=7,
+                    color="white",
+                    alpha=0.9,
+                )
+
+        plt.colorbar(im, ax=ax, label="Logit value")
+        plt.tight_layout()
+        plt.savefig(
+            os.path.join(plots_dir, f"{data_id}_heatmap.png"),
+            dpi=150,
+            bbox_inches="tight",
+        )
+        plt.close(fig)
+
+    except Exception as e:
+        logger.error(f"Visualization failed for {data_id}: {e}")
+
+
 def load_checkpoint(checkpoint_path, device=None):
     """Load checkpoint from path"""
     if device is None:
@@ -115,7 +245,7 @@ def rule_post_processing(msa_list):
 
 def inference(rank, queue_input: mp.Queue, queue_output: mp.Queue, args):
     """Run inference on the input audio"""
-    device = f"cuda:{rank}"
+    device = get_device(getattr(args, 'device', None), rank)
 
     # MuQ model loading (this will automatically fetch the checkpoint from huggingface)
     muq = MuQ.from_pretrained("OpenMuQ/MuQ-large-msd-iter")
@@ -199,7 +329,7 @@ def inference(rank, queue_input: mp.Queue, queue_output: mp.Queue, args):
                     muq_output = muq(audio_seg.unsqueeze(0), output_hidden_states=True)
                     muq_embd_420s = muq_output["hidden_states"][10]
                     del muq_output
-                    torch.cuda.empty_cache()
+                    clear_device_cache(device)
 
                     # MusicFM embedding
                     _, musicfm_hidden_states = musicfm.get_predictions(
@@ -207,7 +337,7 @@ def inference(rank, queue_input: mp.Queue, queue_output: mp.Queue, args):
                     )
                     musicfm_embd_420s = musicfm_hidden_states[10]
                     del musicfm_hidden_states
-                    torch.cuda.empty_cache()
+                    clear_device_cache(device)
 
                     wraped_muq_embd_30s = []
                     wraped_musicfm_embd_30s = []
@@ -229,13 +359,13 @@ def inference(rank, queue_input: mp.Queue, queue_output: mp.Queue, args):
                                 output_hidden_states=True,
                             )["hidden_states"][10]
                         )
-                        torch.cuda.empty_cache()
+                        clear_device_cache(device)
                         wraped_musicfm_embd_30s.append(
                             musicfm.get_predictions(
                                 audio[start_idx_30s:end_idx_30s].unsqueeze(0)
                             )[1][10]
                         )
-                        torch.cuda.empty_cache()
+                        clear_device_cache(device)
 
                     wraped_muq_embd_30s = torch.concatenate(wraped_muq_embd_30s, dim=1)
                     wraped_musicfm_embd_30s = torch.concatenate(
@@ -326,6 +456,11 @@ def inference(rank, queue_input: mp.Queue, queue_output: mp.Queue, args):
                     ensure_ascii=False,
                 )
 
+                if args.save_plots:
+                    save_visualizations(
+                        logits, msa_infer_output, Path(item).stem, args.output_dir
+                    )
+
                 queue_output.put(None)
 
             except Exception as e:
@@ -368,6 +503,8 @@ def main(args):
         checkpoint=args.checkpoint,
         config_path=args.config_path,
         no_rule_post_processing=args.no_rule_post_processing,
+        device=getattr(args, 'device', None),
+        save_plots=args.save_plots,
     )
 
     processes = []
@@ -433,6 +570,12 @@ if __name__ == "__main__":
         help="Disable rule-based post-processing",
     )
     parser.add_argument("--debug", action="store_true", help="Enable debug mode")
+    parser.add_argument(
+        "--save_plots",
+        action="store_true",
+        help="Save activation visualizations (line plots + heatmaps) and raw logits (.npy)",
+    )
+    parser.add_argument("--device", type=str, default=None, help="Force device (mps, cuda, cpu). Auto-detected if not set.")
 
     args = parser.parse_args()
 

src/third_party/MuQ/.gitattributes

@@ -1,2 +0,0 @@
-# Auto detect text files and perform LF normalization
-* text=auto

src/third_party/MuQ/.gitignore

@@ -1,46 +0,0 @@
-# Byte-compiled / optimized / DLL files
-__pycache__/
-*.py[cod]
-*$py.class
-*.egg*/
-*pyc
-
-# Distribution / packaging
-.Python
-env/
-build/
-dist/
-*.log
-
-# pyenv
-.python-version
-
-# dotenv
-.env
-
-# virtualenv
-.venv/
-venv/
-ENV/
-
-# VSCode settings
-.vscode
-
-# IDEA files
-.idea
-
-# OSX dir files
-.DS_Store
-
-# Sublime Text settings
-*.sublime-workspace
-*.sublime-project
-
-# custom
-open/
-src/recipes/pretrain/dataset/music4all/*.json
-src/recipes/contrastive_learning/datasets/mtg-jamendo/*.json
-runs/
-output/
-logs
-outputs/

src/third_party/MuQ/.gitmodules

@@ -1,3 +0,0 @@
-[submodule "src/recipes/pretrain/fairseq"]
-	path = src/recipes/pretrain/fairseq
-	url = https://github.com/facebookresearch/fairseq

src/third_party/MuQ/LICENSE

@@ -1,21 +0,0 @@
-MIT License
-
-Copyright (c) Tencent.
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.

src/third_party/MuQ/LICENSE_weights

@@ -1,399 +0,0 @@
-Attribution-NonCommercial 4.0 International
-
-=======================================================================
-
-Creative Commons Corporation ("Creative Commons") is not a law firm and
-does not provide legal services or legal advice. Distribution of
-Creative Commons public licenses does not create a lawyer-client or
-other relationship. Creative Commons makes its licenses and related
-information available on an "as-is" basis. Creative Commons gives no
-warranties regarding its licenses, any material licensed under their
-terms and conditions, or any related information. Creative Commons
-disclaims all liability for damages resulting from their use to the
-fullest extent possible.
-
-Using Creative Commons Public Licenses
-
-Creative Commons public licenses provide a standard set of terms and
-conditions that creators and other rights holders may use to share
-original works of authorship and other material subject to copyright
-and certain other rights specified in the public license below. The
-following considerations are for informational purposes only, are not
-exhaustive, and do not form part of our licenses.
-
-     Considerations for licensors: Our public licenses are
-     intended for use by those authorized to give the public
-     permission to use material in ways otherwise restricted by
-     copyright and certain other rights. Our licenses are
-     irrevocable. Licensors should read and understand the terms
-     and conditions of the license they choose before applying it.
-     Licensors should also secure all rights necessary before
-     applying our licenses so that the public can reuse the
-     material as expected. Licensors should clearly mark any
-     material not subject to the license. This includes other CC-
-     licensed material, or material used under an exception or
-     limitation to copyright. More considerations for licensors:
-	wiki.creativecommons.org/Considerations_for_licensors
-
-     Considerations for the public: By using one of our public
-     licenses, a licensor grants the public permission to use the
-     licensed material under specified terms and conditions. If
-     the licensor's permission is not necessary for any reason--for
-     example, because of any applicable exception or limitation to
-     copyright--then that use is not regulated by the license. Our
-     licenses grant only permissions under copyright and certain
-     other rights that a licensor has authority to grant. Use of
-     the licensed material may still be restricted for other
-     reasons, including because others have copyright or other
-     rights in the material. A licensor may make special requests,
-     such as asking that all changes be marked or described.
-     Although not required by our licenses, you are encouraged to
-     respect those requests where reasonable. More_considerations
-     for the public: 
-	wiki.creativecommons.org/Considerations_for_licensees
-
-=======================================================================
-
-Creative Commons Attribution-NonCommercial 4.0 International Public
-License
-
-By exercising the Licensed Rights (defined below), You accept and agree
-to be bound by the terms and conditions of this Creative Commons
-Attribution-NonCommercial 4.0 International Public License ("Public
-License"). To the extent this Public License may be interpreted as a
-contract, You are granted the Licensed Rights in consideration of Your
-acceptance of these terms and conditions, and the Licensor grants You
-such rights in consideration of benefits the Licensor receives from
-making the Licensed Material available under these terms and
-conditions.
-
-Section 1 -- Definitions.
-
-  a. Adapted Material means material subject to Copyright and Similar
-     Rights that is derived from or based upon the Licensed Material
-     and in which the Licensed Material is translated, altered,
-     arranged, transformed, or otherwise modified in a manner requiring
-     permission under the Copyright and Similar Rights held by the
-     Licensor. For purposes of this Public License, where the Licensed
-     Material is a musical work, performance, or sound recording,
-     Adapted Material is always produced where the Licensed Material is
-     synched in timed relation with a moving image.
-
-  b. Adapter's License means the license You apply to Your Copyright
-     and Similar Rights in Your contributions to Adapted Material in
-     accordance with the terms and conditions of this Public License.
-
-  c. Copyright and Similar Rights means copyright and/or similar rights
-     closely related to copyright including, without limitation,
-     performance, broadcast, sound recording, and Sui Generis Database
-     Rights, without regard to how the rights are labeled or
-     categorized. For purposes of this Public License, the rights
-     specified in Section 2(b)(1)-(2) are not Copyright and Similar
-     Rights.
-  d. Effective Technological Measures means those measures that, in the
-     absence of proper authority, may not be circumvented under laws
-     fulfilling obligations under Article 11 of the WIPO Copyright
-     Treaty adopted on December 20, 1996, and/or similar international
-     agreements.
-
-  e. Exceptions and Limitations means fair use, fair dealing, and/or
-     any other exception or limitation to Copyright and Similar Rights
-     that applies to Your use of the Licensed Material.
-
-  f. Licensed Material means the artistic or literary work, database,
-     or other material to which the Licensor applied this Public
-     License.
-
-  g. Licensed Rights means the rights granted to You subject to the
-     terms and conditions of this Public License, which are limited to
-     all Copyright and Similar Rights that apply to Your use of the
-     Licensed Material and that the Licensor has authority to license.
-
-  h. Licensor means the individual(s) or entity(ies) granting rights
-     under this Public License.
-
-  i. NonCommercial means not primarily intended for or directed towards
-     commercial advantage or monetary compensation. For purposes of
-     this Public License, the exchange of the Licensed Material for
-     other material subject to Copyright and Similar Rights by digital
-     file-sharing or similar means is NonCommercial provided there is
-     no payment of monetary compensation in connection with the
-     exchange.
-
-  j. Share means to provide material to the public by any means or
-     process that requires permission under the Licensed Rights, such
-     as reproduction, public display, public performance, distribution,
-     dissemination, communication, or importation, and to make material
-     available to the public including in ways that members of the
-     public may access the material from a place and at a time
-     individually chosen by them.
-
-  k. Sui Generis Database Rights means rights other than copyright
-     resulting from Directive 96/9/EC of the European Parliament and of
-     the Council of 11 March 1996 on the legal protection of databases,
-     as amended and/or succeeded, as well as other essentially
-     equivalent rights anywhere in the world.
-
-  l. You means the individual or entity exercising the Licensed Rights
-     under this Public License. Your has a corresponding meaning.
-
-Section 2 -- Scope.
-
-  a. License grant.
-
-       1. Subject to the terms and conditions of this Public License,
-          the Licensor hereby grants You a worldwide, royalty-free,
-          non-sublicensable, non-exclusive, irrevocable license to
-          exercise the Licensed Rights in the Licensed Material to:
-
-            a. reproduce and Share the Licensed Material, in whole or
-               in part, for NonCommercial purposes only; and
-
-            b. produce, reproduce, and Share Adapted Material for
-               NonCommercial purposes only.
-
-       2. Exceptions and Limitations. For the avoidance of doubt, where
-          Exceptions and Limitations apply to Your use, this Public
-          License does not apply, and You do not need to comply with
-          its terms and conditions.
-
-       3. Term. The term of this Public License is specified in Section
-          6(a).
-
-       4. Media and formats; technical modifications allowed. The
-          Licensor authorizes You to exercise the Licensed Rights in
-          all media and formats whether now known or hereafter created,
-          and to make technical modifications necessary to do so. The
-          Licensor waives and/or agrees not to assert any right or
-          authority to forbid You from making technical modifications
-          necessary to exercise the Licensed Rights, including
-          technical modifications necessary to circumvent Effective
-          Technological Measures. For purposes of this Public License,
-          simply making modifications authorized by this Section 2(a)
-          (4) never produces Adapted Material.
-
-       5. Downstream recipients.
-
-            a. Offer from the Licensor -- Licensed Material. Every
-               recipient of the Licensed Material automatically
-               receives an offer from the Licensor to exercise the
-               Licensed Rights under the terms and conditions of this
-               Public License.
-
-            b. No downstream restrictions. You may not offer or impose
-               any additional or different terms or conditions on, or
-               apply any Effective Technological Measures to, the
-               Licensed Material if doing so restricts exercise of the
-               Licensed Rights by any recipient of the Licensed
-               Material.
-
-       6. No endorsement. Nothing in this Public License constitutes or
-          may be construed as permission to assert or imply that You
-          are, or that Your use of the Licensed Material is, connected
-          with, or sponsored, endorsed, or granted official status by,
-          the Licensor or others designated to receive attribution as
-          provided in Section 3(a)(1)(A)(i).
-
-  b. Other rights.
-
-       1. Moral rights, such as the right of integrity, are not
-          licensed under this Public License, nor are publicity,
-          privacy, and/or other similar personality rights; however, to
-          the extent possible, the Licensor waives and/or agrees not to
-          assert any such rights held by the Licensor to the limited
-          extent necessary to allow You to exercise the Licensed
-          Rights, but not otherwise.
-
-       2. Patent and trademark rights are not licensed under this
-          Public License.
-
-       3. To the extent possible, the Licensor waives any right to
-          collect royalties from You for the exercise of the Licensed
-          Rights, whether directly or through a collecting society
-          under any voluntary or waivable statutory or compulsory
-          licensing scheme. In all other cases the Licensor expressly
-          reserves any right to collect such royalties, including when
-          the Licensed Material is used other than for NonCommercial
-          purposes.
-
-Section 3 -- License Conditions.
-
-Your exercise of the Licensed Rights is expressly made subject to the
-following conditions.
-
-  a. Attribution.
-
-       1. If You Share the Licensed Material (including in modified
-          form), You must:
-
-            a. retain the following if it is supplied by the Licensor
-               with the Licensed Material:
-
-                 i. identification of the creator(s) of the Licensed
-                    Material and any others designated to receive
-                    attribution, in any reasonable manner requested by
-                    the Licensor (including by pseudonym if
-                    designated);
-
-                ii. a copyright notice;
-
-               iii. a notice that refers to this Public License;
-
-                iv. a notice that refers to the disclaimer of
-                    warranties;
-
-                 v. a URI or hyperlink to the Licensed Material to the
-                    extent reasonably practicable;
-
-            b. indicate if You modified the Licensed Material and
-               retain an indication of any previous modifications; and
-
-            c. indicate the Licensed Material is licensed under this
-               Public License, and include the text of, or the URI or
-               hyperlink to, this Public License.
-
-       2. You may satisfy the conditions in Section 3(a)(1) in any
-          reasonable manner based on the medium, means, and context in
-          which You Share the Licensed Material. For example, it may be
-          reasonable to satisfy the conditions by providing a URI or
-          hyperlink to a resource that includes the required
-          information.
-
-       3. If requested by the Licensor, You must remove any of the
-          information required by Section 3(a)(1)(A) to the extent
-          reasonably practicable.
-
-       4. If You Share Adapted Material You produce, the Adapter's
-          License You apply must not prevent recipients of the Adapted
-          Material from complying with this Public License.
-
-Section 4 -- Sui Generis Database Rights.
-
-Where the Licensed Rights include Sui Generis Database Rights that
-apply to Your use of the Licensed Material:
-
-  a. for the avoidance of doubt, Section 2(a)(1) grants You the right
-     to extract, reuse, reproduce, and Share all or a substantial
-     portion of the contents of the database for NonCommercial purposes
-     only;
-
-  b. if You include all or a substantial portion of the database
-     contents in a database in which You have Sui Generis Database
-     Rights, then the database in which You have Sui Generis Database
-     Rights (but not its individual contents) is Adapted Material; and
-
-  c. You must comply with the conditions in Section 3(a) if You Share
-     all or a substantial portion of the contents of the database.
-
-For the avoidance of doubt, this Section 4 supplements and does not
-replace Your obligations under this Public License where the Licensed
-Rights include other Copyright and Similar Rights.
-
-Section 5 -- Disclaimer of Warranties and Limitation of Liability.
-
-  a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE
-     EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS
-     AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF
-     ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,
-     IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,
-     WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR
-     PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,
-     ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT
-     KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT
-     ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.
-
-  b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE
-     TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,
-     NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,
-     INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,
-     COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR
-     USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN
-     ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR
-     DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR
-     IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.
-
-  c. The disclaimer of warranties and limitation of liability provided
-     above shall be interpreted in a manner that, to the extent
-     possible, most closely approximates an absolute disclaimer and
-     waiver of all liability.
-
-Section 6 -- Term and Termination.
-
-  a. This Public License applies for the term of the Copyright and
-     Similar Rights licensed here. However, if You fail to comply with
-     this Public License, then Your rights under this Public License
-     terminate automatically.
-
-  b. Where Your right to use the Licensed Material has terminated under
-     Section 6(a), it reinstates:
-
-       1. automatically as of the date the violation is cured, provided
-          it is cured within 30 days of Your discovery of the
-          violation; or
-
-       2. upon express reinstatement by the Licensor.
-
-     For the avoidance of doubt, this Section 6(b) does not affect any
-     right the Licensor may have to seek remedies for Your violations
-     of this Public License.
-
-  c. For the avoidance of doubt, the Licensor may also offer the
-     Licensed Material under separate terms or conditions or stop
-     distributing the Licensed Material at any time; however, doing so
-     will not terminate this Public License.
-
-  d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
-     License.
-
-Section 7 -- Other Terms and Conditions.
-
-  a. The Licensor shall not be bound by any additional or different
-     terms or conditions communicated by You unless expressly agreed.
-
-  b. Any arrangements, understandings, or agreements regarding the
-     Licensed Material not stated herein are separate from and
-     independent of the terms and conditions of this Public License.
-
-Section 8 -- Interpretation.
-
-  a. For the avoidance of doubt, this Public License does not, and
-     shall not be interpreted to, reduce, limit, restrict, or impose
-     conditions on any use of the Licensed Material that could lawfully
-     be made without permission under this Public License.
-
-  b. To the extent possible, if any provision of this Public License is
-     deemed unenforceable, it shall be automatically reformed to the
-     minimum extent necessary to make it enforceable. If the provision
-     cannot be reformed, it shall be severed from this Public License
-     without affecting the enforceability of the remaining terms and
-     conditions.
-
-  c. No term or condition of this Public License will be waived and no
-     failure to comply consented to unless expressly agreed to by the
-     Licensor.
-
-  d. Nothing in this Public License constitutes or may be interpreted
-     as a limitation upon, or waiver of, any privileges and immunities
-     that apply to the Licensor or You, including from the legal
-     processes of any jurisdiction or authority.
-
-=======================================================================
-
-Creative Commons is not a party to its public
-licenses. Notwithstanding, Creative Commons may elect to apply one of
-its public licenses to material it publishes and in those instances
-will be considered the “Licensor.” The text of the Creative Commons
-public licenses is dedicated to the public domain under the CC0 Public
-Domain Dedication. Except for the limited purpose of indicating that
-material is shared under a Creative Commons public license or as
-otherwise permitted by the Creative Commons policies published at
-creativecommons.org/policies, Creative Commons does not authorize the
-use of the trademark "Creative Commons" or any other trademark or logo
-of Creative Commons without its prior written consent including,
-without limitation, in connection with any unauthorized modifications
-to any of its public licenses or any other arrangements,
-understandings, or agreements concerning use of licensed material. For
-the avoidance of doubt, this paragraph does not form part of the
-public licenses.
-
-Creative Commons may be contacted at creativecommons.org.

src/third_party/MuQ/README.md

@@ -1,129 +0,0 @@
-# <img src="images/muq-logo.jpeg" alt="" height="24px"> MuQ & MuQ-MuLan
-
-<div>
-  <a href='#'><img alt="Static Badge" src="https://img.shields.io/badge/Python-3.8%2B-blue?logo=python&logoColor=white"></a>
-  <a href='https://arxiv.org/abs/2501.01108'><img alt="Static Badge" src="https://img.shields.io/badge/arXiv-2501.01108-%23b31b1b?logo=arxiv&link=https%3A%2F%2Farxiv.org%2F"></a>
-  <a href='https://huggingface.co/OpenMuQ'><img alt="Static Badge" src="https://img.shields.io/badge/huggingface-OpenMuQ-%23FFD21E?logo=huggingface&link=https%3A%2F%2Fhuggingface.co%2FOpenMuQ"></a>
-  <a href='https://pytorch.org/'><img alt="Static Badge" src="https://img.shields.io/badge/framework-PyTorch-%23EE4C2C?logo=pytorch"></a>
-  <a href='https://pypi.org/project/muq'><img alt="Static Badge" src="https://img.shields.io/badge/pip%20install-muq-green?logo=PyPI&logoColor=white&link=https%3A%2F%2Fpypi.org%2Fproject%2Fmuq"></a>
-</div>
-
-This is the official repository for the paper *"**MuQ**: Self-Supervised **Mu**sic Representation Learning
- with Mel Residual Vector **Q**uantization"*.
-
-In this repo, the following models are released:
-
-- **MuQ**: A large music foundation model pre-trained via Self-Supervised Learning (SSL), achieving SOTA in various MIR tasks.
-- **MuQ-MuLan**: A music-text joint embedding model trained via contrastive learning, supporting both English and Chinese texts.
-
-## Overview
-
-We develop the **MuQ** for music SSL. MuQ applys our proposed Mel-RVQ as quantitative targets and achieves SOTA performance on many music understanding (or MIR) tasks. 
-
-We also construct the **MuQ-MuLan**, a CLIP-like model trained by contrastive learning, which jointly represents music and text into embeddings.
-
-For more details, please refer to our [paper](https://arxiv.org/abs/2501.01108).
-
-<div>
-  <img src="images/radar.jpg" width="45%" alt="Evaluation on MARBLE Benchmark">
-  <img src="images/tagging.jpg" width="45%" alt="Evaluation on Zero-shot Music Tagging">
-</div>
-
-## Usage
-
-To begin with, please use pip to install the official `muq` lib, and ensure that your `python>=3.8`:
-```bash
-pip3 install muq
-```
-
-
-To extract music audio features using **MuQ**, you can refer to the following code:
-```python
-import torch, librosa
-from muq import MuQ
-
-device = 'cuda'
-wav, sr = librosa.load("path/to/music_audio.wav", sr = 24000)
-wavs = torch.tensor(wav).unsqueeze(0).to(device) 
-
-# This will automatically fetch the checkpoint from huggingface
-muq = MuQ.from_pretrained("OpenMuQ/MuQ-large-msd-iter")
-muq = muq.to(device).eval()
-
-with torch.no_grad():
-    output = muq(wavs, output_hidden_states=True)
-
-print('Total number of layers: ', len(output.hidden_states))
-print('Feature shape: ', output.last_hidden_state.shape)
-
-```
-
-Using **MuQ-MuLan** to extract the music and text embeddings and calculate the similarity:
-```python
-import torch, librosa
-from muq import MuQMuLan
-
-# This will automatically fetch checkpoints from huggingface
-device = 'cuda'
-mulan = MuQMuLan.from_pretrained("OpenMuQ/MuQ-MuLan-large")
-mulan = mulan.to(device).eval()
-
-# Extract music embeddings
-wav, sr = librosa.load("path/to/music_audio.wav", sr = 24000)
-wavs = torch.tensor(wav).unsqueeze(0).to(device) 
-with torch.no_grad():
-    audio_embeds = mulan(wavs = wavs) 
-
-# Extract text embeddings (texts can be in English or Chinese)
-texts = ["classical genres, hopeful mood, piano.", "一首适合海边风景的小提琴曲，节奏欢快"]
-with torch.no_grad():
-    text_embeds = mulan(texts = texts)
-
-# Calculate dot product similarity
-sim = mulan.calc_similarity(audio_embeds, text_embeds)
-print(sim)
-```
-
-> Note that both MuQ and MuQ-MuLan strictly require **24 kHz** audio as input.
-> We recommend using **fp32** during MuQ inference to avoid potential NaN issues.
-
-
-## Performance
-
-<img src="images/tab-marble.jpg" width="100%" style="max-width: 800px" alt="Table MARBLE Benchmark">
-<img src="images/tab-mulan.png" width="50%" style="max-width: 400px; margin: 0 25%" alt="Table Mulan Results">
-
-## Model Checkpoints
-
-| Model Name | Parameters | Data | HuggingFace🤗 |
-| ----------- | --- | ---  | ----------- |
-| MuQ    | ~300M  | MSD dataset | [OpenMuQ/MuQ-large-msd-iter](https://huggingface.co/OpenMuQ/MuQ-large-msd-iter)       |
-| MuQ-MuLan  | ~700M | music-text pairs | [OpenMuQ/MuQ-MuLan-large](https://huggingface.co/OpenMuQ/MuQ-MuLan-large)       |
-
-**Note**: Please note that the open-sourced MuQ was trained on the Million Song Dataset. Due to differences in dataset size, the open-sourced model may not achieve the same level of performance as reported in the paper. The training recipes can be found [here](./src/recipes).
-
-## License
-
-The code in this repository is released under the MIT license as found in the [LICENSE](LICENSE) file.
-
-The model weights (MuQ-large-msd-iter, MuQ-MuLan-large) in this repository are released under the CC-BY-NC 4.0 license, as detailed in the [LICENSE_weights](LICENSE_weights) file. 
-
-## Citation
-
-```
-@article{zhu2025muq,
-      title={MuQ: Self-Supervised Music Representation Learning with Mel Residual Vector Quantization}, 
-      author={Haina Zhu and Yizhi Zhou and Hangting Chen and Jianwei Yu and Ziyang Ma and Rongzhi Gu and Yi Luo and Wei Tan and Xie Chen},
-      journal={arXiv preprint arXiv:2501.01108},
-      year={2025}
-}
-```
-
-## Acknowledgement
-
-We borrow many codes from the following repositories:
-- [lucidrains/musiclm-pytorch](https://github.com/lucidrains/musiclm-pytorch)
-- [minzwon/musicfm](https://github.com/minzwon/musicfm)
-
-
-Also, we are especially grateful to the awesome [MARBLE-Benchmark](https://github.com/a43992899/MARBLE-Benchmark).

src/third_party/MuQ/requirements.txt

@@ -1,11 +0,0 @@
-einops
-librosa
-nnAudio
-numpy
-soundfile
-torch
-torchaudio
-tqdm
-transformers
-easydict
-x_clip

src/third_party/MuQ/setup.py

@@ -1,34 +0,0 @@
-from setuptools import setup, find_packages
-
-setup(
-    name='muq',  # Name of the package
-    version='0.1.0',  # Version of the package
-    packages=find_packages(where='src'),  # Automatically discover packages under the 'src' directory
-    package_dir={'': 'src'},  # Specify the root directory for packages as 'src'
-    include_package_data=True,  # Include additional files, such as static files
-    install_requires=[  # List of dependencies
-        "einops",
-        "librosa",
-        "nnAudio", 
-        "numpy", 
-        "soundfile", 
-        "torch", 
-        "torchaudio", 
-        "tqdm", 
-        "transformers", 
-        "easydict",
-        "x_clip", 
-    ],
-    author='Haina Zhu',  # Author name
-    author_email='juhayna@qq.com',  # Author email address
-    description='MuQ: A deep learning model for music and text',  # Short description of the package
-    long_description=open('README.md', encoding='utf-8').read(),  # Long description from the README file
-    long_description_content_type='text/markdown',  # Format of the long description (Markdown)
-    url='https://github.com/tencent-ailab/MuQ',  # Project URL
-    classifiers=[
-        'Programming Language :: Python :: 3',  # Python 3 support
-        'License :: OSI Approved :: MIT License',  # License type
-        'Operating System :: OS Independent',  # Supports all operating systems
-    ],
-    python_requires='>=3.8',  # Supported Python version
-)

src/third_party/MuQ/src/muq/__init__.py

@@ -1,2 +0,0 @@
-from .muq import MuQ, MuQConfig
-from .muq_mulan import MuQMuLan, MuQMuLanConfig

src/third_party/MuQ/src/muq/muq/__init__.py

@@ -1 +0,0 @@
-from .muq import MuQConfig, MuQ

src/third_party/MuQ/src/muq/muq/models/muq_model.py

@@ -1,366 +0,0 @@
-import json
-import random
-import torch
-from torch import nn
-from einops import rearrange
-import os
-from easydict import EasyDict
-
-from ..modules.random_quantizer import RandomProjectionQuantizer
-from ..modules.features import MelSTFT
-from ..modules.conv import Conv2dSubsampling
-
-class MuQModel(nn.Module):
-
-    def __init__(
-        self,
-        num_codebooks=1,
-        codebook_dim=16,
-        codebook_size=4096,
-        features=["melspec_2048"],
-        hop_length=240,
-        n_mels=128,
-        conv_dim=512,
-        encoder_dim=1024,
-        encoder_depth=12,
-        mask_hop=0.4,
-        mask_prob=0.6,
-        is_flash=False,
-        stat=dict(),
-        w2v2_config=dict(),
-        use_rvq_target=False,
-        use_vq_target=False,
-        use_encodec_target=False,
-        rvq_ckpt_path=None,
-        recon_loss_ratio=None,
-        label_rate=25,
-        rvq_n_codebooks=8,
-        rvq_multi_layer_num=1,
-    ):
-        super().__init__()
-
-        # global variables
-        self.hop_length = hop_length
-        self.mask_hop = mask_hop
-        self.mask_prob = mask_prob
-        self.num_codebooks = num_codebooks
-        self.codebook_size = codebook_size
-        self.features = features
-        self.recon_loss_ratio = recon_loss_ratio
-        self.n_fold = int(100//label_rate)
-        self.label_rate = label_rate
-
-        # load feature mean / std stats
-        self.stat = stat
-
-        # feature extractor
-        self.preprocessor_melspec_2048 = MelSTFT(
-            n_fft=2048, hop_length=hop_length, is_db=True
-        )
-
-        # random quantizer
-        self.use_rvq_target = use_rvq_target
-        self.use_vq_target = use_vq_target
-        self.use_encodec_target = use_encodec_target
-        
-        seed = 142
-        if self.use_rvq_like_target:
-            if use_rvq_target:
-                from ..modules.rvq import ResidualVectorQuantize
-        
-                inp_dim = 128*self.n_fold
-                self.rvq = ResidualVectorQuantize(
-                    input_dim = inp_dim, 
-                    n_codebooks = rvq_n_codebooks, 
-                    codebook_size = 1024, 
-                    codebook_dim = 16, 
-                    quantizer_dropout = 0.0,
-                    use_multi_layer_num = rvq_multi_layer_num,
-                    )
-            elif use_vq_target:
-                from ..modules.rvq import VectorQuantize
-                
-                self.rvq = VectorQuantize(
-                    input_dim = 128*self.n_fold,
-                    codebook_size = 1024,
-                    codebook_dim = 8,
-                    stale_tolerance = 1000,
-                    mfcc_clustering = False
-                )
-            elif use_encodec_target:
-                from encodec import EncodecModel
-                self.rvq = EncodecModel.encodec_model_24khz()
-                self.rvq.set_target_bandwidth(6.0)
-                for param in self.rvq.parameters():
-                    param.requires_grad = False
-                
-            if rvq_ckpt_path is not None and os.path.exists(rvq_ckpt_path):
-                state_dict = torch.load(rvq_ckpt_path, map_location="cpu")
-                self.rvq.load_state_dict(state_dict)
-            else:
-                pass
-                # print(f'Checkpoint for rvq `{rvq_ckpt_path}` not found. Using random initialization.')
-        else:
-            for feature in self.features:
-                for i in range(num_codebooks):
-                    setattr(
-                        self,
-                        f"quantizer_{feature}", # _{i}
-                        RandomProjectionQuantizer(
-                            n_mels * self.n_fold, codebook_dim, codebook_size, seed=seed + i
-                        ),
-                    )
-
-        # two residual convolution layers + one projection layer
-        strides_factory = {
-            4: [2, 2],
-            2: [2, 1]
-        }
-        self.conv = Conv2dSubsampling(
-            1, conv_dim, encoder_dim, strides=strides_factory.get(self.n_fold), n_bands=n_mels
-        )
-
-        # Conformer
-        if is_flash:
-            from modules.flash_conformer import (
-                Wav2Vec2ConformerEncoder,
-                Wav2Vec2ConformerConfig,
-            )
-        else:
-            from transformers.models.wav2vec2_conformer.modeling_wav2vec2_conformer import (
-                Wav2Vec2ConformerEncoder,
-                Wav2Vec2ConformerConfig,
-            )
-        config = EasyDict(w2v2_config)
-        config.num_hidden_layers = encoder_depth
-        config.hidden_size = encoder_dim
-
-        self.conformer = Wav2Vec2ConformerEncoder(config)
-
-        self.linear = nn.Linear(encoder_dim, codebook_size) # projection layer
-
-        # reconstruct melspec
-        if self.recon_loss_ratio is not None and self.recon_loss_ratio > 0:
-            self.recon_proj = nn.Linear(encoder_dim, n_mels * self.n_fold)
-            self.recon_loss = nn.MSELoss()
-
-        # loss function
-        self.loss = nn.CrossEntropyLoss()
-
-        # cls token (used for sequence classification)
-        random.seed(seed)
-        self.cls_token = nn.Parameter(torch.randn(encoder_dim))
-
-
-    @property
-    def use_rvq_like_target(self):
-        return self.use_rvq_target or self.use_vq_target or self.use_encodec_target
-
-    def masking(self, x, attention_mask=None):
-        """random masking of 400ms with given probability"""
-        mx = x.clone()
-        b, t = mx.shape
-        len_masking_raw = int(24000 * self.mask_hop) 
-        len_masking_token = int(24000 / self.hop_length / 2 / 2 * self.mask_hop) 
-
-        # get random mask indices
-        start_indices = torch.rand(b, t // len_masking_raw) < self.mask_prob
-        time_domain_masked_indices = torch.nonzero(
-            start_indices.repeat_interleave(len_masking_raw, dim=1) 
-        )
-        token_domain_masked_indices = torch.nonzero(
-            start_indices.repeat_interleave(len_masking_token, dim=1)
-        )
-
-        # mask with random values
-        masking_noise = (
-            torch.randn(time_domain_masked_indices.shape[0], dtype=x.dtype) * 0.1
-        )  # 0 mean 0.1 std
-        mx[tuple(time_domain_masked_indices.t())] = masking_noise.to(x.device)
-
-        return mx, token_domain_masked_indices
-
-
-    @torch.no_grad()
-    def preprocessing(self, x, features):
-        """extract classic audio features"""
-        # check precision
-        if x.dtype == torch.float16 or x.dtype == torch.bfloat16:
-            precision = 16
-        else:
-            precision = 32
-
-        out = {}
-        for key in features:
-            layer = getattr(self, "preprocessor_%s" % key)
-            layer.to(x.device)
-            dtype = x.dtype
-            out[key] = layer(x.float())[..., :-1]
-            if precision == 16:
-                out[key] = out[key].half()
-            if out[key].dtype != dtype:
-                out[key].to(dtype=dtype)
-        return out
-
-    def encoder(self, x, *, attention_mask=None, is_features_only=False):
-        """2-layer conv + w2v-conformer"""
-        x = self.conv(x)
-        mask_indices = None
-        if attention_mask is None:
-            out = self.conformer(x, output_hidden_states=True)
-        else:
-            attention_mask = attention_mask.bool()
-            skip_n = int(attention_mask.size(-1) / x.size(1))
-            attention_mask = attention_mask[:, ::skip_n]
-            attention_mask = attention_mask[:, :x.size(1)]
-            out = self.conformer(x, attention_mask=attention_mask, output_hidden_states=True)
-        hidden_emb = out["hidden_states"]
-        last_emb = out["last_hidden_state"]
-        logits = self.linear(last_emb)
-        interval = self.codebook_size
-        logits = {
-            key: logits[:, :, i * interval : (i + 1) * interval]
-            for i, key in enumerate(self.features)
-        }
-        return logits, hidden_emb, mask_indices
-
-    @torch.no_grad()
-    def normalize(self, x):
-        """normalize the input audio to have zero mean unit variance"""
-        for key in x.keys():
-            x[key] = (x[key] - self.stat["%s_mean" % key]) / self.stat["%s_std" % key]
-        return x
-
-    @torch.no_grad()
-    def rearrange(self, x):
-        """rearrange the batch to flatten every 4 steps"""
-        for key in x.keys():
-            if key == "chromagram":
-                x[key] = rearrange(x[key], "b f t -> b t f")
-            else:
-                x[key] = rearrange(x[key], "b f (t s) -> b t (s f)", s=self.n_fold)
-        return x
-
-    def get_rvq_codes(self, inp, raw_wav):
-        if self.use_rvq_target:
-            quantized_prompt_embeds, codes, _, commitment_loss, codebook_loss, rvq_usage = self.rvq(inp)
-            return codes
-        if self.use_vq_target:
-            quantized_prompt_embeds, commitment_loss, codebook_loss, codes, _ = self.rvq(inp)
-            return codes.unsqueeze(1)
-        if self.use_encodec_target:
-            encoded_frames = self.rvq.encode(raw_wav.unsqueeze(1)) #list, B,[ 8,T ]
-            codes = torch.cat([encoded[0].detach() for encoded in encoded_frames], dim=-1)
-            if self.label_rate == 25:
-                codes = codes[:, :, ::3]
-            return codes
-
-    @torch.no_grad()
-    def tokenize(self, x, raw_wav):
-        out = {}
-        for key in x.keys():
-            if self.use_rvq_like_target:
-                self.rvq.eval()
-                inp = x[key].permute((0, 2, 1))
-                codes = self.get_rvq_codes(inp, raw_wav)
-                out[key] = torch.cat([codes[:, idx, ...] for idx in range(int(self.codebook_size//1024))], dim=-1)
-            else:
-                layer = getattr(self, "quantizer_%s" % key)
-                out[key] = layer(x[key])
-        return out
-
-    def get_targets(self, x, label=None):
-        if self.use_encodec_target:
-            raw_x = x.clone()
-        else:
-            raw_x = None
-        x = self.preprocessing(x, features=self.features)
-        x = self.normalize(x)
-        x = self.rearrange(x) 
-        melspec = x['melspec_2048']
-        if label is None:
-            # Use labels from Mel-RVQ
-            target_tokens = self.tokenize(x, raw_x) 
-        else:
-            # Use labels pre-extracted for iteration training
-            target_tokens = {'melspec_2048': rearrange(label, "b n s -> b (n s)").long()}
-        return target_tokens, melspec
-
-    def get_predictions(self, x, *, mask=None, attention_mask=None, return_new_mask=False, is_features_only=False):
-        # preprocessing
-        x = self.preprocessing(x, features=["melspec_2048"])
-        x = self.normalize(x)
-
-        # encoding
-        logits, hidden_emb, new_mask = self.encoder(x["melspec_2048"], attention_mask=attention_mask, is_features_only=is_features_only)
-
-        if return_new_mask:
-            return logits, hidden_emb, mask if new_mask is None else new_mask
-        else:
-            return logits, hidden_emb
-
-    def get_latent(self, x, layer_ix=12):
-        _, hidden_states = self.get_predictions(x)
-        emb = hidden_states[layer_ix]
-        return emb
-
-    def compute_nce(self, x, pos, negs):
-        neg_is_pos = (pos == negs).all(-1)
-        pos = pos.unsqueeze(0)
-        targets = torch.cat([pos, negs], dim=0)
-
-        logits = torch.cosine_similarity(x.float(), targets.float(), dim=-1).type_as(x)
-        logits /= 0.1
-        if neg_is_pos.any():
-            logits[1:][neg_is_pos] = float("-inf")
-        logits = logits.transpose(0, 1)
-        return logits
-    
-    def get_loss(self, logits, target_tokens, masked_indices):
-        losses = {}
-        accuracies = {}
-        for key in logits.keys():
-            if not self.use_rvq_like_target:
-                masked_logits = logits[key][tuple(masked_indices.t())]
-                masked_tokens = target_tokens[key][tuple(masked_indices.t())]
-            else:
-                Batch, SeqLen, N_Codebook_x_CodebookSize = logits[key].shape
-                Batch, N_Codebook_x_SeqLen = target_tokens[key].shape 
-                N_Codebook = int(N_Codebook_x_SeqLen // SeqLen)
-                target_tokens[key] = rearrange(target_tokens[key], "b (n s) -> b s n", n=N_Codebook) # Batch, SeqLen=750, N_Codebook=4
-                masked_logits = logits[key][tuple(masked_indices.t())] 
-                masked_tokens = target_tokens[key][tuple(masked_indices.t())] 
-                masked_logits = rearrange(masked_logits, "b (n c) -> (b n) c", n=N_Codebook)
-                masked_tokens = rearrange(masked_tokens, "b n -> (b n)", n=N_Codebook) 
-
-            losses[key] = self.loss(masked_logits, masked_tokens)
-            accuracies[key] = (
-                torch.sum(masked_logits.argmax(-1) == masked_tokens)
-                / masked_tokens.numel()
-            )
-        return losses, accuracies
-
-    def get_recon_loss(self, last_hidden_emb, melspec, masked_indices):
-        pred_melspec = self.recon_proj(last_hidden_emb[tuple(masked_indices.t())])
-        target_melspec = melspec[tuple(masked_indices.t())]
-        recon_loss = self.recon_loss(pred_melspec, target_melspec)
-        return recon_loss
-
-    def forward(self, x, attention_mask=None, label=None):
-        dtype = x.dtype
-        # get target feature tokens
-        target_tokens, melspec = self.get_targets(x, label=label) 
-
-        # masking
-        x, masked_indices = self.masking(x, attention_mask=attention_mask) 
-
-        # forward
-        logits, hidden_emb, masked_indices = self.get_predictions(x, mask=masked_indices, attention_mask=attention_mask, return_new_mask=True) 
-
-        # get loss
-        losses, accuracies = self.get_loss(logits, target_tokens, masked_indices)
-
-        if self.recon_loss_ratio:
-            losses["recon_loss"] = self.get_recon_loss(hidden_emb[-1], melspec, masked_indices) * self.recon_loss_ratio
-
-        return logits, hidden_emb, losses, accuracies

src/third_party/MuQ/src/muq/muq/modules/__init__.py

@@ -1,2 +0,0 @@
-
-

src/third_party/MuQ/src/muq/muq/modules/conv.py

@@ -1,77 +0,0 @@
-from torch import nn
-from einops import rearrange
-
-
-class Res2dModule(nn.Module):
-    def __init__(self, idim, odim, stride=(2, 2)):
-        super(Res2dModule, self).__init__()
-        self.conv1 = nn.Conv2d(idim, odim, 3, padding=1, stride=stride)
-        self.bn1 = nn.BatchNorm2d(odim)
-        self.conv2 = nn.Conv2d(odim, odim, 3, padding=1)
-        self.bn2 = nn.BatchNorm2d(odim)
-        self.relu = nn.ReLU()
-
-        # residual
-        self.diff = False
-        if (idim != odim) or (stride[0] > 1):
-            self.conv3 = nn.Conv2d(idim, odim, 3, padding=1, stride=stride)
-            self.bn3 = nn.BatchNorm2d(odim)
-            self.diff = True
-
-    def forward(self, x):
-        out = self.bn2(self.conv2(self.relu(self.bn1(self.conv1(x)))))
-        if self.diff:
-            x = self.bn3(self.conv3(x))
-        out = x + out
-        out = self.relu(out)
-        return out
-
-
-class Conv2dSubsampling(nn.Module):
-    """Convolutional 2D subsampling (to 1/4 length).
-
-    Args:
-        idim (int): Input dimension.
-        hdim (int): Hidden dimension.
-        odim (int): Output dimension.
-        strides (list): Sizes of strides.
-        n_bands (int): Number of frequency bands.
-    """
-
-    def __init__(self, idim, hdim, odim, strides=[2, 2], n_bands=64):
-        """Construct an Conv2dSubsampling object."""
-        super(Conv2dSubsampling, self).__init__()
-
-        self.conv = nn.Sequential(
-            Res2dModule(idim, hdim, (2, strides[0])),
-            Res2dModule(hdim, hdim, (2, strides[1])),
-        )
-        self.linear = nn.Linear(hdim * n_bands // 2 // 2, odim)
-
-    def forward(self, x):
-        """Subsample x.
-
-        Args:
-            x (torch.Tensor): Input tensor (#batch, idim, time).
-
-        Returns:
-            torch.Tensor: Subsampled tensor (#batch, time', odim),
-                where time' = time // 4.
-        """
-
-        if x.dim() == 3:
-            x = x.unsqueeze(1)  # (b, c, f, t)
-        x = self.conv(x)
-        x = rearrange(x, "b c f t -> b t (c f)")
-        x = self.linear(x)
-        return x
-
-if __name__ == '__main__':
-    import torch
-    conv_dim, encoder_dim = 512, 1024
-    conv = Conv2dSubsampling(
-            1, conv_dim, encoder_dim, strides=[2, 1], n_bands=128
-        )
-    inp = torch.randn((1, 128, 3000))
-    out = conv(inp)
-    print(out.shape)

src/third_party/MuQ/src/muq/muq/modules/features.py

@@ -1,37 +0,0 @@
-import torchaudio
-from torch import nn
-import torch
-
-
-class MelSTFT:
-    def __init__(
-        self,
-        sample_rate=24000,
-        n_fft=2048,
-        hop_length=240,
-        n_mels=128,
-        is_db=False,
-    ):
-        super(MelSTFT, self).__init__()
-
-        # spectrogram
-        self.mel_stft = torchaudio.transforms.MelSpectrogram(
-            sample_rate=sample_rate, n_fft=n_fft, hop_length=hop_length, n_mels=n_mels
-        )
-
-        # amplitude to decibel
-        self.is_db = is_db
-        if is_db:
-            self.amplitude_to_db = torchaudio.transforms.AmplitudeToDB()
-
-    def __call__(self, waveform):
-        if self.is_db:
-            return self.amplitude_to_db(self.mel_stft(waveform))
-        else:
-            return self.mel_stft(waveform)
-
-    def to(self, device):
-        self.mel_stft = self.mel_stft.to(device)
-        if self.is_db:
-            self.amplitude_to_db = self.amplitude_to_db.to(device)
-        return self

src/third_party/MuQ/src/muq/muq/modules/flash_conformer.py

@@ -1,2114 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The Fairseq Authors and the HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch Wav2Vec2-Conformer model."""
-
-import math
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-
-import numpy as np
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import CrossEntropyLoss
-from torch.nn import functional as F
-
-from transformers.activations import ACT2FN
-from transformers.deepspeed import is_deepspeed_zero3_enabled
-from transformers.modeling_outputs import (
-    BaseModelOutput,
-    CausalLMOutput,
-    SequenceClassifierOutput,
-    TokenClassifierOutput,
-    Wav2Vec2BaseModelOutput,
-    XVectorOutput,
-)
-from transformers.modeling_utils import PreTrainedModel
-from transformers.utils import (
-    ModelOutput,
-    add_code_sample_docstrings,
-    add_start_docstrings,
-    add_start_docstrings_to_model_forward,
-    logging,
-    replace_return_docstrings,
-)
-from transformers.models.wav2vec2_conformer.configuration_wav2vec2_conformer import Wav2Vec2ConformerConfig
-
-
-logger = logging.get_logger(__name__)
-
-
-_HIDDEN_STATES_START_POSITION = 2
-
-# General docstring
-_CONFIG_FOR_DOC = "Wav2Vec2ConformerConfig"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "facebook/wav2vec2-conformer-rope-large-960h-ft"
-_EXPECTED_OUTPUT_SHAPE = [1, 292, 1024]
-
-# CTC docstring
-_CTC_EXPECTED_OUTPUT = "'MISTER QUILTER IS THE APOSTLE OF THE MIDDLE CLASSES AND WE ARE GLAD TO WELCOME HIS GOSPEL'"
-_CTC_EXPECTED_LOSS = 64.21
-
-
-WAV2VEC2_CONFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = [
-    "facebook/wav2vec2-conformer-rel-pos-large",
-    # See all Wav2Vec2Conformer models at https://huggingface.co/models?filter=wav2vec2-conformer
-]
-
-
-@dataclass
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForPreTrainingOutput with Wav2Vec2->Wav2Vec2Conformer
-class Wav2Vec2ConformerForPreTrainingOutput(ModelOutput):
-    """
-    Output type of [`Wav2Vec2ConformerForPreTraining`], with potential hidden states and attentions.
-
-    Args:
-        loss (*optional*, returned when `sample_negative_indices` are passed, `torch.FloatTensor` of shape `(1,)`):
-            Total loss as the sum of the contrastive loss (L_m) and the diversity loss (L_d) as stated in the [official
-            paper](https://arxiv.org/pdf/2006.11477.pdf) . (classification) loss.
-        projected_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.proj_codevector_dim)`):
-            Hidden-states of the model projected to *config.proj_codevector_dim* that can be used to predict the masked
-            projected quantized states.
-        projected_quantized_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.proj_codevector_dim)`):
-            Quantized extracted feature vectors projected to *config.proj_codevector_dim* representing the positive
-            target vectors for contrastive loss.
-        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
-            Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
-            shape `(batch_size, sequence_length, hidden_size)`.
-
-            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
-            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
-            sequence_length)`.
-
-            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
-            heads.
-        contrastive_loss (*optional*, returned when `sample_negative_indices` are passed, `torch.FloatTensor` of shape `(1,)`):
-            The contrastive loss (L_m) as stated in the [official paper](https://arxiv.org/pdf/2006.11477.pdf) .
-        diversity_loss (*optional*, returned when `sample_negative_indices` are passed, `torch.FloatTensor` of shape `(1,)`):
-            The diversity loss (L_d) as stated in the [official paper](https://arxiv.org/pdf/2006.11477.pdf) .
-    """
-
-    loss: Optional[torch.FloatTensor] = None
-    projected_states: torch.FloatTensor = None
-    projected_quantized_states: torch.FloatTensor = None
-    codevector_perplexity: torch.FloatTensor = None
-    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
-    attentions: Optional[Tuple[torch.FloatTensor]] = None
-    contrastive_loss: Optional[torch.FloatTensor] = None
-    diversity_loss: Optional[torch.FloatTensor] = None
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2._compute_mask_indices
-def _compute_mask_indices(
-    shape: Tuple[int, int],
-    mask_prob: float,
-    mask_length: int,
-    attention_mask: Optional[torch.LongTensor] = None,
-    min_masks: int = 0,
-) -> np.ndarray:
-    """
-    Computes random mask spans for a given shape. Used to implement [SpecAugment: A Simple Data Augmentation Method for
-    ASR](https://arxiv.org/abs/1904.08779). Note that this method is not optimized to run on TPU and should be run on
-    CPU as part of the preprocessing during training.
-
-    Args:
-        shape: The shape for which to compute masks. This should be of a tuple of size 2 where
-               the first element is the batch size and the second element is the length of the axis to span.
-        mask_prob:  The percentage of the whole axis (between 0 and 1) which will be masked. The number of
-                    independently generated mask spans of length `mask_length` is computed by
-                    `mask_prob*shape[1]/mask_length`. Note that due to overlaps, `mask_prob` is an upper bound and the
-                    actual percentage will be smaller.
-        mask_length: size of the mask
-        min_masks: minimum number of masked spans
-        attention_mask: A (right-padded) attention mask which independently shortens the feature axis of
-                        each batch dimension.
-    """
-    batch_size, sequence_length = shape
-
-    if mask_length < 1:
-        raise ValueError("`mask_length` has to be bigger than 0.")
-
-    if mask_length > sequence_length:
-        raise ValueError(
-            f"`mask_length` has to be smaller than `sequence_length`, but got `mask_length`: {mask_length}"
-            f" and `sequence_length`: {sequence_length}`"
-        )
-
-    # epsilon is used for probabilistic rounding
-    epsilon = np.random.rand(1).item()
-
-    def compute_num_masked_span(input_length):
-        """Given input length, compute how many spans should be masked"""
-        num_masked_span = int(mask_prob * input_length / mask_length + epsilon)
-        num_masked_span = max(num_masked_span, min_masks)
-
-        # make sure num masked span <= sequence_length
-        if num_masked_span * mask_length > sequence_length:
-            num_masked_span = sequence_length // mask_length
-
-        # make sure num_masked span is also <= input_length - (mask_length - 1)
-        if input_length - (mask_length - 1) < num_masked_span:
-            num_masked_span = max(input_length - (mask_length - 1), 0)
-
-        return num_masked_span
-
-    # compute number of masked spans in batch
-    input_lengths = (
-        attention_mask.sum(-1).detach().tolist()
-        if attention_mask is not None
-        else [sequence_length for _ in range(batch_size)]
-    )
-
-    # SpecAugment mask to fill
-    spec_aug_mask = np.zeros((batch_size, sequence_length), dtype=bool)
-    spec_aug_mask_idxs = []
-
-    max_num_masked_span = compute_num_masked_span(sequence_length)
-
-    if max_num_masked_span == 0:
-        return spec_aug_mask
-
-    for input_length in input_lengths:
-        # compute num of masked spans for this input
-        num_masked_span = compute_num_masked_span(input_length)
-
-        # get random indices to mask
-        spec_aug_mask_idx = np.random.choice(
-            np.arange(input_length - (mask_length - 1)), num_masked_span, replace=False
-        )
-
-        # pick first sampled index that will serve as a dummy index to pad vector
-        # to ensure same dimension for all batches due to probabilistic rounding
-        # Picking first sample just pads those vectors twice.
-        if len(spec_aug_mask_idx) == 0:
-            # this case can only happen if `input_length` is strictly smaller then
-            # `sequence_length` in which case the last token has to be a padding
-            # token which we can use as a dummy mask id
-            dummy_mask_idx = sequence_length - 1
-        else:
-            dummy_mask_idx = spec_aug_mask_idx[0]
-
-        spec_aug_mask_idx = np.concatenate(
-            [spec_aug_mask_idx, np.ones(max_num_masked_span - num_masked_span, dtype=np.int32) * dummy_mask_idx]
-        )
-        spec_aug_mask_idxs.append(spec_aug_mask_idx)
-
-    spec_aug_mask_idxs = np.array(spec_aug_mask_idxs)
-
-    # expand masked indices to masked spans
-    spec_aug_mask_idxs = np.broadcast_to(
-        spec_aug_mask_idxs[:, :, None], (batch_size, max_num_masked_span, mask_length)
-    )
-    spec_aug_mask_idxs = spec_aug_mask_idxs.reshape(batch_size, max_num_masked_span * mask_length)
-
-    # add offset to the starting indexes so that indexes now create a span
-    offsets = np.arange(mask_length)[None, None, :]
-    offsets = np.broadcast_to(offsets, (batch_size, max_num_masked_span, mask_length)).reshape(
-        batch_size, max_num_masked_span * mask_length
-    )
-    spec_aug_mask_idxs = spec_aug_mask_idxs + offsets
-
-    # ensure that we cannot have indices larger than sequence_length
-    if spec_aug_mask_idxs.max() > sequence_length - 1:
-        spec_aug_mask_idxs[spec_aug_mask_idxs > sequence_length - 1] = sequence_length - 1
-
-    # scatter indices to mask
-    np.put_along_axis(spec_aug_mask, spec_aug_mask_idxs, 1, -1)
-
-    return spec_aug_mask
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2._sample_negative_indices
-def _sample_negative_indices(
-    features_shape: Tuple, num_negatives: int, mask_time_indices: Optional[np.ndarray] = None
-):
-    """
-    Sample `num_negatives` vectors from feature vectors.
-    """
-    batch_size, sequence_length = features_shape
-
-    # generate indices of the positive vectors themselves, repeat them `num_negatives` times
-    sequence_length_range = np.arange(sequence_length)
-
-    # get `num_negatives` random vector indices from the same utterance
-    sampled_negative_indices = np.zeros(shape=(batch_size, sequence_length, num_negatives), dtype=np.int32)
-
-    mask_time_indices = (
-        mask_time_indices.astype(bool) if mask_time_indices is not None else np.ones(features_shape, dtype=bool)
-    )
-
-    for batch_idx in range(batch_size):
-        high = mask_time_indices[batch_idx].sum() - 1
-        mapped_masked_indices = sequence_length_range[mask_time_indices[batch_idx]]
-
-        feature_indices = np.broadcast_to(np.arange(high + 1)[:, None], (high + 1, num_negatives))
-        sampled_indices = np.random.randint(0, high, size=(high + 1, num_negatives))
-        # avoid sampling the same positive vector, but keep the distribution uniform
-        sampled_indices[sampled_indices >= feature_indices] += 1
-
-        # remap to actual indices
-        sampled_negative_indices[batch_idx][mask_time_indices[batch_idx]] = mapped_masked_indices[sampled_indices]
-
-        # correct for batch size
-        sampled_negative_indices[batch_idx] += batch_idx * sequence_length
-
-    return sampled_negative_indices
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2NoLayerNormConvLayer with Wav2Vec2->Wav2Vec2Conformer
-class Wav2Vec2ConformerNoLayerNormConvLayer(nn.Module):
-    def __init__(self, config, layer_id=0):
-        super().__init__()
-        self.in_conv_dim = config.conv_dim[layer_id - 1] if layer_id > 0 else 1
-        self.out_conv_dim = config.conv_dim[layer_id]
-
-        self.conv = nn.Conv1d(
-            self.in_conv_dim,
-            self.out_conv_dim,
-            kernel_size=config.conv_kernel[layer_id],
-            stride=config.conv_stride[layer_id],
-            bias=config.conv_bias,
-        )
-        self.activation = ACT2FN[config.feat_extract_activation]
-
-    def forward(self, hidden_states):
-        hidden_states = self.conv(hidden_states)
-        hidden_states = self.activation(hidden_states)
-        return hidden_states
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2LayerNormConvLayer with Wav2Vec2->Wav2Vec2Conformer
-class Wav2Vec2ConformerLayerNormConvLayer(nn.Module):
-    def __init__(self, config, layer_id=0):
-        super().__init__()
-        self.in_conv_dim = config.conv_dim[layer_id - 1] if layer_id > 0 else 1
-        self.out_conv_dim = config.conv_dim[layer_id]
-
-        self.conv = nn.Conv1d(
-            self.in_conv_dim,
-            self.out_conv_dim,
-            kernel_size=config.conv_kernel[layer_id],
-            stride=config.conv_stride[layer_id],
-            bias=config.conv_bias,
-        )
-        self.layer_norm = nn.LayerNorm(self.out_conv_dim, elementwise_affine=True)
-        self.activation = ACT2FN[config.feat_extract_activation]
-
-    def forward(self, hidden_states):
-        hidden_states = self.conv(hidden_states)
-
-        hidden_states = hidden_states.transpose(-2, -1)
-        hidden_states = self.layer_norm(hidden_states)
-        hidden_states = hidden_states.transpose(-2, -1)
-
-        hidden_states = self.activation(hidden_states)
-        return hidden_states
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2GroupNormConvLayer with Wav2Vec2->Wav2Vec2Conformer
-class Wav2Vec2ConformerGroupNormConvLayer(nn.Module):
-    def __init__(self, config, layer_id=0):
-        super().__init__()
-        self.in_conv_dim = config.conv_dim[layer_id - 1] if layer_id > 0 else 1
-        self.out_conv_dim = config.conv_dim[layer_id]
-
-        self.conv = nn.Conv1d(
-            self.in_conv_dim,
-            self.out_conv_dim,
-            kernel_size=config.conv_kernel[layer_id],
-            stride=config.conv_stride[layer_id],
-            bias=config.conv_bias,
-        )
-        self.activation = ACT2FN[config.feat_extract_activation]
-
-        self.layer_norm = nn.GroupNorm(num_groups=self.out_conv_dim, num_channels=self.out_conv_dim, affine=True)
-
-    def forward(self, hidden_states):
-        hidden_states = self.conv(hidden_states)
-        hidden_states = self.layer_norm(hidden_states)
-        hidden_states = self.activation(hidden_states)
-        return hidden_states
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2PositionalConvEmbedding with Wav2Vec2->Wav2Vec2Conformer
-class Wav2Vec2ConformerPositionalConvEmbedding(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.conv = nn.Conv1d(
-            config.hidden_size,
-            config.hidden_size,
-            kernel_size=config.num_conv_pos_embeddings,
-            padding=config.num_conv_pos_embeddings // 2,
-            groups=config.num_conv_pos_embedding_groups,
-        )
-
-        if is_deepspeed_zero3_enabled():
-            import deepspeed
-
-            with deepspeed.zero.GatheredParameters(self.conv.weight, modifier_rank=0):
-                self.conv = nn.utils.weight_norm(self.conv, name="weight", dim=2)
-            deepspeed.zero.register_external_parameter(self, self.conv.weight_v)
-            deepspeed.zero.register_external_parameter(self, self.conv.weight_g)
-        else:
-            self.conv = nn.utils.weight_norm(self.conv, name="weight", dim=2)
-
-        self.padding = Wav2Vec2ConformerSamePadLayer(config.num_conv_pos_embeddings)
-        self.activation = ACT2FN[config.feat_extract_activation]
-
-    def forward(self, hidden_states):
-        hidden_states = hidden_states.transpose(1, 2)
-
-        hidden_states = self.conv(hidden_states)
-        hidden_states = self.padding(hidden_states)
-        hidden_states = self.activation(hidden_states)
-
-        hidden_states = hidden_states.transpose(1, 2)
-        return hidden_states
-
-
-class Wav2Vec2ConformerRotaryPositionalEmbedding(nn.Module):
-    """Rotary positional embedding
-    Reference : https://blog.eleuther.ai/rotary-embeddings/ Paper: https://arxiv.org/pdf/2104.09864.pdf
-    """
-
-    def __init__(self, config):
-        super().__init__()
-        dim = config.hidden_size // config.num_attention_heads
-        base = config.rotary_embedding_base
-
-        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
-        self.register_buffer("inv_freq", inv_freq)
-        self.cached_sequence_length = None
-        self.cached_rotary_positional_embedding = None
-
-    def forward(self, hidden_states):
-        sequence_length = hidden_states.shape[1]
-
-        if sequence_length == self.cached_sequence_length and self.cached_rotary_positional_embedding is not None:
-            return self.cached_rotary_positional_embedding
-
-        self.cached_sequence_length = sequence_length
-        time_stamps = torch.arange(sequence_length).type_as(self.inv_freq)
-        freqs = torch.einsum("i,j->ij", time_stamps, self.inv_freq)
-        embeddings = torch.cat((freqs, freqs), dim=-1)
-
-        cos_embeddings = embeddings.cos()[:, None, None, :]
-        sin_embeddings = embeddings.sin()[:, None, None, :]
-        self.cached_rotary_positional_embedding = torch.stack([cos_embeddings, sin_embeddings])
-        return self.cached_rotary_positional_embedding
-
-
-class Wav2Vec2ConformerRelPositionalEmbedding(nn.Module):
-    """Relative positional encoding module."""
-
-    def __init__(self, config):
-        super().__init__()
-        self.max_len = config.max_source_positions
-        self.d_model = config.hidden_size
-        self.pe = None
-        self.extend_pe(torch.tensor(0.0).expand(1, self.max_len))
-
-    def extend_pe(self, x):
-        # Reset the positional encodings
-        if self.pe is not None:
-            # self.pe contains both positive and negative parts
-            # the length of self.pe is 2 * input_len - 1
-            if self.pe.size(1) >= x.size(1) * 2 - 1:
-                if self.pe.dtype != x.dtype or self.pe.device != x.device:
-                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)
-                return
-        # Suppose `i` is the position of query vector and `j` is the
-        # position of key vector. We use positive relative positions when keys
-        # are to the left (i>j) and negative relative positions otherwise (i<j).
-        pe_positive = torch.zeros(x.size(1), self.d_model)
-        pe_negative = torch.zeros(x.size(1), self.d_model)
-        position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)
-        div_term = torch.exp(
-            torch.arange(0, self.d_model, 2, dtype=torch.float32) * -(math.log(10000.0) / self.d_model)
-        )
-        pe_positive[:, 0::2] = torch.sin(position * div_term)
-        pe_positive[:, 1::2] = torch.cos(position * div_term)
-        pe_negative[:, 0::2] = torch.sin(-1 * position * div_term)
-        pe_negative[:, 1::2] = torch.cos(-1 * position * div_term)
-
-        # Reverse the order of positive indices and concat both positive and
-        # negative indices. This is used to support the shifting trick
-        # as in https://arxiv.org/abs/1901.02860
-        pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)
-        pe_negative = pe_negative[1:].unsqueeze(0)
-        pe = torch.cat([pe_positive, pe_negative], dim=1)
-        self.pe = pe.to(device=x.device, dtype=x.dtype)
-
-    def forward(self, hidden_states: torch.Tensor):
-        self.extend_pe(hidden_states)
-        start_idx = self.pe.size(1) // 2 - hidden_states.size(1) + 1
-        end_idx = self.pe.size(1) // 2 + hidden_states.size(1)
-        relative_position_embeddings = self.pe[:, start_idx:end_idx]
-
-        return relative_position_embeddings
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2SamePadLayer with Wav2Vec2->Wav2Vec2Conformer
-class Wav2Vec2ConformerSamePadLayer(nn.Module):
-    def __init__(self, num_conv_pos_embeddings):
-        super().__init__()
-        self.num_pad_remove = 1 if num_conv_pos_embeddings % 2 == 0 else 0
-
-    def forward(self, hidden_states):
-        if self.num_pad_remove > 0:
-            hidden_states = hidden_states[:, :, : -self.num_pad_remove]
-        return hidden_states
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2FeatureEncoder with Wav2Vec2->Wav2Vec2Conformer
-class Wav2Vec2ConformerFeatureEncoder(nn.Module):
-    """Construct the features from raw audio waveform"""
-
-    def __init__(self, config):
-        super().__init__()
-
-        if config.feat_extract_norm == "group":
-            conv_layers = [Wav2Vec2ConformerGroupNormConvLayer(config, layer_id=0)] + [
-                Wav2Vec2ConformerNoLayerNormConvLayer(config, layer_id=i + 1)
-                for i in range(config.num_feat_extract_layers - 1)
-            ]
-        elif config.feat_extract_norm == "layer":
-            conv_layers = [
-                Wav2Vec2ConformerLayerNormConvLayer(config, layer_id=i) for i in range(config.num_feat_extract_layers)
-            ]
-        else:
-            raise ValueError(
-                f"`config.feat_extract_norm` is {config.feat_extract_norm}, but has to be one of ['group', 'layer']"
-            )
-        self.conv_layers = nn.ModuleList(conv_layers)
-        self.gradient_checkpointing = False
-        self._requires_grad = True
-
-    def _freeze_parameters(self):
-        for param in self.parameters():
-            param.requires_grad = False
-        self._requires_grad = False
-
-    def forward(self, input_values):
-        hidden_states = input_values[:, None]
-
-        # make sure hidden_states require grad for gradient_checkpointing
-        if self._requires_grad and self.training:
-            hidden_states.requires_grad = True
-
-        for conv_layer in self.conv_layers:
-            if self._requires_grad and self.gradient_checkpointing and self.training:
-
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        return module(*inputs)
-
-                    return custom_forward
-
-                hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(conv_layer),
-                    hidden_states,
-                )
-            else:
-                hidden_states = conv_layer(hidden_states)
-
-        return hidden_states
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2FeatureProjection with Wav2Vec2->Wav2Vec2Conformer
-class Wav2Vec2ConformerFeatureProjection(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.layer_norm = nn.LayerNorm(config.conv_dim[-1], eps=config.layer_norm_eps)
-        self.projection = nn.Linear(config.conv_dim[-1], config.hidden_size)
-        self.dropout = nn.Dropout(config.feat_proj_dropout)
-
-    def forward(self, hidden_states):
-        # non-projected hidden states are needed for quantization
-        norm_hidden_states = self.layer_norm(hidden_states)
-        hidden_states = self.projection(norm_hidden_states)
-        hidden_states = self.dropout(hidden_states)
-        return hidden_states, norm_hidden_states
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2FeedForward with Wav2Vec2->Wav2Vec2Conformer
-class Wav2Vec2ConformerFeedForward(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.intermediate_dropout = nn.Dropout(config.activation_dropout)
-
-        self.intermediate_dense = nn.Linear(config.hidden_size, config.intermediate_size)
-        if isinstance(config.hidden_act, str):
-            self.intermediate_act_fn = ACT2FN[config.hidden_act]
-        else:
-            self.intermediate_act_fn = config.hidden_act
-
-        self.output_dense = nn.Linear(config.intermediate_size, config.hidden_size)
-        self.output_dropout = nn.Dropout(config.hidden_dropout)
-
-    def forward(self, hidden_states):
-        hidden_states = self.intermediate_dense(hidden_states)
-        hidden_states = self.intermediate_act_fn(hidden_states)
-        hidden_states = self.intermediate_dropout(hidden_states)
-
-        hidden_states = self.output_dense(hidden_states)
-        hidden_states = self.output_dropout(hidden_states)
-        return hidden_states
-
-
-class Wav2Vec2ConformerConvolutionModule(nn.Module):
-    """Convolution block used in the conformer block"""
-
-    def __init__(self, config):
-        super().__init__()
-        if (config.conv_depthwise_kernel_size - 1) % 2 == 1:
-            raise ValueError("`config.conv_depthwise_kernel_size` should be a odd number for 'SAME' padding")
-        self.layer_norm = nn.LayerNorm(config.hidden_size)
-        self.pointwise_conv1 = torch.nn.Conv1d(
-            config.hidden_size,
-            2 * config.hidden_size,
-            kernel_size=1,
-            stride=1,
-            padding=0,
-            bias=False,
-        )
-        self.glu = torch.nn.GLU(dim=1)
-        self.depthwise_conv = torch.nn.Conv1d(
-            config.hidden_size,
-            config.hidden_size,
-            config.conv_depthwise_kernel_size,
-            stride=1,
-            padding=(config.conv_depthwise_kernel_size - 1) // 2,
-            groups=config.hidden_size,
-            bias=False,
-        )
-        self.batch_norm = torch.nn.BatchNorm1d(config.hidden_size)
-        self.activation = ACT2FN[config.hidden_act]
-        self.pointwise_conv2 = torch.nn.Conv1d(
-            config.hidden_size,
-            config.hidden_size,
-            kernel_size=1,
-            stride=1,
-            padding=0,
-            bias=False,
-        )
-        self.dropout = torch.nn.Dropout(config.conformer_conv_dropout)
-
-    def forward(self, hidden_states):
-        hidden_states = self.layer_norm(hidden_states)
-        # exchange the temporal dimension and the feature dimension
-        hidden_states = hidden_states.transpose(1, 2)
-
-        # GLU mechanism
-        # => (batch, 2*channel, dim)
-        hidden_states = self.pointwise_conv1(hidden_states)
-        # => (batch, channel, dim)
-        hidden_states = self.glu(hidden_states)
-
-        # 1D Depthwise Conv
-        hidden_states = self.depthwise_conv(hidden_states)
-        hidden_states = self.batch_norm(hidden_states)
-        hidden_states = self.activation(hidden_states)
-
-        hidden_states = self.pointwise_conv2(hidden_states)
-        hidden_states = self.dropout(hidden_states)
-        hidden_states = hidden_states.transpose(1, 2)
-        return hidden_states
-
-
-class Wav2Vec2ConformerSelfAttention(nn.Module):
-    """Construct an Wav2Vec2ConformerSelfAttention object.
-    Can be enhanced with rotary or relative position embeddings.
-    """
-
-    def __init__(self, config):
-        super().__init__()
-
-        self.head_size = config.hidden_size // config.num_attention_heads
-        self.num_heads = config.num_attention_heads
-        self.position_embeddings_type = config.position_embeddings_type
-
-        self.linear_q = nn.Linear(config.hidden_size, config.hidden_size)
-        self.linear_k = nn.Linear(config.hidden_size, config.hidden_size)
-        self.linear_v = nn.Linear(config.hidden_size, config.hidden_size)
-        self.linear_out = nn.Linear(config.hidden_size, config.hidden_size)
-
-        self.dropout = nn.Dropout(p=config.attention_dropout)
-        self.dropout_p = config.attention_dropout
-
-        self.is_causal = config.is_causal
-
-        if self.position_embeddings_type == "relative":
-            # linear transformation for positional encoding
-            self.linear_pos = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
-            # these two learnable bias are used in matrix c and matrix d
-            # as described in https://arxiv.org/abs/1901.02860 Section 3.3
-            self.pos_bias_u = nn.Parameter(torch.zeros(self.num_heads, self.head_size))
-            self.pos_bias_v = nn.Parameter(torch.zeros(self.num_heads, self.head_size))
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        relative_position_embeddings: Optional[torch.Tensor] = None,
-        output_attentions: bool = False,
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-        # self-attention mechanism
-        batch_size, sequence_length, hidden_size = hidden_states.size()
-
-        # make sure query/key states can be != value states
-        query_key_states = hidden_states
-        value_states = hidden_states
-
-        if self.position_embeddings_type == "rotary":
-            if relative_position_embeddings is None:
-                raise ValueError(
-                    "`relative_position_embeddings` has to be defined when `self.position_embeddings_type == 'rotary'"
-                )
-            query_key_states = self._apply_rotary_embedding(query_key_states, relative_position_embeddings)
-
-        # project query_key_states and value_states
-        query = self.linear_q(query_key_states).view(batch_size, -1, self.num_heads, self.head_size)
-        key = self.linear_k(query_key_states).view(batch_size, -1, self.num_heads, self.head_size)
-        value = self.linear_v(value_states).view(batch_size, -1, self.num_heads, self.head_size)
-
-        # => (batch, head, time1, d_k)
-        query = query.transpose(1, 2)
-        key = key.transpose(1, 2)
-        value = value.transpose(1, 2)
-
-        with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False):
-            hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask, dropout_p=self.dropout_p, is_causal=self.is_causal)
-        probs = None
-
-        # # apply attention_mask if necessary
-        # if attention_mask is not None:
-        #     scores = scores + attention_mask
-
-        # # => (batch, head, time1, time2)
-        # probs = torch.softmax(scores, dim=-1)
-        # probs = self.dropout(probs)
-
-        # # => (batch, head, time1, d_k)
-        # hidden_states = torch.matmul(probs, value)
-
-        # => (batch, time1, hidden_size)
-        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_size)
-        hidden_states = self.linear_out(hidden_states)
-
-        return hidden_states, probs
-
-    def _apply_rotary_embedding(self, hidden_states, relative_position_embeddings):
-        batch_size, sequence_length, hidden_size = hidden_states.size()
-        hidden_states = hidden_states.view(batch_size, sequence_length, self.num_heads, self.head_size)
-
-        cos = relative_position_embeddings[0, :sequence_length, ...]
-        sin = relative_position_embeddings[1, :sequence_length, ...]
-
-        # rotate hidden_states with rotary embeddings
-        hidden_states = hidden_states.transpose(0, 1)
-        rotated_states_begin = hidden_states[..., : self.head_size // 2]
-        rotated_states_end = hidden_states[..., self.head_size // 2 :]
-        rotated_states = torch.cat((-rotated_states_end, rotated_states_begin), dim=rotated_states_begin.ndim - 1)
-        hidden_states = (hidden_states * cos) + (rotated_states * sin)
-        hidden_states = hidden_states.transpose(0, 1)
-
-        hidden_states = hidden_states.view(batch_size, sequence_length, self.num_heads * self.head_size)
-
-        return hidden_states
-
-    def _apply_relative_embeddings(self, query, key, relative_position_embeddings):
-        # 1. project positional embeddings
-        # => (batch, head, 2*time1-1, d_k)
-        proj_relative_position_embeddings = self.linear_pos(relative_position_embeddings)
-        proj_relative_position_embeddings = proj_relative_position_embeddings.view(
-            relative_position_embeddings.size(0), -1, self.num_heads, self.head_size
-        )
-        proj_relative_position_embeddings = proj_relative_position_embeddings.transpose(1, 2)
-        proj_relative_position_embeddings = proj_relative_position_embeddings.transpose(2, 3)
-
-        # 2. Add bias to query
-        # => (batch, head, time1, d_k)
-        query = query.transpose(1, 2)
-        q_with_bias_u = (query + self.pos_bias_u).transpose(1, 2)
-        q_with_bias_v = (query + self.pos_bias_v).transpose(1, 2)
-
-        # 3. attention score: first compute matrix a and matrix c
-        # as described in https://arxiv.org/abs/1901.02860 Section 3.3
-        # => (batch, head, time1, time2)
-        scores_ac = torch.matmul(q_with_bias_u, key.transpose(-2, -1))
-
-        # 4. then compute matrix b and matrix d
-        # => (batch, head, time1, 2*time1-1)
-        scores_bd = torch.matmul(q_with_bias_v, proj_relative_position_embeddings)
-
-        # 5. shift matrix b and matrix d
-        zero_pad = torch.zeros((*scores_bd.size()[:3], 1), device=scores_bd.device, dtype=scores_bd.dtype)
-        scores_bd_padded = torch.cat([zero_pad, scores_bd], dim=-1)
-        scores_bd_padded_shape = scores_bd.size()[:2] + (scores_bd.shape[3] + 1, scores_bd.shape[2])
-        scores_bd_padded = scores_bd_padded.view(*scores_bd_padded_shape)
-        scores_bd = scores_bd_padded[:, :, 1:].view_as(scores_bd)
-        scores_bd = scores_bd[:, :, :, : scores_bd.size(-1) // 2 + 1]
-
-        # 6. sum matrices
-        # => (batch, head, time1, time2)
-        scores = (scores_ac + scores_bd) / math.sqrt(self.head_size)
-
-        return scores
-
-
-class Wav2Vec2ConformerEncoderLayer(nn.Module):
-    """Conformer block based on https://arxiv.org/abs/2005.08100."""
-
-    def __init__(self, config):
-        super().__init__()
-        embed_dim = config.hidden_size
-        dropout = config.attention_dropout
-
-        # Feed-forward 1
-        self.ffn1_layer_norm = nn.LayerNorm(embed_dim)
-        self.ffn1 = Wav2Vec2ConformerFeedForward(config)
-
-        # Self-Attention
-        self.self_attn_layer_norm = nn.LayerNorm(embed_dim)
-        self.self_attn_dropout = torch.nn.Dropout(dropout)
-        self.self_attn = Wav2Vec2ConformerSelfAttention(config)
-
-        # Conformer Convolution
-        self.conv_module = Wav2Vec2ConformerConvolutionModule(config)
-
-        # Feed-forward 2
-        self.ffn2_layer_norm = nn.LayerNorm(embed_dim)
-        self.ffn2 = Wav2Vec2ConformerFeedForward(config)
-        self.final_layer_norm = nn.LayerNorm(embed_dim)
-
-    def forward(
-        self,
-        hidden_states,
-        attention_mask: Optional[torch.Tensor] = None,
-        relative_position_embeddings: Optional[torch.Tensor] = None,
-        output_attentions: bool = False,
-    ):
-        hidden_states = hidden_states
-
-        # 1. Feed-Forward 1 layer
-        residual = hidden_states
-        hidden_states = self.ffn1_layer_norm(hidden_states)
-        hidden_states = self.ffn1(hidden_states)
-        hidden_states = hidden_states * 0.5 + residual
-        residual = hidden_states
-
-        # 2. Self-Attention layer
-        hidden_states = self.self_attn_layer_norm(hidden_states)
-        hidden_states, attn_weigts = self.self_attn(
-            hidden_states=hidden_states,
-            attention_mask=attention_mask,
-            relative_position_embeddings=relative_position_embeddings,
-            output_attentions=output_attentions,
-        )
-        hidden_states = self.self_attn_dropout(hidden_states)
-        hidden_states = hidden_states + residual
-
-        # 3. Convolutional Layer
-        residual = hidden_states
-        hidden_states = self.conv_module(hidden_states)
-        hidden_states = residual + hidden_states
-
-        # 4. Feed-Forward 2 Layer
-        residual = hidden_states
-        hidden_states = self.ffn2_layer_norm(hidden_states)
-        hidden_states = self.ffn2(hidden_states)
-        hidden_states = hidden_states * 0.5 + residual
-        hidden_states = self.final_layer_norm(hidden_states)
-
-        return hidden_states, attn_weigts
-
-
-class Wav2Vec2ConformerEncoder(nn.Module):
-    def __init__(self, config, is_causal=False):
-        super().__init__()
-        config.is_causal = is_causal
-        self.config = config
-
-        if config.position_embeddings_type == "relative":
-            self.embed_positions = Wav2Vec2ConformerRelPositionalEmbedding(config)
-        elif config.position_embeddings_type == "rotary":
-            self.embed_positions = Wav2Vec2ConformerRotaryPositionalEmbedding(config)
-        else:
-            self.embed_positions = None
-
-        self.pos_conv_embed = Wav2Vec2ConformerPositionalConvEmbedding(config)
-        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-        self.dropout = nn.Dropout(config.hidden_dropout)
-        self.layers = nn.ModuleList([Wav2Vec2ConformerEncoderLayer(config) for _ in range(config.num_hidden_layers)])
-        self.gradient_checkpointing = False
-
-    def forward(
-        self,
-        hidden_states,
-        attention_mask=None,
-        output_attentions=False,
-        output_hidden_states=False,
-        return_dict=True,
-    ):
-        all_hidden_states = () if output_hidden_states else None
-        all_self_attentions = () if output_attentions else None
-
-        if attention_mask is not None:
-            # make sure padded tokens output 0
-            hidden_states[~attention_mask] = 0.0
-
-            # extend attention_mask
-            attention_mask = 1.0 - attention_mask[:, None, None, :].to(dtype=hidden_states.dtype)
-            attention_mask = attention_mask * torch.finfo(hidden_states.dtype).min
-            attention_mask = attention_mask.expand(
-                attention_mask.shape[0], 1, attention_mask.shape[-1], attention_mask.shape[-1]
-            )
-
-        hidden_states = self.dropout(hidden_states)
-
-        if self.embed_positions is not None:
-            relative_position_embeddings = self.embed_positions(hidden_states)
-        else:
-            relative_position_embeddings = None
-
-        deepspeed_zero3_is_enabled = is_deepspeed_zero3_enabled()
-
-        for i, layer in enumerate(self.layers):
-            if output_hidden_states:
-                all_hidden_states = all_hidden_states + (hidden_states,)
-
-            # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
-            dropout_probability = np.random.uniform(0, 1)
-
-            skip_the_layer = True if self.training and (dropout_probability < self.config.layerdrop) else False
-            if not skip_the_layer or deepspeed_zero3_is_enabled:
-                # under deepspeed zero3 all gpus must run in sync
-                if self.gradient_checkpointing and self.training:
-                    # create gradient checkpointing function
-                    def create_custom_forward(module):
-                        def custom_forward(*inputs):
-                            return module(*inputs, output_attentions)
-
-                        return custom_forward
-
-                    layer_outputs = torch.utils.checkpoint.checkpoint(
-                        create_custom_forward(layer),
-                        hidden_states,
-                        attention_mask,
-                        relative_position_embeddings,
-                    )
-                else:
-                    layer_outputs = layer(
-                        hidden_states,
-                        attention_mask=attention_mask,
-                        relative_position_embeddings=relative_position_embeddings,
-                        output_attentions=output_attentions,
-                    )
-                hidden_states = layer_outputs[0]
-
-            if skip_the_layer:
-                layer_outputs = (None, None)
-
-            if output_attentions:
-                all_self_attentions = all_self_attentions + (layer_outputs[1],)
-
-        hidden_states = self.layer_norm(hidden_states)
-        if output_hidden_states:
-            all_hidden_states = all_hidden_states + (hidden_states,)
-
-        if not return_dict:
-            return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
-        return BaseModelOutput(
-            last_hidden_state=hidden_states,
-            hidden_states=all_hidden_states,
-            attentions=all_self_attentions,
-        )
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2GumbelVectorQuantizer with Wav2Vec2->Wav2Vec2Conformer
-class Wav2Vec2ConformerGumbelVectorQuantizer(nn.Module):
-    """
-    Vector quantization using gumbel softmax. See `[CATEGORICAL REPARAMETERIZATION WITH
-    GUMBEL-SOFTMAX](https://arxiv.org/pdf/1611.01144.pdf) for more information.
-    """
-
-    def __init__(self, config):
-        super().__init__()
-        self.num_groups = config.num_codevector_groups
-        self.num_vars = config.num_codevectors_per_group
-
-        if config.codevector_dim % self.num_groups != 0:
-            raise ValueError(
-                f"`config.codevector_dim {config.codevector_dim} must be divisible "
-                f"by `config.num_codevector_groups` {self.num_groups} for concatenation"
-            )
-
-        # storage for codebook variables (codewords)
-        self.codevectors = nn.Parameter(
-            torch.FloatTensor(1, self.num_groups * self.num_vars, config.codevector_dim // self.num_groups)
-        )
-        self.weight_proj = nn.Linear(config.conv_dim[-1], self.num_groups * self.num_vars)
-
-        # can be decayed for training
-        self.temperature = 2
-
-    @staticmethod
-    def _compute_perplexity(probs, mask=None):
-        if mask is not None:
-            mask_extended = mask.flatten()[:, None, None].expand(probs.shape)
-            probs = torch.where(mask_extended, probs, torch.zeros_like(probs))
-            marginal_probs = probs.sum(dim=0) / mask.sum()
-        else:
-            marginal_probs = probs.mean(dim=0)
-
-        perplexity = torch.exp(-torch.sum(marginal_probs * torch.log(marginal_probs + 1e-7), dim=-1)).sum()
-        return perplexity
-
-    def forward(self, hidden_states, mask_time_indices=None):
-        batch_size, sequence_length, hidden_size = hidden_states.shape
-
-        # project to codevector dim
-        hidden_states = self.weight_proj(hidden_states)
-        hidden_states = hidden_states.view(batch_size * sequence_length * self.num_groups, -1)
-
-        if self.training:
-            # sample code vector probs via gumbel in differentiateable way
-            codevector_probs = nn.functional.gumbel_softmax(
-                hidden_states.float(), tau=self.temperature, hard=True
-            ).type_as(hidden_states)
-
-            # compute perplexity
-            codevector_soft_dist = torch.softmax(
-                hidden_states.view(batch_size * sequence_length, self.num_groups, -1).float(), dim=-1
-            )
-            perplexity = self._compute_perplexity(codevector_soft_dist, mask_time_indices)
-        else:
-            # take argmax in non-differentiable way
-            # comptute hard codevector distribution (one hot)
-            codevector_idx = hidden_states.argmax(dim=-1)
-            codevector_probs = hidden_states.new_zeros(hidden_states.shape).scatter_(
-                -1, codevector_idx.view(-1, 1), 1.0
-            )
-            codevector_probs = codevector_probs.view(batch_size * sequence_length, self.num_groups, -1)
-
-            perplexity = self._compute_perplexity(codevector_probs, mask_time_indices)
-
-        codevector_probs = codevector_probs.view(batch_size * sequence_length, -1)
-        # use probs to retrieve codevectors
-        codevectors_per_group = codevector_probs.unsqueeze(-1) * self.codevectors
-        codevectors = codevectors_per_group.view(batch_size * sequence_length, self.num_groups, self.num_vars, -1)
-        codevectors = codevectors.sum(-2).view(batch_size, sequence_length, -1)
-
-        return codevectors, perplexity
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2Adapter with Wav2Vec2->Wav2Vec2Conformer
-class Wav2Vec2ConformerAdapter(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-
-        # feature dim might need to be down-projected
-        if config.output_hidden_size != config.hidden_size:
-            self.proj = nn.Linear(config.hidden_size, config.output_hidden_size)
-            self.proj_layer_norm = nn.LayerNorm(config.output_hidden_size)
-        else:
-            self.proj = self.proj_layer_norm = None
-
-        self.layers = nn.ModuleList(Wav2Vec2ConformerAdapterLayer(config) for _ in range(config.num_adapter_layers))
-        self.layerdrop = config.layerdrop
-
-    def forward(self, hidden_states):
-        # down project hidden_states if necessary
-        if self.proj is not None and self.proj_layer_norm is not None:
-            hidden_states = self.proj(hidden_states)
-            hidden_states = self.proj_layer_norm(hidden_states)
-
-        hidden_states = hidden_states.transpose(1, 2)
-
-        for layer in self.layers:
-            layerdrop_prob = np.random.random()
-            if not self.training or (layerdrop_prob > self.layerdrop):
-                hidden_states = layer(hidden_states)
-
-        hidden_states = hidden_states.transpose(1, 2)
-        return hidden_states
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2AdapterLayer with Wav2Vec2->Wav2Vec2Conformer
-class Wav2Vec2ConformerAdapterLayer(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.conv = nn.Conv1d(
-            config.output_hidden_size,
-            2 * config.output_hidden_size,
-            config.adapter_kernel_size,
-            stride=config.adapter_stride,
-            padding=1,
-        )
-
-    def forward(self, hidden_states):
-        hidden_states = self.conv(hidden_states)
-        hidden_states = nn.functional.glu(hidden_states, dim=1)
-
-        return hidden_states
-
-
-class Wav2Vec2ConformerPreTrainedModel(PreTrainedModel):
-    """
-    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
-    models.
-    """
-
-    config_class = Wav2Vec2ConformerConfig
-    base_model_prefix = "wav2vec2_conformer"
-    main_input_name = "input_values"
-    _keys_to_ignore_on_load_missing = [r"position_ids"]
-    supports_gradient_checkpointing = True
-
-    def _init_weights(self, module):
-        """Initialize the weights"""
-        # Wav2Vec2ForPreTraining last 2 linear layers need standard Linear init.
-        if isinstance(module, Wav2Vec2ConformerForPreTraining):
-            module.project_hid.reset_parameters()
-            module.project_q.reset_parameters()
-            module.project_hid._is_hf_initialized = True
-            module.project_q._is_hf_initialized = True
-        # gumbel softmax requires special init
-        elif isinstance(module, Wav2Vec2ConformerGumbelVectorQuantizer):
-            module.weight_proj.weight.data.normal_(mean=0.0, std=1)
-            module.weight_proj.bias.data.zero_()
-            nn.init.uniform_(module.codevectors)
-        elif isinstance(module, Wav2Vec2ConformerSelfAttention):
-            if hasattr(module, "pos_bias_u"):
-                nn.init.xavier_uniform_(module.pos_bias_u)
-            if hasattr(module, "pos_bias_v"):
-                nn.init.xavier_uniform_(module.pos_bias_v)
-        elif isinstance(module, Wav2Vec2ConformerPositionalConvEmbedding):
-            nn.init.normal_(
-                module.conv.weight,
-                mean=0,
-                std=2 * math.sqrt(1 / (module.conv.kernel_size[0] * module.conv.in_channels)),
-            )
-            nn.init.constant_(module.conv.bias, 0)
-        elif isinstance(module, Wav2Vec2ConformerFeatureProjection):
-            k = math.sqrt(1 / module.projection.in_features)
-            nn.init.uniform_(module.projection.weight, a=-k, b=k)
-            nn.init.uniform_(module.projection.bias, a=-k, b=k)
-        elif isinstance(module, nn.Linear):
-            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
-
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):
-            module.bias.data.zero_()
-            module.weight.data.fill_(1.0)
-        elif isinstance(module, nn.Conv1d):
-            nn.init.kaiming_normal_(module.weight)
-
-            if module.bias is not None:
-                k = math.sqrt(module.groups / (module.in_channels * module.kernel_size[0]))
-                nn.init.uniform_(module.bias, a=-k, b=k)
-
-    def _get_feat_extract_output_lengths(
-        self, input_lengths: Union[torch.LongTensor, int], add_adapter: Optional[bool] = None
-    ):
-        """
-        Computes the output length of the convolutional layers
-        """
-
-        add_adapter = self.config.add_adapter if add_adapter is None else add_adapter
-
-        def _conv_out_length(input_length, kernel_size, stride):
-            # 1D convolutional layer output length formula taken
-            # from https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html
-            return torch.div(input_length - kernel_size, stride, rounding_mode="floor") + 1
-
-        for kernel_size, stride in zip(self.config.conv_kernel, self.config.conv_stride):
-            input_lengths = _conv_out_length(input_lengths, kernel_size, stride)
-
-        if add_adapter:
-            for _ in range(self.config.num_adapter_layers):
-                input_lengths = _conv_out_length(input_lengths, 1, self.config.adapter_stride)
-
-        return input_lengths
-
-    def _get_feature_vector_attention_mask(
-        self, feature_vector_length: int, attention_mask: torch.LongTensor, add_adapter=None
-    ):
-        # Effectively attention_mask.sum(-1), but not inplace to be able to run
-        # on inference mode.
-        non_padded_lengths = attention_mask.cumsum(dim=-1)[:, -1]
-
-        output_lengths = self._get_feat_extract_output_lengths(non_padded_lengths, add_adapter=add_adapter)
-        output_lengths = output_lengths.to(torch.long)
-
-        batch_size = attention_mask.shape[0]
-
-        attention_mask = torch.zeros(
-            (batch_size, feature_vector_length), dtype=attention_mask.dtype, device=attention_mask.device
-        )
-        # these two operations makes sure that all values before the output lengths idxs are attended to
-        attention_mask[(torch.arange(attention_mask.shape[0], device=attention_mask.device), output_lengths - 1)] = 1
-        attention_mask = attention_mask.flip([-1]).cumsum(-1).flip([-1]).bool()
-        return attention_mask
-
-    def _set_gradient_checkpointing(self, module, value=False):
-        if isinstance(module, (Wav2Vec2ConformerEncoder, Wav2Vec2ConformerFeatureEncoder)):
-            module.gradient_checkpointing = value
-
-
-WAV2VEC2_CONFORMER_START_DOCSTRING = r"""
-    Wav2Vec2Conformer was proposed in [wav2vec 2.0: A Framework for Self-Supervised Learning of Speech
-    Representations](https://arxiv.org/abs/2006.11477) by Alexei Baevski, Henry Zhou, Abdelrahman Mohamed, Michael
-    Auli.
-
-    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
-    library implements for all its model (such as downloading or saving etc.).
-
-    This model is a PyTorch [nn.Module](https://pytorch.org/docs/stable/nn.html#nn.Module) sub-class. Use it as a
-    regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior.
-
-    Parameters:
-        config ([`Wav2Vec2ConformerConfig`]): Model configuration class with all the parameters of the model.
-            Initializing with a config file does not load the weights associated with the model, only the
-            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-
-WAV2VEC2_CONFORMER_INPUTS_DOCSTRING = r"""
-    Args:
-        input_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
-            Float values of input raw speech waveform. Values can be obtained by loading a `.flac` or `.wav` audio file
-            into an array of type `List[float]` or a `numpy.ndarray`, *e.g.* via the soundfile library (`pip install
-            soundfile`). To prepare the array into `input_values`, the [`AutoProcessor`] should be used for padding and
-            conversion into a tensor of type `torch.FloatTensor`. See [`Wav2Vec2Processor.__call__`] for details.
-        attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
-            Mask to avoid performing convolution and attention on padding token indices. Mask values selected in `[0,
-            1]`:
-
-            - 1 for tokens that are **not masked**,
-            - 0 for tokens that are **masked**.
-
-            [What are attention masks?](../glossary#attention-mask)
-
-            <Tip warning={true}>
-
-            `attention_mask` should only be passed if the corresponding processor has `config.return_attention_mask ==
-            True`. For all models whose processor has `config.return_attention_mask == False`, such as
-            [wav2vec2-conformer-rel-pos-large](https://huggingface.co/facebook/wav2vec2-conformer-rel-pos-large),
-            `attention_mask` should **not** be passed to avoid degraded performance when doing batched inference. For
-            such models `input_values` should simply be padded with 0 and passed without `attention_mask`. Be aware
-            that these models also yield slightly different results depending on whether `input_values` is padded or
-            not.
-
-            </Tip>
-
-        output_attentions (`bool`, *optional*):
-            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
-            tensors for more detail.
-        output_hidden_states (`bool`, *optional*):
-            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
-            more detail.
-        return_dict (`bool`, *optional*):
-            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
-    "The bare Wav2Vec2Conformer Model transformer outputting raw hidden-states without any specific head on top.",
-    WAV2VEC2_CONFORMER_START_DOCSTRING,
-)
-class Wav2Vec2ConformerModel(Wav2Vec2ConformerPreTrainedModel):
-    def __init__(self, config: Wav2Vec2ConformerConfig):
-        super().__init__(config)
-        self.config = config
-        self.feature_extractor = Wav2Vec2ConformerFeatureEncoder(config)
-        self.feature_projection = Wav2Vec2ConformerFeatureProjection(config)
-
-        # model only needs masking vector if mask prob is > 0.0
-        if config.mask_time_prob > 0.0 or config.mask_feature_prob > 0.0:
-            self.masked_spec_embed = nn.Parameter(torch.FloatTensor(config.hidden_size).uniform_())
-
-        self.encoder = Wav2Vec2ConformerEncoder(config)
-
-        self.adapter = Wav2Vec2ConformerAdapter(config) if config.add_adapter else None
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2Model.freeze_feature_encoder
-    def freeze_feature_encoder(self):
-        """
-        Calling this function will disable the gradient computation for the feature encoder so that its parameter will
-        not be updated during training.
-        """
-        self.feature_extractor._freeze_parameters()
-
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2Model._mask_hidden_states
-    def _mask_hidden_states(
-        self,
-        hidden_states: torch.FloatTensor,
-        mask_time_indices: Optional[torch.FloatTensor] = None,
-        attention_mask: Optional[torch.LongTensor] = None,
-    ):
-        """
-        Masks extracted features along time axis and/or along feature axis according to
-        [SpecAugment](https://arxiv.org/abs/1904.08779).
-        """
-
-        # `config.apply_spec_augment` can set masking to False
-        if not getattr(self.config, "apply_spec_augment", True):
-            return hidden_states
-
-        # generate indices & apply SpecAugment along time axis
-        batch_size, sequence_length, hidden_size = hidden_states.size()
-
-        if mask_time_indices is not None:
-            # apply SpecAugment along time axis with given mask_time_indices
-            hidden_states[mask_time_indices] = self.masked_spec_embed.to(hidden_states.dtype)
-        elif self.config.mask_time_prob > 0 and self.training:
-            mask_time_indices = _compute_mask_indices(
-                (batch_size, sequence_length),
-                mask_prob=self.config.mask_time_prob,
-                mask_length=self.config.mask_time_length,
-                attention_mask=attention_mask,
-                min_masks=self.config.mask_time_min_masks,
-            )
-            mask_time_indices = torch.tensor(mask_time_indices, device=hidden_states.device, dtype=torch.bool)
-            hidden_states[mask_time_indices] = self.masked_spec_embed.to(hidden_states.dtype)
-
-        if self.config.mask_feature_prob > 0 and self.training:
-            # generate indices & apply SpecAugment along feature axis
-            mask_feature_indices = _compute_mask_indices(
-                (batch_size, hidden_size),
-                mask_prob=self.config.mask_feature_prob,
-                mask_length=self.config.mask_feature_length,
-                min_masks=self.config.mask_feature_min_masks,
-            )
-            mask_feature_indices = torch.tensor(mask_feature_indices, device=hidden_states.device, dtype=torch.bool)
-            mask_feature_indices = mask_feature_indices[:, None].expand(-1, sequence_length, -1)
-            hidden_states[mask_feature_indices] = 0
-
-        return hidden_states
-
-    @add_start_docstrings_to_model_forward(WAV2VEC2_CONFORMER_INPUTS_DOCSTRING)
-    @add_code_sample_docstrings(
-        checkpoint=_CHECKPOINT_FOR_DOC,
-        output_type=Wav2Vec2BaseModelOutput,
-        config_class=_CONFIG_FOR_DOC,
-        modality="audio",
-        expected_output=_EXPECTED_OUTPUT_SHAPE,
-    )
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2Model.forward with wav2vec2->wav2vec2_conformer
-    def forward(
-        self,
-        input_values: Optional[torch.Tensor],
-        attention_mask: Optional[torch.Tensor] = None,
-        mask_time_indices: Optional[torch.FloatTensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ) -> Union[Tuple, Wav2Vec2BaseModelOutput]:
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        extract_features = self.feature_extractor(input_values)
-        extract_features = extract_features.transpose(1, 2)
-
-        if attention_mask is not None:
-            # compute reduced attention_mask corresponding to feature vectors
-            attention_mask = self._get_feature_vector_attention_mask(
-                extract_features.shape[1], attention_mask, add_adapter=False
-            )
-
-        hidden_states, extract_features = self.feature_projection(extract_features)
-        hidden_states = self._mask_hidden_states(
-            hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask
-        )
-
-        encoder_outputs = self.encoder(
-            hidden_states,
-            attention_mask=attention_mask,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-        hidden_states = encoder_outputs[0]
-
-        if self.adapter is not None:
-            hidden_states = self.adapter(hidden_states)
-
-        if not return_dict:
-            return (hidden_states, extract_features) + encoder_outputs[1:]
-
-        return Wav2Vec2BaseModelOutput(
-            last_hidden_state=hidden_states,
-            extract_features=extract_features,
-            hidden_states=encoder_outputs.hidden_states,
-            attentions=encoder_outputs.attentions,
-        )
-
-
-@add_start_docstrings(
-    """Wav2Vec2Conformer Model with a quantizer and `VQ` head on top.""", WAV2VEC2_CONFORMER_START_DOCSTRING
-)
-class Wav2Vec2ConformerForPreTraining(Wav2Vec2ConformerPreTrainedModel):
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForPreTraining.__init__ with Wav2Vec2->Wav2Vec2Conformer,wav2vec2->wav2vec2_conformer
-    def __init__(self, config: Wav2Vec2ConformerConfig):
-        super().__init__(config)
-        self.wav2vec2_conformer = Wav2Vec2ConformerModel(config)
-        self.dropout_features = nn.Dropout(config.feat_quantizer_dropout)
-
-        self.quantizer = Wav2Vec2ConformerGumbelVectorQuantizer(config)
-
-        self.project_hid = nn.Linear(config.hidden_size, config.proj_codevector_dim)
-        self.project_q = nn.Linear(config.codevector_dim, config.proj_codevector_dim)
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForPreTraining.set_gumbel_temperature
-    def set_gumbel_temperature(self, temperature: int):
-        """
-        Set the Gumbel softmax temperature to a given value. Only necessary for training
-        """
-        self.quantizer.temperature = temperature
-
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForPreTraining.freeze_feature_encoder with wav2vec2->wav2vec2_conformer
-    def freeze_feature_encoder(self):
-        """
-        Calling this function will disable the gradient computation for the feature encoder so that its parameter will
-        not be updated during training.
-        """
-        self.wav2vec2_conformer.feature_extractor._freeze_parameters()
-
-    @staticmethod
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForPreTraining.compute_contrastive_logits
-    def compute_contrastive_logits(
-        target_features: torch.FloatTensor,
-        negative_features: torch.FloatTensor,
-        predicted_features: torch.FloatTensor,
-        temperature: int = 0.1,
-    ):
-        """
-        Compute logits for contrastive loss based using cosine similarity as the distance measure between
-        `[positive_feature, negative_features]` and `[predicted_features]`. Additionally, temperature can be applied.
-        """
-        target_features = torch.cat([target_features, negative_features], dim=0)
-
-        logits = torch.cosine_similarity(predicted_features.float(), target_features.float(), dim=-1).type_as(
-            target_features
-        )
-
-        # apply temperature
-        logits = logits / temperature
-        return logits
-
-    @add_start_docstrings_to_model_forward(WAV2VEC2_CONFORMER_INPUTS_DOCSTRING)
-    @replace_return_docstrings(output_type=Wav2Vec2ConformerForPreTrainingOutput, config_class=_CONFIG_FOR_DOC)
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForPreTraining.forward with Wav2Vec2->Wav2Vec2Conformer,wav2vec2->wav2vec2_conformer,wav2vec2_conformer-base->wav2vec2-conformer-rel-pos-large
-    def forward(
-        self,
-        input_values: Optional[torch.Tensor],
-        attention_mask: Optional[torch.Tensor] = None,
-        mask_time_indices: Optional[torch.BoolTensor] = None,
-        sampled_negative_indices: Optional[torch.BoolTensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ) -> Union[Tuple, Wav2Vec2ConformerForPreTrainingOutput]:
-        r"""
-        mask_time_indices (`torch.BoolTensor` of shape `(batch_size, sequence_length)`, *optional*):
-            Indices to mask extracted features for contrastive loss. When in training mode, model learns to predict
-            masked extracted features in *config.proj_codevector_dim* space.
-        sampled_negative_indices (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_negatives)`, *optional*):
-            Indices indicating which quantized target vectors are used as negative sampled vectors in contrastive loss.
-            Required input for pre-training.
-
-        Returns:
-
-        Example:
-
-        ```python
-        >>> import torch
-        >>> from transformers import AutoFeatureExtractor, Wav2Vec2ConformerForPreTraining
-        >>> from transformers.models.wav2vec2_conformer.modeling_wav2vec2_conformer import (
-        ...     _compute_mask_indices,
-        ...     _sample_negative_indices,
-        ... )
-        >>> from datasets import load_dataset
-
-        >>> feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-conformer-rel-pos-large")
-        >>> model = Wav2Vec2ConformerForPreTraining.from_pretrained("facebook/wav2vec2-conformer-rel-pos-large")
-
-        >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
-        >>> input_values = feature_extractor(ds[0]["audio"]["array"], return_tensors="pt").input_values  # Batch size 1
-
-        >>> # compute masked indices
-        >>> batch_size, raw_sequence_length = input_values.shape
-        >>> sequence_length = model._get_feat_extract_output_lengths(raw_sequence_length).item()
-        >>> mask_time_indices = _compute_mask_indices(
-        ...     shape=(batch_size, sequence_length), mask_prob=0.2, mask_length=2
-        ... )
-        >>> sampled_negative_indices = _sample_negative_indices(
-        ...     features_shape=(batch_size, sequence_length),
-        ...     num_negatives=model.config.num_negatives,
-        ...     mask_time_indices=mask_time_indices,
-        ... )
-        >>> mask_time_indices = torch.tensor(data=mask_time_indices, device=input_values.device, dtype=torch.long)
-        >>> sampled_negative_indices = torch.tensor(
-        ...     data=sampled_negative_indices, device=input_values.device, dtype=torch.long
-        ... )
-
-        >>> with torch.no_grad():
-        ...     outputs = model(input_values, mask_time_indices=mask_time_indices)
-
-        >>> # compute cosine similarity between predicted (=projected_states) and target (=projected_quantized_states)
-        >>> cosine_sim = torch.cosine_similarity(outputs.projected_states, outputs.projected_quantized_states, dim=-1)
-
-        >>> # show that cosine similarity is much higher than random
-        >>> cosine_sim[mask_time_indices.to(torch.bool)].mean() > 0.5
-        tensor(True)
-
-        >>> # for contrastive loss training model should be put into train mode
-        >>> model = model.train()
-        >>> loss = model(
-        ...     input_values, mask_time_indices=mask_time_indices, sampled_negative_indices=sampled_negative_indices
-        ... ).loss
-        ```"""
-
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        if mask_time_indices is not None:
-            mask_time_indices = mask_time_indices.to(torch.bool)
-
-        outputs = self.wav2vec2_conformer(
-            input_values,
-            attention_mask=attention_mask,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            mask_time_indices=mask_time_indices,
-            return_dict=return_dict,
-        )
-
-        # 1. project all transformed features (including masked) to final vq dim
-        transformer_features = self.project_hid(outputs[0])
-
-        # 2. quantize all (unmasked) extracted features and project to final vq dim
-        extract_features = self.dropout_features(outputs[1])
-
-        if attention_mask is not None:
-            # compute reduced attention_mask correponding to feature vectors
-            attention_mask = self._get_feature_vector_attention_mask(
-                extract_features.shape[1], attention_mask, add_adapter=False
-            )
-
-        quantized_features, codevector_perplexity = self.quantizer(
-            extract_features, mask_time_indices=mask_time_indices
-        )
-        quantized_features = self.project_q(quantized_features)
-
-        loss = contrastive_loss = diversity_loss = None
-        if sampled_negative_indices is not None:
-            batch_size, sequence_length, hidden_size = quantized_features.shape
-
-            # for training, we sample negatives
-            # 3. sample K negatives (distractors) quantized states for contrastive loss
-            # if attention_mask is passed, make sure that padded feature vectors cannot be sampled
-            # sample negative quantized vectors BTC => (BxT)C
-            negative_quantized_features = quantized_features.view(-1, hidden_size)[
-                sampled_negative_indices.long().view(-1)
-            ]
-            negative_quantized_features = negative_quantized_features.view(
-                batch_size, sequence_length, -1, hidden_size
-            ).permute(2, 0, 1, 3)
-
-            # 4. compute logits, corresponding to `logs = sim(c_t, [q_t, \sim{q}_t]) / \kappa`
-            # of equation (3) in https://arxiv.org/pdf/2006.11477.pdf
-            logits = self.compute_contrastive_logits(
-                quantized_features[None, :],
-                negative_quantized_features,
-                transformer_features,
-                self.config.contrastive_logits_temperature,
-            )
-
-            # 5. if a negative vector is identical to the positive (i.e. when codebook utilization is low),
-            # its cosine similarity will be masked
-            neg_is_pos = (quantized_features == negative_quantized_features).all(-1)
-
-            if neg_is_pos.any():
-                logits[1:][neg_is_pos] = float("-inf")
-
-            # 6. compute contrastive loss \mathbf{L}_m = cross_entropy(logs) =
-            # -log(exp(sim(c_t, q_t)/\kappa) / \sum_{\sim{q}} exp(sim(c_t, \sim{q})/\kappa))
-            logits = logits.transpose(0, 2).reshape(-1, logits.size(0))
-            target = ((1 - mask_time_indices.long()) * -100).transpose(0, 1).flatten()
-
-            contrastive_loss = nn.functional.cross_entropy(logits.float(), target, reduction="sum")
-            # 7. compute diversity loss: \mathbf{L}_d
-            num_codevectors = self.config.num_codevectors_per_group * self.config.num_codevector_groups
-            diversity_loss = ((num_codevectors - codevector_perplexity) / num_codevectors) * mask_time_indices.sum()
-
-            # 8. \mathbf{L} = \mathbf{L}_m + \alpha * \mathbf{L}_d
-            loss = contrastive_loss + self.config.diversity_loss_weight * diversity_loss
-
-        if not return_dict:
-            if loss is not None:
-                return (loss, transformer_features, quantized_features, codevector_perplexity) + outputs[2:]
-            return (transformer_features, quantized_features, codevector_perplexity) + outputs[2:]
-
-        return Wav2Vec2ConformerForPreTrainingOutput(
-            loss=loss,
-            projected_states=transformer_features,
-            projected_quantized_states=quantized_features,
-            codevector_perplexity=codevector_perplexity,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-            contrastive_loss=contrastive_loss,
-            diversity_loss=diversity_loss,
-        )
-
-
-@add_start_docstrings(
-    """Wav2Vec2Conformer Model with a `language modeling` head on top for Connectionist Temporal Classification (CTC).""",
-    WAV2VEC2_CONFORMER_START_DOCSTRING,
-)
-class Wav2Vec2ConformerForCTC(Wav2Vec2ConformerPreTrainedModel):
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForCTC.__init__ with Wav2Vec2->Wav2Vec2Conformer,wav2vec2->wav2vec2_conformer
-    def __init__(self, config):
-        super().__init__(config)
-
-        self.wav2vec2_conformer = Wav2Vec2ConformerModel(config)
-        self.dropout = nn.Dropout(config.final_dropout)
-
-        if config.vocab_size is None:
-            raise ValueError(
-                f"You are trying to instantiate {self.__class__} with a configuration that "
-                "does not define the vocabulary size of the language model head. Please "
-                "instantiate the model as follows: `Wav2Vec2ConformerForCTC.from_pretrained(..., vocab_size=vocab_size)`. "
-                "or define `vocab_size` of your model's configuration."
-            )
-        output_hidden_size = (
-            config.output_hidden_size if hasattr(config, "add_adapter") and config.add_adapter else config.hidden_size
-        )
-        self.lm_head = nn.Linear(output_hidden_size, config.vocab_size)
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForCTC.freeze_feature_encoder with wav2vec2->wav2vec2_conformer
-    def freeze_feature_encoder(self):
-        """
-        Calling this function will disable the gradient computation for the feature encoder so that its parameter will
-        not be updated during training.
-        """
-        self.wav2vec2_conformer.feature_extractor._freeze_parameters()
-
-    @add_start_docstrings_to_model_forward(WAV2VEC2_CONFORMER_INPUTS_DOCSTRING)
-    @add_code_sample_docstrings(
-        checkpoint=_CHECKPOINT_FOR_DOC,
-        output_type=CausalLMOutput,
-        config_class=_CONFIG_FOR_DOC,
-        expected_output=_CTC_EXPECTED_OUTPUT,
-        expected_loss=_CTC_EXPECTED_LOSS,
-    )
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForCTC.forward with Wav2Vec2->Wav2Vec2Conformer,wav2vec2->wav2vec2_conformer
-    def forward(
-        self,
-        input_values: Optional[torch.Tensor],
-        attention_mask: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        labels: Optional[torch.Tensor] = None,
-    ) -> Union[Tuple, CausalLMOutput]:
-        r"""
-        labels (`torch.LongTensor` of shape `(batch_size, target_length)`, *optional*):
-            Labels for connectionist temporal classification. Note that `target_length` has to be smaller or equal to
-            the sequence length of the output logits. Indices are selected in `[-100, 0, ..., config.vocab_size - 1]`.
-            All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ...,
-            config.vocab_size - 1]`.
-        """
-
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        outputs = self.wav2vec2_conformer(
-            input_values,
-            attention_mask=attention_mask,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-        hidden_states = outputs[0]
-        hidden_states = self.dropout(hidden_states)
-
-        logits = self.lm_head(hidden_states)
-
-        loss = None
-        if labels is not None:
-            if labels.max() >= self.config.vocab_size:
-                raise ValueError(f"Label values must be <= vocab_size: {self.config.vocab_size}")
-
-            # retrieve loss input_lengths from attention_mask
-            attention_mask = (
-                attention_mask if attention_mask is not None else torch.ones_like(input_values, dtype=torch.long)
-            )
-            input_lengths = self._get_feat_extract_output_lengths(attention_mask.sum(-1)).to(torch.long)
-
-            # assuming that padded tokens are filled with -100
-            # when not being attended to
-            labels_mask = labels >= 0
-            target_lengths = labels_mask.sum(-1)
-            flattened_targets = labels.masked_select(labels_mask)
-
-            # ctc_loss doesn't support fp16
-            log_probs = nn.functional.log_softmax(logits, dim=-1, dtype=torch.float32).transpose(0, 1)
-
-            with torch.backends.cudnn.flags(enabled=False):
-                loss = nn.functional.ctc_loss(
-                    log_probs,
-                    flattened_targets,
-                    input_lengths,
-                    target_lengths,
-                    blank=self.config.pad_token_id,
-                    reduction=self.config.ctc_loss_reduction,
-                    zero_infinity=self.config.ctc_zero_infinity,
-                )
-
-        if not return_dict:
-            output = (logits,) + outputs[_HIDDEN_STATES_START_POSITION:]
-            return ((loss,) + output) if loss is not None else output
-
-        return CausalLMOutput(
-            loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions
-        )
-
-
-@add_start_docstrings(
-    """
-    Wav2Vec2Conformer Model with a sequence classification head on top (a linear layer over the pooled output) for
-    tasks like SUPERB Keyword Spotting.
-    """,
-    WAV2VEC2_CONFORMER_START_DOCSTRING,
-)
-class Wav2Vec2ConformerForSequenceClassification(Wav2Vec2ConformerPreTrainedModel):
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForSequenceClassification.__init__ with Wav2Vec2->Wav2Vec2Conformer,wav2vec2->wav2vec2_conformer
-    def __init__(self, config):
-        super().__init__(config)
-
-        if hasattr(config, "add_adapter") and config.add_adapter:
-            raise ValueError(
-                "Sequence classification does not support the use of Wav2Vec2Conformer adapters (config.add_adapter=True)"
-            )
-        self.wav2vec2_conformer = Wav2Vec2ConformerModel(config)
-        num_layers = config.num_hidden_layers + 1  # transformer layers + input embeddings
-        if config.use_weighted_layer_sum:
-            self.layer_weights = nn.Parameter(torch.ones(num_layers) / num_layers)
-        self.projector = nn.Linear(config.hidden_size, config.classifier_proj_size)
-        self.classifier = nn.Linear(config.classifier_proj_size, config.num_labels)
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForSequenceClassification.freeze_feature_encoder with wav2vec2->wav2vec2_conformer
-    def freeze_feature_encoder(self):
-        """
-        Calling this function will disable the gradient computation for the feature encoder so that its parameter will
-        not be updated during training.
-        """
-        self.wav2vec2_conformer.feature_extractor._freeze_parameters()
-
-    def freeze_base_model(self):
-        """
-        Calling this function will disable the gradient computation for the base model so that its parameters will not
-        be updated during training. Only the classification head will be updated.
-        """
-        for param in self.wav2vec2_conformer.parameters():
-            param.requires_grad = False
-
-    @add_start_docstrings_to_model_forward(WAV2VEC2_CONFORMER_INPUTS_DOCSTRING)
-    @add_code_sample_docstrings(
-        checkpoint=_CHECKPOINT_FOR_DOC,
-        output_type=SequenceClassifierOutput,
-        config_class=_CONFIG_FOR_DOC,
-        modality="audio",
-    )
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForSequenceClassification.forward with Wav2Vec2->Wav2Vec2Conformer,wav2vec2->wav2vec2_conformer,WAV_2_VEC_2->WAV2VEC2_CONFORMER
-    def forward(
-        self,
-        input_values: Optional[torch.Tensor],
-        attention_mask: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        labels: Optional[torch.Tensor] = None,
-    ) -> Union[Tuple, SequenceClassifierOutput]:
-        r"""
-        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
-            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
-            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
-            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
-        """
-
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        output_hidden_states = True if self.config.use_weighted_layer_sum else output_hidden_states
-
-        outputs = self.wav2vec2_conformer(
-            input_values,
-            attention_mask=attention_mask,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-        if self.config.use_weighted_layer_sum:
-            hidden_states = outputs[_HIDDEN_STATES_START_POSITION]
-            hidden_states = torch.stack(hidden_states, dim=1)
-            norm_weights = nn.functional.softmax(self.layer_weights, dim=-1)
-            hidden_states = (hidden_states * norm_weights.view(-1, 1, 1)).sum(dim=1)
-        else:
-            hidden_states = outputs[0]
-
-        hidden_states = self.projector(hidden_states)
-        if attention_mask is None:
-            pooled_output = hidden_states.mean(dim=1)
-        else:
-            padding_mask = self._get_feature_vector_attention_mask(hidden_states.shape[1], attention_mask)
-            hidden_states[~padding_mask] = 0.0
-            pooled_output = hidden_states.sum(dim=1) / padding_mask.sum(dim=1).view(-1, 1)
-
-        logits = self.classifier(pooled_output)
-
-        loss = None
-        if labels is not None:
-            loss_fct = CrossEntropyLoss()
-            loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1))
-
-        if not return_dict:
-            output = (logits,) + outputs[_HIDDEN_STATES_START_POSITION:]
-            return ((loss,) + output) if loss is not None else output
-
-        return SequenceClassifierOutput(
-            loss=loss,
-            logits=logits,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )
-
-
-@add_start_docstrings(
-    """
-    Wav2Vec2Conformer Model with a frame classification head on top for tasks like Speaker Diarization.
-    """,
-    WAV2VEC2_CONFORMER_START_DOCSTRING,
-)
-class Wav2Vec2ConformerForAudioFrameClassification(Wav2Vec2ConformerPreTrainedModel):
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForAudioFrameClassification.__init__ with Wav2Vec2->Wav2Vec2Conformer,wav2vec2->wav2vec2_conformer,WAV_2_VEC_2->WAV2VEC2_CONFORMER
-    def __init__(self, config):
-        super().__init__(config)
-
-        if hasattr(config, "add_adapter") and config.add_adapter:
-            raise ValueError(
-                "Audio frame classification does not support the use of Wav2Vec2Conformer adapters (config.add_adapter=True)"
-            )
-        self.wav2vec2_conformer = Wav2Vec2ConformerModel(config)
-        num_layers = config.num_hidden_layers + 1  # transformer layers + input embeddings
-        if config.use_weighted_layer_sum:
-            self.layer_weights = nn.Parameter(torch.ones(num_layers) / num_layers)
-        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-        self.num_labels = config.num_labels
-
-        self.init_weights()
-
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForAudioFrameClassification.freeze_feature_encoder with wav2vec2->wav2vec2_conformer
-    def freeze_feature_encoder(self):
-        """
-        Calling this function will disable the gradient computation for the feature encoder so that its parameter will
-        not be updated during training.
-        """
-        self.wav2vec2_conformer.feature_extractor._freeze_parameters()
-
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForAudioFrameClassification.freeze_base_model with wav2vec2->wav2vec2_conformer
-    def freeze_base_model(self):
-        """
-        Calling this function will disable the gradient computation for the base model so that its parameters will not
-        be updated during training. Only the classification head will be updated.
-        """
-        for param in self.wav2vec2_conformer.parameters():
-            param.requires_grad = False
-
-    @add_start_docstrings_to_model_forward(WAV2VEC2_CONFORMER_INPUTS_DOCSTRING)
-    @add_code_sample_docstrings(
-        checkpoint=_CHECKPOINT_FOR_DOC,
-        output_type=TokenClassifierOutput,
-        config_class=_CONFIG_FOR_DOC,
-        modality="audio",
-    )
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForAudioFrameClassification.forward with wav2vec2->wav2vec2_conformer
-    def forward(
-        self,
-        input_values: Optional[torch.Tensor],
-        attention_mask: Optional[torch.Tensor] = None,
-        labels: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ) -> Union[Tuple, TokenClassifierOutput]:
-        r"""
-        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
-            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
-            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
-            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
-        """
-
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        output_hidden_states = True if self.config.use_weighted_layer_sum else output_hidden_states
-
-        outputs = self.wav2vec2_conformer(
-            input_values,
-            attention_mask=attention_mask,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-        if self.config.use_weighted_layer_sum:
-            hidden_states = outputs[_HIDDEN_STATES_START_POSITION]
-            hidden_states = torch.stack(hidden_states, dim=1)
-            norm_weights = nn.functional.softmax(self.layer_weights, dim=-1)
-            hidden_states = (hidden_states * norm_weights.view(-1, 1, 1)).sum(dim=1)
-        else:
-            hidden_states = outputs[0]
-
-        logits = self.classifier(hidden_states)
-
-        loss = None
-        if labels is not None:
-            loss_fct = CrossEntropyLoss()
-            loss = loss_fct(logits.view(-1, self.num_labels), torch.argmax(labels.view(-1, self.num_labels), axis=1))
-
-        if not return_dict:
-            output = (logits,) + outputs[_HIDDEN_STATES_START_POSITION:]
-            return output
-
-        return TokenClassifierOutput(
-            loss=loss,
-            logits=logits,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.AMSoftmaxLoss
-class AMSoftmaxLoss(nn.Module):
-    def __init__(self, input_dim, num_labels, scale=30.0, margin=0.4):
-        super(AMSoftmaxLoss, self).__init__()
-        self.scale = scale
-        self.margin = margin
-        self.num_labels = num_labels
-        self.weight = nn.Parameter(torch.randn(input_dim, num_labels), requires_grad=True)
-        self.loss = nn.CrossEntropyLoss()
-
-    def forward(self, hidden_states, labels):
-        labels = labels.flatten()
-        weight = nn.functional.normalize(self.weight, dim=0)
-        hidden_states = nn.functional.normalize(hidden_states, dim=1)
-        cos_theta = torch.mm(hidden_states, weight)
-        psi = cos_theta - self.margin
-
-        onehot = nn.functional.one_hot(labels, self.num_labels)
-        logits = self.scale * torch.where(onehot.bool(), psi, cos_theta)
-        loss = self.loss(logits, labels)
-
-        return loss
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.TDNNLayer
-class TDNNLayer(nn.Module):
-    def __init__(self, config, layer_id=0):
-        super().__init__()
-        self.in_conv_dim = config.tdnn_dim[layer_id - 1] if layer_id > 0 else config.tdnn_dim[layer_id]
-        self.out_conv_dim = config.tdnn_dim[layer_id]
-        self.kernel_size = config.tdnn_kernel[layer_id]
-        self.dilation = config.tdnn_dilation[layer_id]
-
-        self.kernel = nn.Linear(self.in_conv_dim * self.kernel_size, self.out_conv_dim)
-        self.activation = nn.ReLU()
-
-    def forward(self, hidden_states):
-        hidden_states = hidden_states.unsqueeze(1)
-        hidden_states = nn.functional.unfold(
-            hidden_states,
-            (self.kernel_size, self.in_conv_dim),
-            stride=(1, self.in_conv_dim),
-            dilation=(self.dilation, 1),
-        )
-        hidden_states = hidden_states.transpose(1, 2)
-        hidden_states = self.kernel(hidden_states)
-
-        hidden_states = self.activation(hidden_states)
-        return hidden_states
-
-
-@add_start_docstrings(
-    """
-    Wav2Vec2Conformer Model with an XVector feature extraction head on top for tasks like Speaker Verification.
-    """,
-    WAV2VEC2_CONFORMER_START_DOCSTRING,
-)
-class Wav2Vec2ConformerForXVector(Wav2Vec2ConformerPreTrainedModel):
-    def __init__(self, config):
-        super().__init__(config)
-
-        self.wav2vec2_conformer = Wav2Vec2ConformerModel(config)
-        num_layers = config.num_hidden_layers + 1  # transformer layers + input embeddings
-        if config.use_weighted_layer_sum:
-            self.layer_weights = nn.Parameter(torch.ones(num_layers) / num_layers)
-        self.projector = nn.Linear(config.hidden_size, config.tdnn_dim[0])
-
-        tdnn_layers = [TDNNLayer(config, i) for i in range(len(config.tdnn_dim))]
-        self.tdnn = nn.ModuleList(tdnn_layers)
-
-        self.feature_extractor = nn.Linear(config.tdnn_dim[-1] * 2, config.xvector_output_dim)
-        self.classifier = nn.Linear(config.xvector_output_dim, config.xvector_output_dim)
-
-        self.objective = AMSoftmaxLoss(config.xvector_output_dim, config.num_labels)
-
-        self.init_weights()
-
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForXVector.freeze_feature_encoder with wav2vec2->wav2vec2_conformer
-    def freeze_feature_encoder(self):
-        """
-        Calling this function will disable the gradient computation for the feature encoder so that its parameter will
-        not be updated during training.
-        """
-        self.wav2vec2_conformer.feature_extractor._freeze_parameters()
-
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForXVector.freeze_base_model with wav2vec2->wav2vec2_conformer
-    def freeze_base_model(self):
-        """
-        Calling this function will disable the gradient computation for the base model so that its parameters will not
-        be updated during training. Only the classification head will be updated.
-        """
-        for param in self.wav2vec2_conformer.parameters():
-            param.requires_grad = False
-
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForXVector._get_tdnn_output_lengths with wav2vec2->wav2vec2_conformer
-    def _get_tdnn_output_lengths(self, input_lengths: Union[torch.LongTensor, int]):
-        """
-        Computes the output length of the TDNN layers
-        """
-
-        def _conv_out_length(input_length, kernel_size, stride):
-            # 1D convolutional layer output length formula taken
-            # from https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html
-            return (input_length - kernel_size) // stride + 1
-
-        for kernel_size in self.config.tdnn_kernel:
-            input_lengths = _conv_out_length(input_lengths, kernel_size, 1)
-
-        return input_lengths
-
-    @add_start_docstrings_to_model_forward(WAV2VEC2_CONFORMER_INPUTS_DOCSTRING)
-    @add_code_sample_docstrings(
-        checkpoint=_CHECKPOINT_FOR_DOC,
-        output_type=XVectorOutput,
-        config_class=_CONFIG_FOR_DOC,
-        modality="audio",
-    )
-    # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForXVector.forward with Wav2Vec2->Wav2Vec2Conformer,wav2vec2->wav2vec2_conformer,WAV_2_VEC_2->WAV2VEC2_CONFORMER
-    def forward(
-        self,
-        input_values: Optional[torch.Tensor],
-        attention_mask: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        labels: Optional[torch.Tensor] = None,
-    ) -> Union[Tuple, XVectorOutput]:
-        r"""
-        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
-            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
-            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
-            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
-        """
-
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        output_hidden_states = True if self.config.use_weighted_layer_sum else output_hidden_states
-
-        outputs = self.wav2vec2_conformer(
-            input_values,
-            attention_mask=attention_mask,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-        if self.config.use_weighted_layer_sum:
-            hidden_states = outputs[_HIDDEN_STATES_START_POSITION]
-            hidden_states = torch.stack(hidden_states, dim=1)
-            norm_weights = nn.functional.softmax(self.layer_weights, dim=-1)
-            hidden_states = (hidden_states * norm_weights.view(-1, 1, 1)).sum(dim=1)
-        else:
-            hidden_states = outputs[0]
-
-        hidden_states = self.projector(hidden_states)
-
-        for tdnn_layer in self.tdnn:
-            hidden_states = tdnn_layer(hidden_states)
-
-        # Statistic Pooling
-        if attention_mask is None:
-            mean_features = hidden_states.mean(dim=1)
-            std_features = hidden_states.std(dim=1)
-        else:
-            feat_extract_output_lengths = self._get_feat_extract_output_lengths(attention_mask.sum(dim=1))
-            tdnn_output_lengths = self._get_tdnn_output_lengths(feat_extract_output_lengths)
-            mean_features = []
-            std_features = []
-            for i, length in enumerate(tdnn_output_lengths):
-                mean_features.append(hidden_states[i, :length].mean(dim=0))
-                std_features.append(hidden_states[i, :length].std(dim=0))
-            mean_features = torch.stack(mean_features)
-            std_features = torch.stack(std_features)
-        statistic_pooling = torch.cat([mean_features, std_features], dim=-1)
-
-        output_embeddings = self.feature_extractor(statistic_pooling)
-        logits = self.classifier(output_embeddings)
-
-        loss = None
-        if labels is not None:
-            loss = self.objective(logits, labels)
-
-        if not return_dict:
-            output = (logits, output_embeddings) + outputs[_HIDDEN_STATES_START_POSITION:]
-            return ((loss,) + output) if loss is not None else output
-
-        return XVectorOutput(
-            loss=loss,
-            logits=logits,
-            embeddings=output_embeddings,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )

src/third_party/MuQ/src/muq/muq/modules/random_quantizer.py

@@ -1,68 +0,0 @@
-import torch
-from torch import nn, einsum
-from einops import rearrange
-
-
-class RandomProjectionQuantizer(nn.Module):
-    """
-    Random projection and codebook lookup module
-
-    Some code is borrowed from:
-     https://github.com/lucidrains/vector-quantize-pytorch/blob/master/vector_quantize_pytorch/random_projection_quantizer.py
-    But I did normalization using pre-computed global mean & variance instead of using layer norm.
-    """
-
-    def __init__(
-        self,
-        input_dim,
-        codebook_dim,
-        codebook_size,
-        seed=142,
-    ):
-        super().__init__()
-
-        # random seed
-        torch.manual_seed(seed)
-
-        # randomly initialized projection
-        random_projection = torch.empty(input_dim, codebook_dim)
-        nn.init.xavier_normal_(random_projection)
-        self.register_buffer("random_projection", random_projection)
-
-        # randomly initialized codebook
-        codebook = torch.empty(codebook_size, codebook_dim)
-        nn.init.normal_(codebook)
-        self.register_buffer("codebook", codebook)
-
-    def codebook_lookup(self, x):
-        # reshape
-        b = x.shape[0]
-        x = rearrange(x, "b n e -> (b n) e")
-
-        # L2 normalization
-        normalized_x = nn.functional.normalize(x, dim=1, p=2)
-        normalized_codebook = nn.functional.normalize(self.codebook, dim=1, p=2)
-
-        # compute distances
-        distances = torch.cdist(normalized_codebook, normalized_x)
-
-        # get nearest
-        nearest_indices = torch.argmin(distances, dim=0)
-
-        # reshape
-        xq = rearrange(nearest_indices, "(b n) -> b n", b=b)
-
-        return xq
-
-    @torch.no_grad()
-    def forward(self, x):
-        # always eval
-        self.eval()
-
-        # random projection [batch, length, input_dim] -> [batch, length, codebook_dim]
-        x = einsum("b n d, d e -> b n e", x, self.random_projection)
-
-        # codebook lookup
-        xq = self.codebook_lookup(x)
-
-        return xq

src/third_party/MuQ/src/muq/muq/modules/rvq.py

@@ -1,314 +0,0 @@
-
-from typing import Union
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange
-try:
-    from torch.nn.utils import weight_norm
-except:
-    try:
-        from torch.nn.utils.parametrizations import weight_norm
-    except:
-        from torch.nn.utils.parametrize import weight_norm
-
-def WNConv1d(*args, **kwargs):
-    return weight_norm(nn.Conv1d(*args, **kwargs))
-
-
-class VectorQuantize(nn.Module):
-    """
-    Implementation of VQ similar to Karpathy's repo:
-    https://github.com/karpathy/deep-vector-quantization
-    Additionally uses following tricks from Improved VQGAN
-    (https://arxiv.org/pdf/2110.04627.pdf):
-        1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space
-            for improved codebook usage
-        2. l2-normalized codes: Converts euclidean distance to cosine similarity which
-            improves training stability
-    """
-
-    def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int, stale_tolerance: int = 1000, mfcc_clustering=False, n_layer=1):
-        super().__init__()
-        self.codebook_size = codebook_size
-        self.codebook_dim = codebook_dim
-        self.mfcc_clustering = mfcc_clustering
-
-        ProjClass = nn.Identity if mfcc_clustering else WNConv1d
-        if n_layer==1:
-            self.in_proj = ProjClass(input_dim, codebook_dim, kernel_size=1)
-            self.out_proj = ProjClass(codebook_dim, input_dim, kernel_size=1)
-        elif n_layer >= 2:
-            ndim_hidden = 128
-            self.in_proj = nn.Sequential(
-                ProjClass(input_dim, ndim_hidden, kernel_size=1),
-                *[nn.Sequential(nn.ReLU(), ProjClass(ndim_hidden, ndim_hidden, kernel_size=1),) for _ in range(n_layer-2)],
-                nn.ReLU(),
-                ProjClass(ndim_hidden, codebook_dim, kernel_size=1)
-            )
-            self.out_proj = nn.Sequential(
-                ProjClass(codebook_dim, ndim_hidden, kernel_size=1),
-                nn.ReLU(),
-                *[nn.Sequential(ProjClass(ndim_hidden, ndim_hidden, kernel_size=1), nn.ReLU()) for _ in range(n_layer-2)],
-                ProjClass(ndim_hidden, input_dim, kernel_size=1),
-            )
-        self.codebook = nn.Embedding(codebook_size, codebook_dim)
-        self.register_buffer("stale_counter", torch.zeros(self.codebook_size,))
-        self.stale_tolerance = stale_tolerance
-
-    def forward(self, z):
-        """Quantized the input tensor using a fixed codebook and returns
-        the corresponding codebook vectors
-
-        Parameters
-        ----------
-        z : Tensor[B x D x T]
-
-        Returns
-        -------
-        Tensor[B x D x T]
-            Quantized continuous representation of input
-        Tensor[1]
-            Commitment loss to train encoder to predict vectors closer to codebook
-            entries
-        Tensor[1]
-            Codebook loss to update the codebook
-        Tensor[B x T]
-            Codebook indices (quantized discrete representation of input)
-        Tensor[B x D x T]
-            Projected latents (continuous representation of input before quantization)
-        """
-
-        # Factorized codes (ViT-VQGAN) Project input into low-dimensional space
-
-        z_e = self.in_proj(z)  # z_e : (B x D x T)
-        z_q, indices = self.decode_latents(z_e)
-
-        commitment_loss = F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2])
-        codebook_loss = F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2])
-
-        z_q = (
-            z_e + (z_q - z_e).detach()
-        )  # noop in forward pass, straight-through gradient estimator in backward pass
-
-        z_q = self.out_proj(z_q)
-
-        return z_q, commitment_loss, codebook_loss, indices, z_e
-
-    def embed_code(self, embed_id):
-        return F.embedding(embed_id, self.codebook.weight)
-
-    def decode_code(self, embed_id):
-        return self.embed_code(embed_id).transpose(1, 2)
-
-    def decode_latents(self, latents):
-        encodings = rearrange(latents, "b d t -> (b t) d")
-        codebook = self.codebook.weight  # codebook: (N x D)
-
-        # L2 normalize encodings and codebook (ViT-VQGAN)
-        encodings = F.normalize(encodings)
-        codebook = F.normalize(codebook)
-
-        # Compute euclidean distance with codebook
-        dist = (
-            encodings.pow(2).sum(1, keepdim=True)
-            - 2 * encodings @ codebook.t()
-            + codebook.pow(2).sum(1, keepdim=True).t()
-        )
-        indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
-        z_q = self.decode_code(indices)
-
-        if(self.training):
-            onehots = torch.nn.functional.one_hot(indices, self.codebook_size).float()  # B, T, codebook_size
-            stale_codes = (onehots.sum(0).sum(0) == 0).float()
-            self.stale_counter = self.stale_counter * stale_codes + stale_codes
-
-            # random replace codes that haven't been used for a while
-            replace_code = (self.stale_counter == self.stale_tolerance).float() # codebook_size
-            if replace_code.sum(-1) > 0:
-                print("Replace {} codes".format(replace_code.sum(-1)))
-                random_input_idx = torch.randperm(encodings.shape[0])
-                random_input = encodings[random_input_idx].view(encodings.shape)
-                if random_input.shape[0] < self.codebook_size:
-                    random_input = torch.cat([random_input]*(self.codebook_size // random_input.shape[0] + 1), 0)
-                random_input = random_input[:self.codebook_size,:].contiguous()  # codebook_size, dim
-
-                self.codebook.weight.data = self.codebook.weight.data * (1 - replace_code).unsqueeze(-1) + random_input * replace_code.unsqueeze(-1)
-                self.stale_counter = self.stale_counter * (1 - replace_code)
-
-        return z_q, indices
-
-
-class ResidualVectorQuantize(nn.Module):
-    """
-    Introduced in SoundStream: An end2end neural audio codec
-    https://arxiv.org/abs/2107.03312
-    """
-
-    def __init__(
-        self,
-        input_dim: int = 512,
-        n_codebooks: int = 9,
-        codebook_size: int = 1024,
-        codebook_dim: Union[int, list] = 8,
-        quantizer_dropout: float = 0.0,
-        stale_tolerance: int = 100,
-        use_multi_layer_num:int = 1,
-    ):
-        super().__init__()
-        if isinstance(codebook_dim, int):
-            codebook_dim = [codebook_dim for _ in range(n_codebooks)]
-
-        self.n_codebooks = n_codebooks
-        self.codebook_dim = codebook_dim
-        self.codebook_size = codebook_size
-
-        self.quantizers = nn.ModuleList(
-            [
-                VectorQuantize(input_dim, codebook_size, codebook_dim[i], stale_tolerance=stale_tolerance, n_layer=use_multi_layer_num)
-                for i in range(n_codebooks)
-            ]
-        )
-        self.quantizer_dropout = quantizer_dropout
-
-    def forward(self, z, n_quantizers: int = None):
-        """Quantized the input tensor using a fixed set of `n` codebooks and returns
-        the corresponding codebook vectors
-        Parameters
-        ----------
-        z : Tensor[B x D x T]
-        n_quantizers : int, optional
-            No. of quantizers to use
-            (n_quantizers < self.n_codebooks ex: for quantizer dropout)
-            Note: if `self.quantizer_dropout` is True, this argument is ignored
-                when in training mode, and a random number of quantizers is used.
-        Returns
-        -------
-        dict
-            A dictionary with the following keys:
-
-            "z" : Tensor[B x D x T]
-                Quantized continuous representation of input
-            "codes" : Tensor[B x N x T]
-                Codebook indices for each codebook
-                (quantized discrete representation of input)
-            "latents" : Tensor[B x N*D x T]
-                Projected latents (continuous representation of input before quantization)
-            "vq/commitment_loss" : Tensor[1]
-                Commitment loss to train encoder to predict vectors closer to codebook
-                entries
-            "vq/codebook_loss" : Tensor[1]
-                Codebook loss to update the codebook
-        """
-        z_q = 0
-        residual = z
-        commitment_loss = 0
-        codebook_loss = 0
-
-        codebook_indices = []
-        latents = []
-
-        if n_quantizers is None:
-            n_quantizers = self.n_codebooks
-        if self.training:
-            n_quantizers = torch.ones((z.shape[0],)) * self.n_codebooks + 1
-            dropout = torch.randint(1, self.n_codebooks + 1, (z.shape[0],))
-            n_dropout = int(z.shape[0] * self.quantizer_dropout)
-            n_quantizers[:n_dropout] = dropout[:n_dropout]
-            n_quantizers = n_quantizers.to(z.device)
-        else:
-            n_quantizers = torch.ones((z.shape[0],)) * n_quantizers + 1
-            n_quantizers = n_quantizers.to(z.device)
-
-        for i, quantizer in enumerate(self.quantizers):
-            # if self.training is False and i >= n_quantizers:
-            #     break
-
-            z_q_i, commitment_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer(
-                residual
-            )
-
-            # Create mask to apply quantizer dropout
-            mask = (
-                torch.full((z.shape[0],), fill_value=i, device=z.device) < n_quantizers
-            ) 
-            z_q = z_q + z_q_i * mask[:, None, None]
-            residual = residual - z_q_i
-
-            # Sum losses
-            commitment_loss += (commitment_loss_i * mask).mean()
-            codebook_loss += (codebook_loss_i * mask).mean()
-
-            codebook_indices.append(indices_i)
-            latents.append(z_e_i)
-
-        codes = torch.stack(codebook_indices, dim=1)
-        latents = torch.cat(latents, dim=1)
-
-        encodings = F.one_hot(codes, self.codebook_size).float() # B N T 1024
-
-        return z_q, codes, latents, commitment_loss, codebook_loss, n_quantizers.clamp(max=self.n_codebooks).long() - 1
-    
-    def get_loss(self, x, quantized_prompt_embeds, commitment_loss, codebook_loss):
-        final_loss = commitment_loss * 0.25 + codebook_loss * 1.0 + (x - quantized_prompt_embeds).abs().mean()
-        return final_loss
-
-    def from_codes(self, codes: torch.Tensor):
-        """Given the quantized codes, reconstruct the continuous representation
-        Parameters
-        ----------
-        codes : Tensor[B x N x T]
-            Quantized discrete representation of input
-        Returns
-        -------
-        Tensor[B x D x T]
-            Quantized continuous representation of input
-        """
-        z_q = 0.0
-        z_p = []
-        n_codebooks = codes.shape[1]
-        for i in range(n_codebooks):
-            z_p_i = self.quantizers[i].decode_code(codes[:, i, :])
-            z_p.append(z_p_i)
-
-            z_q_i = self.quantizers[i].out_proj(z_p_i)
-            z_q = z_q + z_q_i
-        return z_q, torch.cat(z_p, dim=1), codes
-
-    def from_latents(self, latents: torch.Tensor):
-        """Given the unquantized latents, reconstruct the
-        continuous representation after quantization.
-
-        Parameters
-        ----------
-        latents : Tensor[B x N x T]
-            Continuous representation of input after projection
-
-        Returns
-        -------
-        Tensor[B x D x T]
-            Quantized representation of full-projected space
-        Tensor[B x D x T]
-            Quantized representation of latent space
-        """
-        z_q = 0
-        z_p = []
-        codes = []
-        dims = np.cumsum([0] + [q.codebook_dim for q in self.quantizers])
-
-        n_codebooks = np.where(dims <= latents.shape[1])[0].max(axis=0, keepdims=True)[
-            0
-        ]
-        for i in range(n_codebooks):
-            j, k = dims[i], dims[i + 1]
-            z_p_i, codes_i = self.quantizers[i].decode_latents(latents[:, j:k, :])
-            z_p.append(z_p_i)
-            codes.append(codes_i)
-
-            z_q_i = self.quantizers[i].out_proj(z_p_i)
-            z_q = z_q + z_q_i
-
-        return z_q, torch.cat(z_p, dim=1), torch.stack(codes, dim=1)
-

src/third_party/MuQ/src/muq/muq/muq.py

@@ -1,90 +0,0 @@
-import torch.nn as nn
-import torch
-from .models.muq_model import MuQModel
-from dataclasses import dataclass, field
-from typing import List, Optional
-from transformers.modeling_outputs import BaseModelOutput
-from huggingface_hub import PyTorchModelHubMixin
-
-@dataclass
-class MuQConfig:
-    label_rate:int = field(default=25)
-    num_codebooks:int = field(default=1)
-    codebook_dim:int = field(default=16)
-    codebook_size:int = field(default=4096)
-    features:List[str] = field(default_factory=lambda:["melspec_2048"])
-    hop_length:int = field(default=240)
-    n_mels:int = field(default=128)
-    conv_dim:int = field(default=512)
-    encoder_dim:int = field(default=1024)
-    encoder_depth:int = field(default=12)
-    mask_hop:float = field(default=0.4)
-    mask_prob:float = field(default=0.6)
-    is_flash:bool = field(default=False)
-    stat:Optional[dict] = field(default_factory=dict)
-    w2v2_config:Optional[dict] = field(default_factory=dict)
-    use_rvq_target:bool = field(default=False)
-    use_vq_target:bool = field(default=False)
-    use_encodec_target:bool = field(default=False)
-    rvq_ckpt_path: Optional[str] = field(default=None)
-    recon_loss_ratio: Optional[float] = field(default=None)
-    resume_checkpoint: Optional[str] = None
-    rvq_n_codebooks:int = field(default=8)
-    rvq_multi_layer_num:int = field(default=1)
-
-class MuQ(nn.Module, PyTorchModelHubMixin):
-    def __init__(self, config: MuQConfig):
-        super().__init__()
-        if isinstance(config, dict):
-            config = MuQConfig(**config)
-        self.config = config
-        self.model = MuQModel(
-            num_codebooks=config.num_codebooks,
-            codebook_dim=config.codebook_dim,
-            codebook_size=config.codebook_size,
-            features=config.features,
-            hop_length=config.hop_length,
-            n_mels=config.n_mels,
-            conv_dim=config.conv_dim,
-            encoder_dim=config.encoder_dim,
-            encoder_depth=config.encoder_depth,
-            mask_hop=config.mask_hop,
-            mask_prob=config.mask_prob,
-            is_flash=config.is_flash,
-            stat=config.stat,
-            w2v2_config=config.w2v2_config,
-            use_rvq_target=config.use_rvq_target,
-            use_vq_target=config.use_vq_target,
-            use_encodec_target=config.use_encodec_target,
-            rvq_ckpt_path=config.rvq_ckpt_path,
-            recon_loss_ratio=config.recon_loss_ratio,
-            label_rate=config.label_rate,
-            rvq_n_codebooks=config.rvq_n_codebooks,
-            rvq_multi_layer_num=config.rvq_multi_layer_num,
-        )
-    
-    def forward(self, x, attention_mask:Optional[torch.Tensor]=None, output_hidden_states:bool=True) ->BaseModelOutput:
-        """
-        Forward pass through the MuQ model and extract features.
-
-        Args:
-            x (torch.Tensor): Input waveform tensor of shape (batch_size, time).
-            attention_mask (torch.Tensor, optional): Mask to avoid performing attention on padding token indices.
-                Default is None.
-            output_hidden_states (bool, optional): Whether to return all hidden states or only the last one.
-                Default is False.
-
-        Returns:
-            BaseModelOutput: An object containing the last hidden state and optionally all hidden states.
-                - last_hidden_state (torch.Tensor): The last hidden state of the model, i.e. extracted MuQ features, of shape (batch_size, sequence_length, hidden_size).
-                - hidden_states (tuple(torch.Tensor), optional): A tuple containing all hidden states produced by the model,
-                each of shape (batch_size, sequence_length, hidden_size). Only returned if output_hidden_states is True.
-        """ 
-        _, hidden_states = self.model.get_predictions(x, attention_mask=attention_mask, is_features_only=True)
-        last_hidden_state = hidden_states[-1]
-        if not output_hidden_states:
-            return BaseModelOutput(last_hidden_state=last_hidden_state)
-        return BaseModelOutput(
-            last_hidden_state=last_hidden_state,
-            hidden_states=hidden_states
-        )

src/third_party/MuQ/src/muq/muq_mulan/__init__.py

@@ -1 +0,0 @@
-from .muq_mulan import MuQMuLan, MuQMuLanConfig, MuLanConfig, ModalModelConfig, TextTransformerConfig, AudioTransformerConfig

src/third_party/MuQ/src/muq/muq_mulan/models/audio.py

@@ -1,294 +0,0 @@
-from contextlib import suppress
-import torch
-import torch.nn as nn
-from einops import rearrange, repeat, reduce
-from einops.layers.torch import Rearrange
-from torchaudio.transforms import Spectrogram, TimeStretch, FrequencyMasking, TimeMasking
-from transformers import Wav2Vec2FeatureExtractor,AutoModel
-from ..modules.transformer import Transformer, LayerNorm, posemb_sincos_2d
-from ..modules.utils import print_once, round_down_nearest_multiple, frozen_params, Sequential
-
-
-def pair(t):
-    return (t, t) if not isinstance(t, tuple) else t
-
-class AudioSpectrogramTransformer(nn.Module):
-    def __init__(
-        self,
-        dim,
-        depth,
-        patch_size = 16,
-        dim_head = 64,
-        heads = 8,
-        attn_dropout = 0.,
-        ff_mult = 4,
-        ff_dropout = 0.,
-        accept_spec = False,
-        accept_spec_time_first = True,
-        spec_n_fft = 128,
-        spec_power = 2,
-        spec_win_length = 24,
-        spec_hop_length = None,
-        spec_pad = 0,
-        spec_center = True,
-        spec_pad_mode = 'reflect',
-        spec_aug_stretch_factor = 0.8,
-        spec_aug_freq_mask = 80,
-        spec_aug_time_mask = 80,
-        patch_dropout_prob = 0.25
-    ):
-        super().__init__()
-        self.dim = dim
-        self.depth = depth
-
-        self.patch_size = pair(patch_size)
-        patch_input_dim = self.patch_size[0] * self.patch_size[1]
-
-        self.to_patch_tokens = Sequential(
-            Rearrange('b (h p1) (w p2) -> b h w (p1 p2)', p1 = self.patch_size[0], p2 = self.patch_size[1]),
-            nn.LayerNorm(patch_input_dim),
-            nn.Linear(patch_input_dim, dim),
-            nn.LayerNorm(dim)
-        )
-
-        self.accept_spec = accept_spec
-        self.accept_spec_time_first = accept_spec_time_first
-
-        self.spec = Spectrogram(
-            n_fft = spec_n_fft,
-            power = spec_power,
-            win_length = spec_win_length,
-            hop_length = spec_hop_length,
-            pad = spec_pad,
-            center = spec_center,
-            pad_mode = spec_pad_mode
-        )
-
-        # SpecAugment - seems to be widely used in audio field https://arxiv.org/abs/1904.08779
-
-        self.aug = torch.nn.Sequential(
-            TimeStretch(spec_aug_stretch_factor, fixed_rate = True),
-            FrequencyMasking(freq_mask_param = spec_aug_freq_mask),
-            TimeMasking(time_mask_param = spec_aug_time_mask),
-        )
-
-        self.transformer = Transformer(
-            dim = dim,
-            depth = depth,
-            dim_head = dim_head,
-            heads = heads,
-            attn_dropout = attn_dropout,
-            ff_mult = ff_mult,
-            ff_dropout = ff_dropout
-        )
-
-        self.norm = LayerNorm(dim)
-
-        # patch dropout
-
-        self.patch_dropout_prob = patch_dropout_prob
-
-        # 2d dynamic positional bias
-
-        mlp_hidden_dim = dim // 4
-
-        self.dynamic_pos_bias_mlp = nn.Sequential(
-            nn.Linear(2, mlp_hidden_dim),
-            nn.SiLU(),
-            nn.Linear(mlp_hidden_dim, mlp_hidden_dim),
-            nn.SiLU(),
-            nn.Linear(mlp_hidden_dim, heads),
-            Rearrange('... i j h -> ... h i j')
-        )
-
-    def forward(
-        self,
-        x,
-        force_no_patch_dropout = False,
-        return_all_layers = False
-    ):
-        batch, device = x.shape[0], x.device
-        assert (self.accept_spec and x.ndim == 3) or (not self.accept_spec and x.ndim == 2)
-
-        if self.accept_spec and self.accept_spec_time_first:
-            x = rearrange(x, 'b t f -> b f t')
-
-        if not self.accept_spec:
-            x = self.spec(x)
-
-        if self.training:
-            x = self.aug(x)
-
-        # automatically crop if audio does not yield a 2d spectrogram that is divisible by patch sizes
-
-        height, width = x.shape[-2:]
-        patch_height, patch_width = self.patch_size
-
-        rounded_height, rounded_width = map(lambda args: round_down_nearest_multiple(*args), ((height, patch_height), (width, patch_width)))
-
-        if (height, width) != (rounded_height, rounded_width): # just keep printing to be annoying until it is fixed
-            print_once(f'spectrogram yielded shape of {(height, width)}, but had to be cropped to {(rounded_height, rounded_width)} to be patchified for transformer')
-
-        x = x[..., :rounded_height, :rounded_width]
-
-        # to patches
-
-        x = self.to_patch_tokens(x)
-
-        # get number of patches along height and width
-
-        _, num_patch_height, num_patch_width, _ = x.shape
-
-        # get 2d relative positions
-
-        grid = torch.stack(torch.meshgrid(
-            torch.arange(num_patch_height, device = device),
-            torch.arange(num_patch_width, device = device)
-        , indexing = 'ij'), dim = -1)
-
-        grid = rearrange(grid, '... c -> (...) c')
-
-        # 2d sinusoidal positional embedding
-
-        x = x + posemb_sincos_2d(x)
-
-        x = rearrange(x, 'b ... c -> b (...) c')
-
-        # patch dropout
-
-        if self.training and self.patch_dropout_prob > 0. and not force_no_patch_dropout:
-            n, device = x.shape[1], x.device
-
-            batch_indices = torch.arange(batch, device = device)
-            batch_indices = rearrange(batch_indices, '... -> ... 1')
-            num_patches_keep = max(1, int(n * (1 - self.patch_dropout_prob)))
-            patch_indices_keep = torch.randn(batch, n, device = device).topk(num_patches_keep, dim = -1).indices
-
-            x = x[batch_indices, patch_indices_keep]
-
-            grid = repeat(grid, '... -> b ...', b = batch)
-            grid = grid[batch_indices, patch_indices_keep]
-
-        # 2d relative positional bias
-
-        rel_dist = rearrange(grid, '... i c -> ... i 1 c') - rearrange(grid, '... j c -> ... 1 j c')
-        rel_pos_bias = self.dynamic_pos_bias_mlp(rel_dist.float())
-
-        # attention, what else
-
-        x, all_layers = self.transformer(x, rel_pos_bias = rel_pos_bias, return_all_layers = True)
-
-        # final global average and norm (most recent papers show this is superior to CLS token)
-
-        x = reduce(x, 'b n d -> b d', 'mean')
-
-        out = self.norm(x)
-
-        if not return_all_layers:
-            return out
-
-        return out, all_layers
-
-class AudioSpectrogramTransformerPretrained(nn.Module):
-    def __init__(
-        self,
-        model_name = 'm-a-p/MERT-v1-330M',
-        dim = 768,
-        model_dim = 1024,
-        sr = 24000,
-        tf_depth = 12,
-        dim_head = 64,
-        heads = 8,
-        attn_dropout = 0.,
-        ff_dropout = 0.,
-        ff_mult = 4,
-        use_layer_idx = -1,
-        frozen_pretrained = True,
-        hf_hub_cache_dir = None,
-    ):
-        super().__init__()
-        self.model_name = model_name
-        self.dim = dim
-        self.sr = sr
-        self.use_layer_idx = use_layer_idx # Which layer's features should be used
-        self.hf_hub_cache_dir = hf_hub_cache_dir
-
-        self.model_name
-
-        self._init_pretrained_model(model_name)
-
-        self.aggregator = nn.Conv1d(in_channels=25, out_channels=1, kernel_size=1)
-
-
-        self.transformer = Transformer(
-            dim = dim,
-            depth = tf_depth,
-            dim_head = dim_head,
-            heads = heads,
-            attn_dropout = attn_dropout,
-            ff_dropout = ff_dropout,
-            ff_mult = ff_mult
-        ) # if tf_depth > 0 else torch.nn.Identity()
-
-        self.proj = nn.Linear(model_dim, dim)
-
-        if frozen_pretrained:
-            frozen_params(self.model)
-            frozen_params(self.aggregator)
-        self.frozen_pretrained = frozen_pretrained
-
-    def _init_pretrained_model(self, model_name):
-        if 'muq' in model_name.lower():
-            from muq import MuQ
-            self.model = MuQ.from_pretrained(model_name, cache_dir=self.hf_hub_cache_dir)
-        else:
-            self.model = AutoModel.from_pretrained(model_name, trust_remote_code=True, cache_dir=self.hf_hub_cache_dir)
-            self.processor = Wav2Vec2FeatureExtractor.from_pretrained(model_name,trust_remote_code=True, cache_dir=self.hf_hub_cache_dir)
-            
-            assert self.processor.sampling_rate == self.sr
-
-    @property
-    def device(self):
-        return next(self.model.parameters()).device
-    
-    @property
-    def dtype(self):
-        return next(self.model.parameters()).dtype
-
-    def _forward_pretrained_model(self, x):
-        if 'muq' in self.model_name.lower():
-            outputs = self.model(x, output_hidden_states=True)
-            return outputs.hidden_states # 13 layer x [batch_size, Time steps, 1024 feature_dim]
-        else:
-            inputs = self.processor(x, sampling_rate=self.sr, return_tensors="pt")
-            input_values = inputs['input_values'].squeeze(0).to(self.device, dtype = self.dtype)
-            outputs = self.model(input_values, output_hidden_states=True) # [25 layer, batch_size, Time steps, 1024 feature_dim]
-            return outputs.hidden_states
-    
-    def forward(
-        self,
-        x,
-        return_all_layers = False,
-        return_mean = True,
-        no_proj = False,
-    ):
-        batch, device = x.shape[0], x.device
-        
-        with torch.no_grad() if self.frozen_pretrained else suppress():
-            outputs = self._forward_pretrained_model(x)
-        layer_hidden_states = outputs[self.use_layer_idx]
-
-        if no_proj:
-            outputs = layer_hidden_states
-        else:
-            outputs = self.proj(layer_hidden_states)
-        outputs, layer_results = self.transformer(outputs, return_all_layers=True)
-
-        if return_mean:
-            outputs = outputs.mean(dim = -2)
-
-        if return_all_layers:
-            return outputs, layer_results
-        return outputs
-
-  

src/third_party/MuQ/src/muq/muq_mulan/models/mulan.py

@@ -1,148 +0,0 @@
-import math
-from typing import List, Optional, Union
-from collections import OrderedDict
-from functools import partial
-
-import torch
-from torch import nn, einsum
-
-from .audio import AudioSpectrogramTransformer, AudioSpectrogramTransformerPretrained
-from .text import TextTransformer, TextTransformerPretrained
-from ..modules.contrastive import RankSoftmaxContrastiveLearning, SoftmaxContrastiveLearning, SigmoidContrastiveLearning, MultiLayerContrastiveLoss, interspersed_indices
-from ..modules.utils import exists, default, l2norm
-
-class MuLanModel(nn.Module):
-    def __init__(
-        self,
-        audio_transformer: Union[AudioSpectrogramTransformer, AudioSpectrogramTransformerPretrained],
-        text_transformer: Union[TextTransformer, TextTransformerPretrained],
-        dim_latent = 128,                       # they use 128
-        decoupled_contrastive_learning = True,  # think this was used, make it optional
-        hierarchical_contrastive_loss = False,
-        hierarchical_contrastive_loss_layers = None,
-        sigmoid_contrastive_loss = False,
-        rank_contrast = False,    # apply contrast on rank dimension
-        proj_to_latent = True,
-        norm_type = 'l2norm',
-        **kwargs,
-    ):
-        super().__init__()
-        self.dim_latent = dim_latent
-
-        # audio and text transformer
-        self.audio = audio_transformer
-        self.text = text_transformer
-
-        # two linear layers to project embeddings to latent space
-        if proj_to_latent:
-            self.text_to_latents = nn.Linear(self.text.dim, dim_latent)
-            self.audio_to_latents = nn.Linear(self.audio.dim, dim_latent)
-
-        self.sigmoid_contrastive_loss = sigmoid_contrastive_loss
-        self.decoupled_contrastive_learning = decoupled_contrastive_learning
-        self.rank_contrast = rank_contrast
-        self.norm_type = norm_type
-
-        # use decoupled contrastive learning or not, where self.contrast is loss module for contrastive learning
-        if sigmoid_contrastive_loss:
-            klass = SigmoidContrastiveLearning
-        else: 
-            if rank_contrast:
-                klass = partial(RankSoftmaxContrastiveLearning,  decoupled_contrastive_learning = decoupled_contrastive_learning) 
-            else:
-                klass = partial(SoftmaxContrastiveLearning, decoupled_contrastive_learning = decoupled_contrastive_learning)
-
-        self.contrast = klass() 
-
-        self.multi_layer_contrastive_learning = None
-
-        if hierarchical_contrastive_loss:
-            num_layers = default(hierarchical_contrastive_loss_layers, min(audio_transformer.depth, text_transformer.depth) - 1) 
-            assert num_layers > 0
-
-            self.register_buffer('text_layers_indices', interspersed_indices(num_layers, text_transformer.depth)) 
-            self.register_buffer('audio_layers_indices', interspersed_indices(num_layers, audio_transformer.depth))
-
-            self.multi_layer_contrastive_learning = MultiLayerContrastiveLoss(
-                audio_dim = self.audio.dim,
-                text_dim = self.text.dim,
-                dim_latent = dim_latent,
-                layers = num_layers,
-                decoupled_contrastive_learning = decoupled_contrastive_learning,
-                sigmoid_contrastive_loss = sigmoid_contrastive_loss
-            )
-
-    def get_audio_latents(
-        self,
-        wavs,
-        return_all_layers = False,
-    ):
-        audio_embeds, audio_layers = self.audio(wavs, return_all_layers = True)
-        audio_latents = self.audio_to_latents(audio_embeds)
-        out = self._norm_latents(audio_latents) #->[Batch, Feat=128]
-
-        if not return_all_layers:
-            return out
-
-        return out, audio_layers #[nLayer=5, Batch=2, 15, 512]
-
-    def get_text_latents(
-        self,
-        texts = None,
-        raw_texts: Optional[List[str]] = None,
-        return_all_layers = False
-    ):
-        text_embeds, text_layers = self.text(texts, raw_texts = raw_texts, return_all_layers = True)
-        text_latents = self.text_to_latents(text_embeds)
-        out = self._norm_latents(text_latents)
-
-        if not return_all_layers:
-            return out
-
-        return out, text_layers
-    
-    def _norm_latents(self, latents):
-        if self.norm_type == 'l2norm':
-            return l2norm(latents)
-        else:
-            return self.norm(latents)
-
-    def forward(
-        self,
-        wavs,
-        texts = None,
-        raw_texts: Optional[List[str]] = None,
-        return_latents = False,
-        return_similarities = False,
-        return_pairwise_similarities = False,
-    ):
-        batch, device = wavs.shape[0], wavs.device
-        
-        # both latents are of [Batch, Feat=128]
-        audio_latents, audio_layers = self.get_audio_latents(wavs, return_all_layers = True)
-        text_latents, text_layers = self.get_text_latents(texts, raw_texts = raw_texts, return_all_layers = True)
-
-        if return_latents: # used in inference
-            return audio_latents, text_latents
-
-        if return_similarities:
-            return einsum('i d, i d -> i', audio_latents, text_latents)
-
-        if return_pairwise_similarities:
-            cosine_sim = einsum('i d, j d -> i j', audio_latents, text_latents) 
-            return cosine_sim
-
-        cl_loss = self.contrast(audio_latents, text_latents) #contrastive loss
-
-        if not exists(self.multi_layer_contrastive_learning):
-            return cl_loss
-
-        audio_layers = audio_layers[self.audio_layers_indices]
-        text_layers = text_layers[self.text_layers_indices]
-
-        hierarchical_cl_loss = self.multi_layer_contrastive_learning(
-            audio_layers = audio_layers,
-            text_layers = text_layers
-        )
-
-        return cl_loss + hierarchical_cl_loss 

src/third_party/MuQ/src/muq/muq_mulan/models/text.py

@@ -1,241 +0,0 @@
-from typing import Optional, List
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from x_clip.tokenizer import tokenizer
-from einops import rearrange, repeat, reduce, pack, unpack
-from transformers import AutoTokenizer,XLMRobertaModel,AutoModelForCausalLM
-
-from ..modules.utils import *
-from ..modules.transformer import Transformer, LayerNorm
-from ..modules.utils import frozen_params
-
-
-# text transformer
-
-class TextTransformer(nn.Module):
-    def __init__(
-        self,
-        dim,
-        depth,
-        num_tokens = tokenizer.vocab_size,
-        max_seq_len = 256,
-        dim_head = 64,
-        heads = 8,
-        attn_dropout = 0.,
-        ff_dropout = 0.,
-        ff_mult = 4,
-        pad_id = 0
-    ):
-        super().__init__()
-        self.dim = dim
-
-        self.token_emb = nn.Embedding(num_tokens, dim)
-        self.pos_emb = nn.Embedding(max_seq_len, dim)
-
-        self.depth = depth
-        self.max_seq_len = max_seq_len
-
-        self.cls_token = nn.Parameter(torch.randn(dim))
-
-        self.transformer = Transformer(
-            dim = dim,
-            depth = depth,
-            dim_head = dim_head,
-            heads = heads,
-            attn_dropout = attn_dropout,
-            ff_dropout = ff_dropout,
-            ff_mult = ff_mult
-        )
-
-        self.pad_id = pad_id
-        self.norm = LayerNorm(dim)
-
-    @property
-    def device(self):
-        return next(self.parameters()).device
-
-    def forward(
-        self,
-        x = None,
-        raw_texts: Optional[List[str]] = None,
-        mask = None,
-        return_all_layers = False
-    ):
-        assert exists(x) ^ exists(raw_texts)
-
-        if exists(raw_texts):
-            x = tokenizer.tokenize(raw_texts).to(self.device)
-
-        if not exists(mask):
-            mask = x != self.pad_id
-
-        b, n, device = *x.shape, x.device
-
-        # token embedding + positional embedding
-
-        x = self.token_emb(x)
-
-        assert n <= self.max_seq_len, f'text sequence length {n} must be less than {self.max_seq_len}'
-
-        x = x + self.pos_emb(torch.arange(n, device = device))
-
-        # cls tokens, as in bert
-
-        cls_tokens = repeat(self.cls_token, 'd -> b d', b = b)
-        x, ps = pack([cls_tokens, x], 'b * d')
-
-        # account for attending to cls token with self attention mask
-
-        mask = F.pad(mask, (1, 0), value = True)
-
-        # attention
-
-        x, all_layers = self.transformer(x, mask = mask, return_all_layers = True)
-
-        # unpack the cls tokens
-
-        cls_tokens, _ = unpack(x, ps, 'b * d')
-
-        out = self.norm(cls_tokens)
-
-        if not return_all_layers:
-            return out
-
-        return out, all_layers
-
-class TextPretrainedModelType:
-    Roberta = 'roberta'
-    Qwen = 'qwen'
-
-class TextTransformerPretrained(nn.Module):
-    def __init__(
-        self,
-        model_name = 'xlm-roberta-base',
-        dim = 768,
-        model_dim = None,
-        max_seq_len = 256,
-        tf_depth = 12,
-        dim_head = 64,
-        heads = 8,
-        attn_dropout = 0.,
-        ff_dropout = 0.,
-        ff_mult = 4,
-        frozen_pretrained = True,
-        hf_hub_cache_dir = None,
-    ):
-        super().__init__()
-        self.dim = dim 
-
-        self.model_name = model_name
-
-        self.hf_hub_cache_dir = hf_hub_cache_dir
-
-        self.pretrained_model_type = self._get_pretrained_model_type(model_name)
-
-        self.model = self._init_pretrained_model()
-
-        self._tokenizer = None
-
-        self.max_seq_len = max_seq_len
-
-        self.transformer = Transformer(
-            dim = dim,
-            depth = tf_depth,
-            dim_head = dim_head,
-            heads = heads,
-            attn_dropout = attn_dropout,
-            ff_dropout = ff_dropout,
-            ff_mult = ff_mult
-        ) # if tf_depth > 0 else torch.nn.Identity()
-
-        is_proj = exists(model_dim) and model_dim != dim
-
-        self.proj = nn.Linear(model_dim, dim) if is_proj else torch.nn.Identity()
-        if frozen_pretrained:
-            frozen_params(self.model)
-        self.frozen_pretrained = frozen_pretrained
-    
-    @staticmethod
-    def _get_pretrained_model_type(model_name):
-        if 'xlm-roberta' in model_name:
-            return TextPretrainedModelType.Roberta
-        elif 'Qwen' in model_name:
-            return TextPretrainedModelType.Qwen
-        else:
-            raise ValueError(f"Unknown pretrained model named: {model_name}")
-    
-    def _init_pretrained_model(self):
-        if self.pretrained_model_type == TextPretrainedModelType.Roberta:
-            model = XLMRobertaModel.from_pretrained(self.model_name, trust_remote_code=True, cache_dir=self.hf_hub_cache_dir)
-        elif self.pretrained_model_type == TextPretrainedModelType.Qwen:
-            model = AutoModelForCausalLM.from_pretrained(self.model_name, trust_remote_code=True, fp16=True, cache_dir=self.hf_hub_cache_dir)
-        else:
-            raise ValueError(f"Failed to init pretrained model type: {self.pretrained_model_type}")
-        return model
-    
-    def _init_tokenizer(self):
-        if self.pretrained_model_type == TextPretrainedModelType.Roberta:
-            tokenizer = AutoTokenizer.from_pretrained(self.model_name, trust_remote_code=True, cache_dir=self.hf_hub_cache_dir)
-        elif self.pretrained_model_type == TextPretrainedModelType.Qwen:
-            tokenizer = AutoTokenizer.from_pretrained(self.model_name, trust_remote_code=True, cache_dir=self.hf_hub_cache_dir)
-            tokenizer.pad_token = '<|im_end|>'
-        else:
-            raise ValueError(f"Failed to init tokenizer of pretrained model type: {self.pretrained_model_type}")
-        return tokenizer
-
-    @property
-    def tokenizer(self):
-        if not exists(self._tokenizer):
-            self._tokenizer = self._init_tokenizer()
-        return self._tokenizer
-    
-    @property
-    def device(self):
-        return next(self.model.parameters()).device
-    
-    @property
-    def dtype(self):
-        return next(self.transformer.parameters()).dtype
-    
-    
-    def pred_pretrained_model_hidden(self, **kw):
-        if self.pretrained_model_type == TextPretrainedModelType.Roberta:
-            outputs = self.model(**kw)
-            outputs = outputs.last_hidden_state
-        elif self.pretrained_model_type == TextPretrainedModelType.Qwen:
-            last_hidden_state = self.model(**kw, output_hidden_states=True)['hidden_states'][-1]
-            outputs = last_hidden_state.to(dtype = self.dtype)
-        else:
-            raise ValueError(f"Unknown pretrained model type: {self.pretrained_model_type}")
-        return outputs
-
-    def forward(
-        self,
-        x = None,
-        raw_texts: Optional[List[str]] = None,
-        mask = None,
-        return_all_layers = False,
-        return_mean = True
-    ):
-        assert exists(x) ^ exists(raw_texts)
-        with torch.no_grad():
-            if exists(raw_texts):
-                inputs = self.tokenizer(raw_texts, return_tensors='pt', padding=True)
-                inputs = inputs.to(self.device)
-
-            if exists(mask):
-                inputs['attention_mask'] = mask
-            
-            outputs = self.pred_pretrained_model_hidden(**inputs)
-
-        outputs = self.proj(outputs)
-
-        outputs, layer_results = self.transformer(outputs, return_all_layers=True)
-        if return_mean:
-            outputs = outputs.mean(dim = -2)
-        
-        if return_all_layers:
-            return outputs, layer_results
-        return outputs

src/third_party/MuQ/src/muq/muq_mulan/modules/contrastive.py

@@ -1,238 +0,0 @@
-import math
-from functools import partial
-
-import torch.nn as nn
-import torch
-import torch.nn.functional as F
-from einops import rearrange, reduce
-from torch import einsum
-import torch.distributed as dist
-
-from .utils import exists, l2norm, log, print_once
-from .distributed import AllGather
-from .extend_distributed import all_gather
-from .transformer import LayerNorm
-
-def matrix_diag(t):
-    device = t.device
-    i, j = t.shape[-2:]
-    num_diag_el = min(i, j)
-    i_range = torch.arange(i, device = device)
-    j_range = torch.arange(j, device = device)
-    diag_mask = rearrange(i_range, 'i -> i 1') == rearrange(j_range, 'j -> 1 j')
-    diag_el = t.masked_select(diag_mask)
-    return rearrange(diag_el, '(b d) -> b d', d = num_diag_el)
-
-# contrastive losses
-
-class SoftmaxContrastiveLearning(nn.Module):
-    def __init__(
-        self,
-        *,
-        layers = 1,
-        decoupled_contrastive_learning = False,
-        init_temp = 10
-    ):
-        super().__init__()
-        self.temperatures = nn.Parameter(torch.ones(layers, 1, 1) * math.log(init_temp))
-        self.decoupled_contrastive_learning = decoupled_contrastive_learning
-
-        self.all_gather = AllGather(dim = 2)
-
-    @property
-    def device(self):
-        return next(self.parameters()).device
-    
-    def forward(self, audio_latents, text_latents):
-        if audio_latents.ndim == 2:
-            audio_latents = rearrange(audio_latents, '... -> 1 ...')
-
-        if text_latents.ndim == 2:
-            text_latents = rearrange(text_latents, '... -> 1 ...')
-
-        batch = audio_latents.shape[1]
-
-        if self.all_gather.is_distributed:
-            latents = torch.stack((audio_latents, text_latents))
-            latents, _ = self.all_gather(latents)
-            audio_latents, text_latents = latents
-
-        sims = einsum('l i d, l j d -> l i j', audio_latents, text_latents)
-
-        sims = sims * self.temperatures.exp()
-
-        cosine_sims_exp = sims.exp() # Similarity matrix  [Rank, N, N]
-
-        numerator = matrix_diag(cosine_sims_exp) # Take diagonal elements, that is, for t [l, i, j], take all elements of i==j to obtain a array of l * min (i, j)
-
-        if self.decoupled_contrastive_learning:
-            eye = torch.eye(batch, device = self.device, dtype = torch.bool)
-            cosine_sims_exp = cosine_sims_exp.masked_fill(eye, 0.) # Set the diagonal to 0
-
-        denominator_i = reduce(cosine_sims_exp, 'l i j -> l i', 'sum')
-        denominator_j = reduce(cosine_sims_exp, 'l i j -> l j', 'sum')
-
-        contrastive_loss = -log(numerator) + 0.5 * (log(denominator_i) + log(denominator_j))
-
-        contrastive_loss = reduce(contrastive_loss, 'l n -> l', 'mean')
-        return contrastive_loss.sum()
-
-
-class RankSoftmaxContrastiveLearning(nn.Module):
-    def __init__(
-        self,
-        *,
-        layers = 1,
-        decoupled_contrastive_learning = False,
-        init_temp = 10,
-    ):
-        super().__init__()
-        self.temperatures = nn.Parameter(torch.ones(layers, 1, 1) * math.log(init_temp))
-        self.decoupled_contrastive_learning = decoupled_contrastive_learning
-
-    
-    @property
-    def device(self):
-        return next(self.parameters()).device
-
-    def forward(self, audio_latents, text_latents):
-        if audio_latents.ndim == 2:
-            audio_latents = rearrange(audio_latents, '... -> 1 ...')
-
-        if text_latents.ndim == 2:
-            text_latents = rearrange(text_latents, '... -> 1 ...')
-
-        audio_latents = all_gather(audio_latents, None)
-        text_latents = all_gather(text_latents, None)
-
-        print_once("audio_latents:"+str(audio_latents.shape) + "text_latents:" + str(text_latents.shape))
-        
-        
-        batch = audio_latents.shape[1]
-        rank = audio_latents.shape[0]
-
-        audio_latents = rearrange(audio_latents, 'l i d -> (l i) d')
-        text_latents = rearrange(text_latents, 'l j d -> (l j) d')
-
-        sims = einsum('i d, j d -> i j', audio_latents, text_latents)
-
-        sims = sims * self.temperatures.exp()
-
-        sims = rearrange(sims, '1 i j -> i j')
-
-        cosine_sims_exp = sims.exp() # Similarity matrix  [Rank, N, N]
-
-
-        numerator = matrix_diag(cosine_sims_exp) # Take diagonal elements, that is, for t [l, i, j], take all elements of i==j to obtain a array of l * min (i, j)
-
-        if self.decoupled_contrastive_learning:
-            eye = torch.eye(batch*rank, device = self.device, dtype = torch.bool)
-            cosine_sims_exp = cosine_sims_exp.masked_fill(eye, 0.) # Set the diagonal to 0
-
-        denominator_i = reduce(cosine_sims_exp, 'i j -> i', 'sum')
-        denominator_j = reduce(cosine_sims_exp, 'i j -> j', 'sum')
-
-        contrastive_loss = -log(numerator) + 0.5 * (log(denominator_i) + log(denominator_j))
-
-        contrastive_loss = reduce(contrastive_loss, '1 n -> 1', 'mean')
-        return contrastive_loss
-
-
-class SigmoidContrastiveLearning(nn.Module):
-    """ https://arxiv.org/abs/2303.15343 """
-
-    def __init__(
-        self,
-        *,
-        layers = 1,
-        init_temp = 10,
-        init_bias = -10
-    ):
-        super().__init__()
-        self.temperatures = nn.Parameter(torch.ones(layers, 1, 1) * math.log(init_temp))
-        self.bias = nn.Parameter(torch.ones(layers, 1, 1) * init_bias)
-
-        self.all_gather = AllGather(dim = 1, all_reduce_grads = True)
-
-    @property
-    def device(self):
-        return next(self.parameters()).device
-
-    def forward(self, audio_latents, text_latents):
-        device = self.device
-
-        if audio_latents.ndim == 2:
-            audio_latents = rearrange(audio_latents, '... -> 1 ...') # To [Rank, Batch, Latent]
-
-        if text_latents.ndim == 2:
-            text_latents = rearrange(text_latents, '... -> 1 ...')
-
-        text_latents, rank_sizes = self.all_gather(text_latents)
-
-        n = text_latents.shape[1]
-
-        sims = einsum('l i d, l j d -> l i j', audio_latents, text_latents) # Calculate dot product similarity between pairs
-
-        sims = sims * self.temperatures.exp() + self.bias
-
-        labels = torch.eye(n, device = device)
-
-        if exists(rank_sizes):
-            labels_by_ranks = labels.split(rank_sizes.tolist(), dim = 0) 
-            labels = labels_by_ranks[dist.get_rank()] # labels to the n elements of the current rank
-
-        labels = 2 * rearrange(labels, 'i j -> 1 i j') - torch.ones_like(sims)
-
-        return -F.logsigmoid(labels * sims).sum() / n
-
-
-
-
-# hierarchical cl loss
-
-def interspersed_indices(layers, total_layers):
-    assert total_layers >= layers
-    step = total_layers / layers
-    return (torch.arange(0, layers) * step).floor().long()
-
-class MultiLayerContrastiveLoss(nn.Module):
-    def __init__(
-        self,
-        *,
-        audio_dim,
-        text_dim,
-        dim_latent,
-        layers,
-        decoupled_contrastive_learning = False,
-        sigmoid_contrastive_loss = False
-    ):
-        super().__init__()
-        self.layers = layers
-
-        self.audio_norm = LayerNorm(audio_dim, scale = False)
-        self.audio_gamma = nn.Parameter(torch.ones(layers, 1, audio_dim))
-        self.audio_latent_weight = nn.Parameter(torch.randn(layers, audio_dim, dim_latent))
-        self.audio_latent_bias = nn.Parameter(torch.randn(layers, 1, dim_latent))
-
-        self.text_norm = LayerNorm(text_dim, scale = False)
-        self.text_gamma = nn.Parameter(torch.ones(layers, 1, text_dim))
-        self.text_latent_weight = nn.Parameter(torch.randn(layers, text_dim, dim_latent))
-        self.text_latent_bias = nn.Parameter(torch.randn(layers, 1, dim_latent))
-
-        klass = SigmoidContrastiveLearning if sigmoid_contrastive_loss else partial(SoftmaxContrastiveLearning, decoupled_contrastive_learning = decoupled_contrastive_learning)
-        self.contrast = klass(layers = layers)
-
-    def forward(self, *, audio_layers, text_layers):
-        device, batch = audio_layers.device, audio_layers.shape[1]
-
-        audio_gap = reduce(audio_layers, 'l b n d -> l b d', 'mean')
-        audio_embeds = self.audio_norm(audio_gap) * self.audio_gamma
-        audio_latents = einsum('l b d, l d e -> l b e', audio_embeds, self.audio_latent_weight) + self.audio_latent_bias
-        audio_latents = l2norm(audio_latents)
-
-        text_cls_tokens = text_layers[:, :, 0]
-        text_embeds = self.text_norm(text_cls_tokens) * self.text_gamma
-        text_latents = einsum('l b d, l d e -> l b e', text_embeds, self.text_latent_weight) + self.text_latent_bias
-        text_latents = l2norm(text_latents)
-
-        return self.contrast(audio_latents, text_latents)

src/third_party/MuQ/src/muq/muq_mulan/modules/distributed.py

@@ -1,83 +0,0 @@
-import torch
-from torch import nn
-from torch.autograd import Function
-import torch.distributed as dist
-
-from einops import rearrange
-
-def exists(val):
-    return val is not None
-
-# distributed helpers
-
-def all_gather_same_dim(t):
-    world_size = dist.get_world_size()
-    gathered_tensors = [torch.empty_like(t, device = t.device, dtype = t.dtype) for i in range(world_size)]
-    dist.all_gather(gathered_tensors, t)
-    return gathered_tensors
-
-def all_gather_variable_dim(t, dim = 0, sizes = None):
-    device, rank, world_size = t.device, dist.get_rank(), dist.get_world_size()
-
-    if not exists(sizes):
-        size = torch.tensor(t.shape[dim], device = device, dtype = torch.long)
-        sizes = all_gather_same_dim(size)
-        sizes = torch.stack(sizes)
-
-    if torch.unique(sizes).numel() == 1:
-        gathered_tensors = all_gather_same_dim(t)
-        return torch.cat(gathered_tensors, dim = dim), sizes
-
-    max_size = sizes.amax().item()
-
-    padded_t = pad_dim_to(t, max_size, dim = dim)
-    gathered_tensors = all_gather_same_dim(padded_t)
-
-    gathered_tensor = torch.cat(gathered_tensors, dim = dim)
-    seq = torch.arange(max_size, device = device)
-
-    mask = rearrange(seq, 'j -> 1 j') < rearrange(sizes, 'i -> i 1')
-    mask = rearrange(mask, 'i j -> (i j)')
-    seq = torch.arange(mask.shape[-1], device = device)
-    indices = seq[mask]
-
-    gathered_tensor = gathered_tensor.index_select(dim, indices)
-
-    return gathered_tensor, sizes
-
-class AllGatherFunction(Function):
-    @staticmethod
-    def forward(ctx, x, dim, sizes, all_reduce_grads):
-        x, batch_sizes = all_gather_variable_dim(x, dim = dim, sizes = sizes)
-        ctx.dim = dim
-        ctx.all_reduce_grads = all_reduce_grads
-        ctx.batch_sizes = batch_sizes.tolist()
-        return x, batch_sizes
-
-    @staticmethod
-    def backward(ctx, grads, _):
-        batch_sizes, rank = ctx.batch_sizes, dist.get_rank()
-        if ctx.all_reduce_grads:
-            dist.all_reduce(grads)
-
-        grads_by_rank = grads.split(batch_sizes, dim = ctx.dim)
-        return grads_by_rank[rank], None, None, None
-
-class AllGather(nn.Module):
-    def __init__(
-        self,
-        dim,
-        *,
-        all_reduce_grads = False
-    ):
-        super().__init__()
-        self.dim = dim
-        self.all_reduce_grads = all_reduce_grads
-        self.is_distributed = dist.is_initialized() and dist.get_world_size() > 1
-
-    def forward(
-        self,
-        x,
-        sizes = None
-    ):
-        return AllGatherFunction.apply(x, self.dim, sizes, self.all_reduce_grads)

src/third_party/MuQ/src/muq/muq_mulan/modules/extend_distributed.py

@@ -1,604 +0,0 @@
-
-import builtins
-import os
-import sys
-
-import torch
-import torch.distributed as dist
-from torch.autograd import Function
-from torch.autograd.profiler import record_function
-from torch.nn.parallel import DistributedDataParallel as DDP
-
-import torch.distributed as dist
-
-try:
-    import torch_ccl
-except ImportError as e:
-    # print(e)
-    torch_ccl = False
-
-try:
-    import torch_ucc
-except ImportError as e:
-    torch_ucc = False
-
-
-my_rank = -1
-my_size = -1
-my_local_rank = 1
-my_local_size = 1
-alltoall_supported = False
-a2a_impl = os.environ.get("DLRM_ALLTOALL_IMPL", "")
-
-myreq = None
-
-
-def env2int(env_list, default=-1):
-    for e in env_list:
-        val = int(os.environ.get(e, -1))
-        if val >= 0:
-            return val
-    return default
-
-
-def get_my_slice(n):
-    k, m = divmod(n, my_size)
-    return slice(
-        my_rank * k + min(my_rank, m), (my_rank + 1) * k + min(my_rank + 1, m), 1
-    )
-
-
-def get_split_lengths(n):
-    k, m = divmod(n, my_size)
-    if m == 0:
-        splits = None
-        my_len = k
-    else:
-        splits = [(k + 1) if i < m else k for i in range(my_size)]
-        my_len = splits[my_rank]
-    return (my_len, splits)
-
-
-def init_distributed(rank=-1, local_rank=-1, size=-1, use_gpu=False, backend=""):
-    global myreq
-    global my_rank
-    global my_size
-    global my_local_rank
-    global my_local_size
-    global a2a_impl
-    global alltoall_supported
-
-    # guess MPI ranks from env (works for IMPI, OMPI and MVAPICH2)
-    num_mpi_ranks = env2int(
-        ["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"]
-    )
-    if backend == "" and num_mpi_ranks > 1:
-        if torch_ccl and env2int(["CCL_WORKER_COUNT"]) > 0:
-            backend = "ccl"
-        elif use_gpu and dist.is_nccl_available():
-            backend = "nccl"
-        elif dist.is_mpi_available():
-            backend = "mpi"
-        else:
-            print(
-                "WARNING: MPI multi-process launch detected but PyTorch MPI backend not available."
-            )
-            backend = "gloo"
-
-    if backend != "":
-        # guess Rank and size
-        if rank == -1:
-            rank = env2int(
-                ["PMI_RANK", "OMPI_COMM_WORLD_RANK", "MV2_COMM_WORLD_RANK", "RANK"], 0
-            )
-        if size == -1:
-            size = env2int(
-                [
-                    "PMI_SIZE",
-                    "OMPI_COMM_WORLD_SIZE",
-                    "MV2_COMM_WORLD_SIZE",
-                    "WORLD_SIZE",
-                ],
-                1,
-            )
-        if not os.environ.get("RANK", None) and rank != -1:
-            os.environ["RANK"] = str(rank)
-        if not os.environ.get("WORLD_SIZE", None) and size != -1:
-            os.environ["WORLD_SIZE"] = str(size)
-        if not os.environ.get("MASTER_PORT", None):
-            os.environ["MASTER_PORT"] = "29500"
-        if not os.environ.get("MASTER_ADDR", None):
-            local_size = env2int(
-                [
-                    "MPI_LOCALNRANKS",
-                    "OMPI_COMM_WORLD_LOCAL_SIZE",
-                    "MV2_COMM_WORLD_LOCAL_SIZE",
-                ],
-                1,
-            )
-            if local_size != size and backend != "mpi":
-                print(
-                    "Warning: Looks like distributed multinode run but MASTER_ADDR env not set, using '127.0.0.1' as default"
-                )
-                print(
-                    "If this run hangs, try exporting rank 0's hostname as MASTER_ADDR"
-                )
-            os.environ["MASTER_ADDR"] = "127.0.0.1"
-
-    if size > 1:
-        if local_rank == -1:
-            my_local_rank = env2int(
-                [
-                    "MPI_LOCALRANKID",
-                    "OMPI_COMM_WORLD_LOCAL_RANK",
-                    "MV2_COMM_WORLD_LOCAL_RANK",
-                    "LOCAL_RANK",
-                ],
-                0,
-            )
-        else:
-            my_local_rank = local_rank
-        my_local_size = env2int(
-            [
-                "MPI_LOCALNRANKS",
-                "OMPI_COMM_WORLD_LOCAL_SIZE",
-                "MV2_COMM_WORLD_LOCAL_SIZE",
-            ],
-            1,
-        )
-        if use_gpu:
-            if my_local_size > torch.cuda.device_count():
-                print(
-                    "Not sufficient GPUs available... local_size = %d, ngpus = %d"
-                    % (my_local_size, torch.cuda.device_count())
-                )
-                sys.exit(1)
-            torch.cuda.set_device(my_local_rank)
-        dist.init_process_group(backend, rank=rank, world_size=size)
-        my_rank = dist.get_rank()
-        my_size = dist.get_world_size()
-        if my_rank == 0:
-            print("Running on %d ranks using %s backend" % (my_size, backend))
-        if hasattr(dist, "all_to_all_single"):
-            try:
-                t = torch.zeros([4])
-                if use_gpu:
-                    t = t.cuda()
-                dist.all_to_all_single(t, t)
-                alltoall_supported = True
-            except RuntimeError as err:
-                print("fail to enable all_to_all_single primitive: %s" % err)
-        if a2a_impl == "alltoall" and alltoall_supported == False:
-            print(
-                "Requested DLRM_ALLTOALL_IMPL=%s but backend %s does not support it, use scatter/gather based alltoall"
-                % (a2a_impl, backend)
-            )
-            a2a_impl = "scatter"
-        if a2a_impl != "":
-            print("Using DLRM_ALLTOALL_IMPL=%s" % a2a_impl)
-    else:
-        my_rank = 0
-        my_size = 1
-        my_local_rank = 0
-        my_local_size = 1
-    print_all(
-        "world size: %d, current rank: %d, local rank: %d"
-        % (my_size, my_rank, my_local_rank)
-    )
-    myreq = Request()
-
-
-class Request(object):
-    def __init__(self):
-        self.req = None
-        self.tensor = None
-        self.WaitFunction = All2All_Scatter_Wait
-
-    def wait(self):
-        ret = self.WaitFunction.apply(*self.tensor)
-        self.req = None
-        self.tensor = None
-        return ret
-
-
-class All2All_ScatterList_Req(Function):
-    @staticmethod
-    def forward(ctx, a2a_info, *inputs):
-        global myreq
-        batch_split_lengths = (
-            a2a_info.global_batch_partition_slices
-            if a2a_info.global_batch_partition_slices
-            else a2a_info.local_batch_num
-        )
-        table_split_lengths = (
-            a2a_info.global_table_wise_parition_slices
-            if a2a_info.global_table_wise_parition_slices
-            else [a2a_info.local_table_num] * my_size
-        )
-        gather_list = []
-        req_list = []
-        for i in range(my_size):
-            for j in range(table_split_lengths[i]):
-                out_tensor = inputs[0].new_empty(
-                    [a2a_info.local_batch_num, a2a_info.emb_dim]
-                )
-                scatter_list = (
-                    list(inputs[j].split(batch_split_lengths, dim=0))
-                    if i == my_rank
-                    else []
-                )
-                req = dist.scatter(out_tensor, scatter_list, src=i, async_op=True)
-                gather_list.append(out_tensor)
-                req_list.append(req)
-        myreq.req = req_list
-        myreq.tensor = tuple(gather_list)
-        myreq.a2a_info = a2a_info
-        return myreq.tensor
-
-    @staticmethod
-    def backward(ctx, *grad_output):
-        global myreq
-        for r in myreq.req:
-            r.wait()
-        myreq.req = None
-        grad_inputs = myreq.tensor
-        myreq.tensor = None
-        return (None, *grad_inputs)
-
-
-class All2All_ScatterList_Wait(Function):
-    @staticmethod
-    def forward(ctx, *output):
-        global myreq
-        ctx.a2a_info = myreq.a2a_info
-        for r in myreq.req:
-            r.wait()
-        myreq.req = None
-        myreq.tensor = None
-        return output
-
-    @staticmethod
-    def backward(ctx, *grad_output):
-        global myreq
-        a2a_info = ctx.a2a_info
-        grad_output = [t.contiguous() for t in grad_output]
-        batch_split_lengths = (
-            a2a_info.global_batch_partition_slices
-            if a2a_info.global_batch_partition_slices
-            else [a2a_info.local_batch_num] * my_size
-        )
-        per_rank_table_splits = (
-            a2a_info.global_table_wise_parition_slices
-            if a2a_info.global_table_wise_parition_slices
-            else [a2a_info.local_table_num] * my_size
-        )
-        grad_inputs = [
-            grad_output[0].new_empty([ctx.a2a_info.batch_size, ctx.a2a_info.emb_dim])
-            for _ in range(a2a_info.local_table_num)
-        ]
-        req_list = []
-        ind = 0
-        for i in range(my_size):
-            for j in range(per_rank_table_splits[i]):
-                gather_list = (
-                    list(grad_inputs[j].split(batch_split_lengths, dim=0))
-                    if i == my_rank
-                    else None
-                )
-                req = dist.gather(grad_output[ind], gather_list, dst=i, async_op=True)
-                req_list.append(req)
-                ind += 1
-        myreq.req = req_list
-        myreq.tensor = grad_inputs
-        return tuple(grad_output)
-
-
-class All2All_Scatter_Req(Function):
-    @staticmethod
-    def forward(ctx, a2a_info, *inputs):
-        global myreq
-        batch_split_lengths = (
-            a2a_info.global_batch_partition_slices
-            if a2a_info.global_batch_partition_slices
-            else a2a_info.local_batch_num
-        )
-        table_split_lengths = (
-            a2a_info.global_table_wise_parition_slices
-            if a2a_info.global_table_wise_parition_slices
-            else [a2a_info.local_table_num] * my_size
-        )
-        input = torch.cat(inputs, dim=1)
-        scatter_list = list(input.split(batch_split_lengths, dim=0))
-        gather_list = []
-        req_list = []
-        for i in range(my_size):
-            out_tensor = input.new_empty(
-                [a2a_info.local_batch_num, table_split_lengths[i] * a2a_info.emb_dim]
-            )
-            req = dist.scatter(
-                out_tensor, scatter_list if i == my_rank else [], src=i, async_op=True
-            )
-            gather_list.append(out_tensor)
-            req_list.append(req)
-        myreq.req = req_list
-        myreq.tensor = tuple(gather_list)
-        myreq.a2a_info = a2a_info
-        ctx.a2a_info = a2a_info
-        return myreq.tensor
-
-    @staticmethod
-    def backward(ctx, *grad_output):
-        global myreq
-        for r in myreq.req:
-            r.wait()
-        myreq.req = None
-        grad_input = myreq.tensor
-        grad_inputs = grad_input.split(ctx.a2a_info.emb_dim, dim=1)
-        myreq.tensor = None
-        return (None, *grad_inputs)
-
-
-class All2All_Scatter_Wait(Function):
-    @staticmethod
-    def forward(ctx, *output):
-        global myreq
-        ctx.a2a_info = myreq.a2a_info
-        for r in myreq.req:
-            r.wait()
-        myreq.req = None
-        myreq.tensor = None
-        return output
-
-    @staticmethod
-    def backward(ctx, *grad_output):
-        global myreq
-        assert len(grad_output) == my_size
-        scatter_list = [t.contiguous() for t in grad_output]
-        a2a_info = ctx.a2a_info
-        batch_split_lengths = (
-            a2a_info.global_batch_partition_slices
-            if a2a_info.global_batch_partition_slices
-            else a2a_info.local_batch_num
-        )
-        table_split_lengths = (
-            a2a_info.global_table_wise_parition_slices
-            if a2a_info.global_table_wise_parition_slices
-            else [a2a_info.local_table_num] * my_size
-        )
-        grad_input = grad_output[0].new_empty(
-            [a2a_info.batch_size, a2a_info.emb_dim * a2a_info.local_table_num]
-        )
-        gather_list = list(grad_input.split(batch_split_lengths, dim=0))
-        req_list = []
-        for i in range(my_size):
-            req = dist.gather(
-                scatter_list[i],
-                gather_list if i == my_rank else [],
-                dst=i,
-                async_op=True,
-            )
-            req_list.append(req)
-        myreq.req = req_list
-        myreq.tensor = grad_input
-        return grad_output
-
-
-class All2All_Req(Function):
-    @staticmethod
-    def forward(ctx, a2a_info, *inputs):
-        global myreq
-        with record_function("DLRM alltoall_req_fwd_single"):
-            batch_split_lengths = a2a_info.global_batch_partition_slices
-            if batch_split_lengths:
-                batch_split_lengths = [
-                    m * a2a_info.emb_dim * a2a_info.local_table_num
-                    for m in batch_split_lengths
-                ]
-            table_split_lengths = a2a_info.global_table_wise_parition_slices
-            if table_split_lengths:
-                table_split_lengths = [
-                    a2a_info.local_batch_num * e * a2a_info.emb_dim
-                    for e in table_split_lengths
-                ]
-            input = torch.cat(inputs, dim=1).view([-1])
-            output = input.new_empty(
-                [
-                    a2a_info.global_table_num
-                    * a2a_info.local_batch_num
-                    * a2a_info.emb_dim
-                ]
-            )
-            req = dist.all_to_all_single(
-                output, input, table_split_lengths, batch_split_lengths, async_op=True
-            )
-
-            myreq.req = req
-            myreq.tensor = []
-            myreq.tensor.append(output)
-            myreq.tensor = tuple(myreq.tensor)
-            a2a_info.batch_split_lengths = batch_split_lengths
-            a2a_info.table_split_lengths = table_split_lengths
-            myreq.a2a_info = a2a_info
-            ctx.a2a_info = a2a_info
-            return myreq.tensor
-
-    @staticmethod
-    def backward(ctx, *grad_output):
-        global myreq
-        with record_function("DLRM alltoall_req_bwd_single"):
-            a2a_info = ctx.a2a_info
-            myreq.req.wait()
-            myreq.req = None
-            grad_input = myreq.tensor
-            grad_inputs = grad_input.view([a2a_info.batch_size, -1]).split(
-                a2a_info.emb_dim, dim=1
-            )
-            grad_inputs = [gin.contiguous() for gin in grad_inputs]
-            myreq.tensor = None
-            return (None, *grad_inputs)
-
-
-class All2All_Wait(Function):
-    @staticmethod
-    def forward(ctx, *output):
-        global myreq
-        with record_function("DLRM alltoall_wait_fwd_single"):
-            a2a_info = myreq.a2a_info
-            ctx.a2a_info = a2a_info
-            myreq.req.wait()
-            myreq.req = None
-            myreq.tensor = None
-            table_split_lengths = (
-                a2a_info.table_split_lengths
-                if a2a_info.table_split_lengths
-                else a2a_info.local_table_num
-                * a2a_info.local_batch_num
-                * a2a_info.emb_dim
-            )
-            outputs = output[0].split(table_split_lengths)
-            outputs = tuple(
-                [out.view([a2a_info.local_batch_num, -1]) for out in outputs]
-            )
-            return outputs
-
-    @staticmethod
-    def backward(ctx, *grad_outputs):
-        global myreq
-        with record_function("DLRM alltoall_wait_bwd_single"):
-            a2a_info = ctx.a2a_info
-            grad_outputs = [gout.contiguous().view([-1]) for gout in grad_outputs]
-            grad_output = torch.cat(grad_outputs)
-            grad_input = grad_output.new_empty(
-                [a2a_info.batch_size * a2a_info.local_table_num * a2a_info.emb_dim]
-            )
-            req = dist.all_to_all_single(
-                grad_input,
-                grad_output,
-                a2a_info.batch_split_lengths,
-                a2a_info.table_split_lengths,
-                async_op=True,
-            )
-            myreq.req = req
-            myreq.tensor = grad_input
-            return (grad_output,)
-
-
-class AllGather(Function):
-    @staticmethod
-    def forward(ctx, input, global_lengths, dim=0):
-        if not isinstance(global_lengths, (list, tuple)):
-            global_lengths = [global_lengths] * my_size
-
-        assert len(global_lengths) == my_size
-        assert global_lengths[my_rank] == input.size(dim)
-        local_start = sum(global_lengths[:my_rank])
-
-        output_size = list(input.size())
-
-        ctx.dim = dim
-        ctx.local_start = local_start
-        ctx.local_length = global_lengths[my_rank]
-
-        input = input.contiguous()
-        if dim == 0:
-            out_len = sum(global_lengths)
-            output_size[dim] = out_len
-            output = input.new_empty(output_size)
-            gather_list = list(output.split(global_lengths, dim=0))
-        else:
-            gather_list = [torch.empty_like(input) for _ in range(my_size)]
-            gather_list = []
-            for length in global_lengths:
-                output_size[dim] = length
-                gather_list.append(input.new_empty(output_size))
-
-        dist.all_gather(gather_list, input)
-
-        if dim != 0:
-            output = torch.cat(gather_list, dim=dim)
-
-        return output
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        # print("Inside All2AllBackward")
-        dim = ctx.dim
-        start = ctx.local_start
-        length = ctx.local_length
-
-        grad_input = grad_output.narrow(dim, start, length)
-
-        return (grad_input, None, None)
-
-
-class All2AllInfo(object):
-    pass
-
-
-def alltoall(inputs, per_rank_table_splits):
-    global myreq
-    batch_size, emb_dim = inputs[0].size()
-    a2a_info = All2AllInfo()
-    a2a_info.local_table_num = len(inputs)
-    a2a_info.global_table_wise_parition_slices = per_rank_table_splits
-    (
-        a2a_info.local_batch_num,
-        a2a_info.global_batch_partition_slices,
-    ) = get_split_lengths(batch_size)
-    a2a_info.emb_dim = emb_dim
-    a2a_info.batch_size = batch_size
-    a2a_info.global_table_num = (
-        sum(per_rank_table_splits)
-        if per_rank_table_splits
-        else a2a_info.local_table_num * my_size
-    )
-
-    if a2a_impl == "" and alltoall_supported or a2a_impl == "alltoall":
-        # print("Using All2All_Req")
-        output = All2All_Req.apply(a2a_info, *inputs)
-        myreq.WaitFunction = All2All_Wait
-    elif a2a_impl == "" or a2a_impl == "scatter":
-        # print("Using All2All_Scatter_Req")
-        output = All2All_Scatter_Req.apply(a2a_info, *inputs)
-        myreq.WaitFunction = All2All_Scatter_Wait
-    elif a2a_impl == "scatter_list":
-        # print("Using All2All_ScatterList_Req")
-        output = All2All_ScatterList_Req.apply(a2a_info, *inputs)
-        myreq.WaitFunction = All2All_ScatterList_Wait
-    else:
-        print(
-            "Unknown value set for DLRM_ALLTOALL_IMPL (%s), "
-            "please use one of [alltoall, scatter, scatter_list]" % a2a_impl
-        )
-    return myreq
-
-
-def all_gather(input, lengths, dim=0):
-    global my_rank, my_size
-    if my_size == -1:
-        my_size = dist.get_world_size()
-        my_rank = dist.get_rank()
-    if not lengths:
-        lengths = [input.size(0)] * my_size
-    return AllGather.apply(input, lengths, dim)
-
-
-def barrier():
-    if my_size > 1:
-        dist.barrier()
-
-
-# Override builtin print function to print only from rank 0
-orig_print = builtins.print
-
-
-def rank0_print(*args, **kwargs):
-    if my_rank <= 0 or kwargs.get("print_all", False):
-        orig_print(*args, **kwargs)
-
-
-# builtins.print = rank0_print
-
-# Allow printing from all rank with explicit print_all
-def print_all(*args, **kwargs):
-    orig_print(*args, **kwargs)

src/third_party/MuQ/src/muq/muq_mulan/modules/transformer.py

@@ -1,185 +0,0 @@
-# attention
-from torch import einsum
-import torch.nn as nn
-import torch.nn.functional as F
-import torch
-from einops import rearrange
-from .utils import l2norm, default, exists
-
-# 2d sinusoidal positional embedding
-# simple vit paper shows it is good enough compared to learned
-
-def posemb_sincos_2d(patches, temperature = 10000, dtype = torch.float32):
-    _, h, w, dim, device, dtype = *patches.shape, patches.device, patches.dtype
-
-    y, x = torch.meshgrid(torch.arange(h, device = device), torch.arange(w, device = device), indexing = 'ij')
-    assert (dim % 4) == 0, 'feature dimension must be multiple of 4 for sincos emb'
-
-    omega = torch.arange(dim // 4, device = device) / (dim // 4 - 1)
-    omega = 1. / (temperature ** omega)
-
-    y = y.flatten()[:, None] * omega[None, :]
-    x = x.flatten()[:, None] * omega[None, :] 
-
-    pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim = 1)
-    pe = pe.type(dtype)
-
-    return rearrange(pe, '(h w) d -> h w d', h = h, w = w)
-
-# biasless layernorm
-
-class LayerNorm(nn.Module):
-    def __init__(self, dim, scale = True):
-        super().__init__()
-        self.learned_gamma = nn.Parameter(torch.ones(dim)) if scale else None
-
-        self.register_buffer('gamma', torch.ones(dim), persistent = False)
-        self.register_buffer('beta', torch.zeros(dim), persistent = False)
-
-    def forward(self, x):
-        return F.layer_norm(x, x.shape[-1:], default(self.learned_gamma, self.gamma), self.beta)
-
-# feedforward
-
-class GEGLU(nn.Module):
-    def forward(self, x):
-        x, gate = x.chunk(2, dim = -1)
-        return F.gelu(gate) * x
-
-def FeedForward(dim, mult = 4, dropout = 0.):
-    dim_hidden = int(dim * mult * 2 / 3)
-
-    return nn.Sequential(
-        LayerNorm(dim),
-        nn.Linear(dim, dim_hidden * 2, bias = False),
-        GEGLU(),
-        nn.Dropout(dropout),
-        nn.Linear(dim_hidden, dim, bias = False)
-    )
-
-class Attention(nn.Module):
-    def __init__(
-        self,
-        dim,
-        causal = False,
-        dim_head = 64,
-        heads = 8,
-        dropout = 0.,
-        scale = 8
-    ):
-        super().__init__()
-        self.heads = heads
-        self.scale = scale
-        self.causal = causal
-        inner_dim = dim_head * heads
-
-        self.norm = LayerNorm(dim)
-
-        self.attn_dropout = nn.Dropout(dropout)
-
-        self.to_q = nn.Linear(dim, inner_dim, bias = False)
-        self.to_kv = nn.Linear(dim, inner_dim * 2, bias = False)
-
-        self.q_scale = nn.Parameter(torch.ones(dim_head))
-        self.k_scale = nn.Parameter(torch.ones(dim_head))
-
-        self.to_out = nn.Sequential(
-            nn.Linear(inner_dim, dim, bias = False),
-            nn.Dropout(dropout)
-        )
-
-    def forward(
-        self,
-        x,
-        rel_pos_bias = None,
-        mask = None
-    ):
-        b, n, _, device = *x.shape, x.device
-
-        # prenorm
-
-        x = self.norm(x)
-
-        # project for queries, keys, values
-
-        q, k, v = self.to_q(x), *self.to_kv(x).chunk(2, dim = -1)
-
-        # split for multi-headed attention
-
-        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), (q, k, v))
-
-        # qk rmsnorm, technique circulating within brain used to stabilize a 22B parameter vision model training
-
-        q, k = map(l2norm, (q, k))
-        q = q * self.q_scale
-        k = k * self.k_scale
-
-        # similarities
-
-        sim = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
-
-        if exists(rel_pos_bias):
-            sim = sim + rel_pos_bias
-
-        if exists(mask):
-            mask = rearrange(mask, 'b j -> b 1 1 j')
-            sim = sim.masked_fill(~mask, -torch.finfo(sim.dtype).max)
-
-        if self.causal:
-            i, j = sim.shape[-2:]
-            causal_mask = torch.ones((i, j), dtype = torch.bool, device = x.device).triu(j - i + 1)
-            sim = sim.masked_fill(causal_mask, -torch.finfo(sim.dtype).max)
-
-        # attention
-
-        attn = sim.softmax(dim = -1)
-        attn = self.attn_dropout(attn)
-
-        # aggregate
-
-        out = einsum('b h i j, b h j d -> b h i d', attn, v)
-
-        # merge heads
-
-        out = rearrange(out, 'b h n d -> b n (h d)')
-        return self.to_out(out)
-
-# transformer
-
-class Transformer(nn.Module):
-    def __init__(
-        self,
-        dim,
-        depth,
-        dim_head = 64,
-        heads = 8,
-        attn_dropout = 0.,
-        ff_mult = 4,
-        ff_dropout = 0.
-    ):
-        super().__init__()
-        self.layers = nn.ModuleList([])
-        for _ in range(depth):
-            self.layers.append(nn.ModuleList([
-                Attention(dim = dim, dim_head = dim_head, heads = heads, dropout = attn_dropout),
-                FeedForward(dim = dim, mult = ff_mult, dropout = ff_dropout),
-            ]))
-
-    def forward(
-        self,
-        x,
-        rel_pos_bias = None,
-        mask = None,
-        return_all_layers = False
-    ):
-        layers = []
-
-        for attn, ff in self.layers:
-            x = attn(x, rel_pos_bias = rel_pos_bias, mask = mask) + x
-            x = ff(x) + x
-            layers.append(x)
-
-        if not return_all_layers:
-            return x
-
-        return x, torch.stack(layers[:-1]) if len(self.layers)>1 else None

src/third_party/MuQ/src/muq/muq_mulan/modules/utils.py

@@ -1,45 +0,0 @@
-from functools import wraps
-import torch
-from torch import nn
-import torch.nn.functional as F
-
-def exists(val):
-    return val is not None
-
-def first(it):
-    return it[0]
-
-def default(val, d):
-    return val if exists(val) else d
-
-def round_down_nearest_multiple(n, divisor):
-    return n // divisor * divisor
-
-def Sequential(*modules):
-    return nn.Sequential(*filter(exists, modules))
-
-
-def once(fn):
-    called = False
-    @wraps(fn)
-    def inner(x):
-        nonlocal called
-        if called:
-            return
-        called = True
-        return fn(x)
-    return inner
-
-print_once = once(print)
-
-# tensor functions
-
-def log(t, eps = 1e-20):
-    return torch.log(t.clamp(min = eps))
-
-def l2norm(t):
-    return F.normalize(t, p = 2, dim = -1)
-
-def frozen_params(model:nn.Module):
-    for param in model.parameters():
-        param.requires_grad = False

src/third_party/MuQ/src/muq/muq_mulan/muq_mulan.py

@@ -1,271 +0,0 @@
-from typing import List, Optional
-from dataclasses import dataclass, field
-import os
-
-from torch.nn.parallel.distributed import DistributedDataParallel
-import torch
-import torch.nn as nn
-from torch import einsum
-from einops import rearrange
-from huggingface_hub import PyTorchModelHubMixin
-from easydict import EasyDict
-
-from .models.mulan import MuLanModel
-from .models.audio import AudioSpectrogramTransformerPretrained
-from .models.text import TextTransformerPretrained
-from .modules.utils import exists, frozen_params
-
-
-@dataclass
-class MuLanConfig:
-    sr:int = field(default=24000)
-    clip_secs:float = field(default=10)
-    dim_latent:int = field(default=512)
-    decoupled_contrastive_learning:bool = field(default=True)
-    hierarchical_contrastive_loss:bool = field(default=False)
-    hierarchical_contrastive_loss_layers:Optional[List] = field(default=None)
-    sigmoid_contrastive_loss:bool = field(default=False)
-    rank_contrast:bool = field(default=True)
-
-@dataclass
-class AudioTransformerConfig:
-    dim:int = field(default=768)
-    tf_depth:int = field(default=8)
-    heads:int = field(default=8)
-    dim_head:int = field(default=64)
-    attn_dropout:float = field(default=0.)
-    ff_dropout:float = field(default=0.)
-    ff_mult:int = field(default=4)
-
-@dataclass
-class TextTransformerConfig:
-    dim:int = field(default=768)
-    tf_depth:int = field(default=8)
-    max_seq_len:int = field(default=1024)
-    dim_head:int = field(default=64)
-    heads:int = field(default=8)
-    attn_dropout:float = field(default=0.)
-    ff_dropout:float = field(default=0.)
-    ff_mult:int = field(default=4)
-
-@dataclass
-class ModalModelConfig:
-    name:str = field(default='')
-    model_dim: Optional[int] = field(default=None)
-    use_layer_idx: int = field(default=-1)
-
-
-@dataclass
-class MuQMuLanConfig:
-    mulan: MuLanConfig
-    audio_model: ModalModelConfig
-    text_model: ModalModelConfig
-    audio_transformer: AudioTransformerConfig
-    text_transformer: TextTransformerConfig
-
-class MuQMuLan(nn.Module, PyTorchModelHubMixin):
-    def __init__(self, config: MuQMuLanConfig, hf_hub_cache_dir=None):
-        super().__init__()
-        config = self._to_obj(config)
-        self.config = config
-        self.mulan = self.create_MuLan_from_config(config, hf_hub_cache_dir)
-        self.sr = config.mulan.sr
-        self.clip_secs = config.mulan.clip_secs
-    
-    def _to_obj(self, config):
-        if isinstance(config, MuQMuLanConfig):
-            config = EasyDict(
-                mulan = config.mulan,
-                audio_model = config.audio_model,
-                text_model = config.text_model,
-                audio_transformer = config.audio_transformer,
-                text_transformer = config.text_transformer,
-            )
-        else:
-            config = EasyDict(config)
-        return config
-    
-    @classmethod
-    def from_pretrained(cls, *args, cache_dir=None, **kwargs):
-        kwargs['hf_hub_cache_dir'] = cache_dir
-        return super().from_pretrained(*args, cache_dir=cache_dir, **kwargs)
-
-
-    @classmethod
-    def create_MuLan_from_config(cls, config:MuQMuLanConfig, hf_hub_cache_dir=None) -> MuLanModel:
-
-        audio_transformer = AudioSpectrogramTransformerPretrained(
-            model_name = config.audio_model.name, 
-            model_dim = config.audio_model.model_dim,
-            use_layer_idx = config.audio_model.use_layer_idx,
-            **config.audio_transformer,
-            frozen_pretrained = False,
-            hf_hub_cache_dir = hf_hub_cache_dir,
-        )
-        text_transformer = TextTransformerPretrained(
-            model_name = config.text_model.name, 
-            model_dim = config.text_model.model_dim,
-            **config.text_transformer,
-            frozen_pretrained = False,
-            hf_hub_cache_dir = hf_hub_cache_dir,
-        )
-
-        mulan = MuLanModel(
-            audio_transformer = audio_transformer,
-            text_transformer = text_transformer,
-            **config.mulan
-        )
-
-        return mulan
-    
-    def frozen(self):
-        frozen_params(self)
-
-    @property
-    def device(self):
-        return next(self.parameters()).device
-    
-    @property
-    def mulan_module(self):
-        if isinstance(self.mulan, DistributedDataParallel):
-            return self.mulan.module
-        else:
-            return self.mulan
-    
-    def forward(self,
-        wavs: Optional[torch.Tensor] = None,
-        texts: Optional[List[str]] = None,
-        *,
-        parallel_processing = False, 
-    ) -> torch.Tensor:    
-        """
-        Extract audio or text features, takes audio OR texts batch as input.
-        Note that if the audio is longer than 10s, it will be crop to multi cips and returns the average latent.
-        The param `parallel_processing` is used to control whether to use parallel processing or not.
-        If set to True, it uses parallel processing extractraction, which is faster but uses more GPU memory.
-        If set to False(the default), it uses serial processing extraction, which is slower but memory-friendly.
-
-        Args:
-            wavs (Optional[torch.Tensor]): Audio waveform tensor. Defaults to None.
-            texts (Optional[List[str]]): List of text strings. Defaults to None.
-            parallel_processing (bool): Whether to use parallel processing. Defaults to False.
-
-        Returns:
-            torch.Tensor: Latent representation of audio or text input.
-
-        Raises:
-            AssertionError: If both wavs and texts are provided or if neither is provided.
-
-        Note:
-            - Either wavs or texts must be provided, but not both.
-            - If wavs is provided, it calls extract_audio_latents method to process audio.
-            - If texts is provided, it calls extract_text_latents method to process text.
-        """
-        assert exists(wavs) ^ exists(texts), "Please provide either wavs or texts, but not both"
-        
-        if exists(wavs):
-            return self.extract_audio_latents(wavs = wavs, parallel_processing = parallel_processing)
-        else: 
-            return self.extract_text_latents(texts = texts)
-    
-    def calc_similarity(self, audio_latents: torch.Tensor, text_latents: torch.Tensor) -> torch.Tensor:
-        """
-        Calculate the dot-product similarity between audio and text latent representations.
-        It supports various dimensions of input tensors (with/without batch dimension) for both audio and text.
-        
-        Note:
-            The effect of this function is basically equivalent to the dot product.
-            mulan.calc_similarity(lat_a, lat_t) <==> einsum('i d, j d -> i j', lat_a, lat_t)
-
-        Args:
-            audio_latents (torch.Tensor): Latent representation of audio.
-            text_latents (torch.Tensor): Latent representation of text.
-
-        Returns:
-            torch.Tensor: Similarity scores between audio and text latent representations.
-
-        """
-        dim_a, dim_t = len(audio_latents.shape), len(text_latents.shape)
-        if dim_a == 2 and dim_t == 2:
-            return einsum('i d, j d -> i j', audio_latents, text_latents)
-        elif dim_a == 1 and dim_t == 1:
-            return torch.dot(audio_latents, text_latents)
-        elif dim_a == 2 and dim_t == 1:
-            return einsum('i d, d -> i', audio_latents, text_latents)
-        elif dim_a == 1 and dim_t == 2:
-            return einsum('d, j d -> j', audio_latents, text_latents)
-        
-        raise RuntimeError(f"Invalid dimensions: audio {dim_a}, text {dim_t}")
-        
-    
-    def extract_audio_latents(self, wavs:torch.Tensor, *, parallel_processing = False) -> torch.Tensor:
-        """
-        Extract latent representations from audio waveforms.
-
-        This function processes a batch of audio waveforms and extracts their latent representations.
-        It supports parallel processing for faster computation but uses more GPU memory.
-
-        Args:
-            wavs (torch.Tensor): A batch of audio waveform tensors.
-            parallel_processing (bool): Flag to enable parallel processing. Defaults to False.
-
-        Returns:
-            torch.Tensor: A tensor containing the latent representations of the input audio waveforms.
-        """
-        audio_latents = []
-
-        def audio_to_latent(wav):
-            return self.mulan_module.get_audio_latents(wav)
-        for wav in wavs:
-            wav_tensors = []
-            if isinstance(wav, torch.Tensor):
-                wav_tensors = self._get_all_clips(wav)
-            else: 
-                raise TypeError('wavs must be a Tensor')
-            
-            if parallel_processing:
-                wav_tensors = wav_tensors.to(self.device)
-                audio_latent = audio_to_latent(wav_tensors)
-                audio_latent = audio_latent.mean(dim=0)
-            else:  
-                wav_tensors = rearrange(wav_tensors, "i j -> i 1 j")
-                audio_latent = []
-                for wav_tensor in wav_tensors:
-                    audio_latent.append(audio_to_latent(wav_tensor).squeeze(0))
-                    del wav_tensor
-                audio_latent = torch.stack(audio_latent, dim=0)
-                audio_latent = audio_latent.mean(dim=0).to(self.device)      
-                
-            audio_latents.append(audio_latent)
-        audio_latents = torch.stack(audio_latents, dim=0)
-        return audio_latents
-
-    def extract_text_latents(self, texts: List[str]) -> torch.Tensor:
-        """
-        Extract latent representations from text inputs.
-
-        This function processes a list of text strings and extracts their latent representations
-        using the MuLan model's text tower.
-
-        Args:
-            texts (List[str]): A list of text strings to be processed.
-
-        Returns:
-            torch.Tensor: A tensor containing the latent representations of the input texts.
-        """
-        return self.mulan_module.get_text_latents(raw_texts=texts)
-    
-    def _get_all_clips(self, audio):
-        origin_length = len(audio)
-        accum_length = 0
-        delta = self.sr * self.clip_secs
-        audio_clips = []
-        while accum_length + delta <= origin_length:
-            clip = audio[accum_length:accum_length + delta]
-            audio_clips.append(clip)
-            accum_length += delta
-        if accum_length < origin_length:
-            audio_clips.append(torch.cat([audio[accum_length:], audio[0:delta - (origin_length - accum_length)]]))
-
-        return torch.stack(audio_clips, dim=0)
-

src/third_party/MuQ/src/recipes/contrastive_learning/README.md

@@ -1,107 +0,0 @@
-`
-# Guidance on MuQ-MuLan Training (Contrastive Learning)
-
-This guide provides instructions for training **MuQ-MuLan**, a contrastive learning model that jointly encodes music and text.
-
-We recommend training on **32 GPUs**, each with **at least 32 GB of memory**. Training typically takes **1–2 days**, depending on hardware environment.
-
----
-
-## Step 1: Environment Setup
-
-First, install all required dependencies listed in [requirements.txt](./requirements.txt):
-
-```bash
-pip install -r requirements.txt
-```
-
-Then, install this repository as a library in editable mode:
-
-```bash
-pip install -e .
-```
-
-
----
-
-## Step 2: Data Preparation
-
-We provide an example setup using the [MTG-Jamendo](https://github.com/MTG/mtg-jamendo-dataset) open-source music dataset.
-
-Please download our preprocessed data split files from [this link](https://drive.google.com/file/d/1PMCUpmtw8JwUv9Y-bNl8WIAZDtX8UaFY/view?usp=sharing) and place them in the appropriate directory.
-
-If you wish to use your own dataset, ensure that your data follow the required format.
-
----
-
-## Step 3: Run Training
-
-### Option 1: Manual Configuration
-
-This method will prompt you to configure distributed training manually via `accelerate`:
-
-```bash
-accelerate config
-accelerate launch train.py
-```
-
-### Option 2: Using a Predefined Config File
-
-We provide example config files for multi-node multi-GPU training. For example, to launch training on **4 nodes**, each with **8 GPUs**, run:
-
-```bash
-accelerate launch \
-  --config_file config/accelerate/32gpu4node_fp16.yaml \
-  --machine_rank $NODE_RANK \
-  --main_process_ip $CHIEF_IP \
-  --main_process_port 29500 \
-  train.py
-```
-
-* `$NODE_RANK`: Index of the current machine (0 for the main node).
-* `$CHIEF_IP`: IP address of the main (rank-0) node.
-* Adjust the config file as needed for different hardware setups.
-
-If you are training on a **single machine**, simply set `--main_process_ip=127.0.0.1` and `--machine_rank=0`.
-
-If you wish to use your own pretrained MuQ model for initialization, simply modify the `model.mulan.audio_model.name` field in `config/model/muq_mulan.yaml`.
-
----
-
-## Step 4: Convert to HF Checkpoint
-
-After training, convert your Fairseq-style checkpoint to HuggingFace format:
-
-```bash
-python scripts/convert_muqmulan_fairseq_ckpt_to_huggingface.py \
-  --checkpoint_path outputs/YYYY-MM-DD/hh-mm-ss/ckpt/mulan.1100.pt \
-  --save_dir outputs/hf-username/My-MuQ-MuLan-large
-```
-
-You can then load and use the model via the HuggingFace-style interface:
-
-```python
-from muq import MuQMuLan
-
-# Load from local HuggingFace-style checkpoint
-mulan = MuQMuLan.from_pretrained("outputs/hf-username/My-MuQ-MuLan-large")
-mulan = mulan.to(device).eval()
-
-# Extract music embeddings
-audio_embeds = mulan(wavs=wavs)
-
-# Extract text embeddings (in English or Chinese)
-text_embeds = mulan(texts=texts)
-
-# Compute similarity
-sim = mulan.calc_similarity(audio_embeds, text_embeds)
-```
-
-You can also upload your converted checkpoint to the Hugging Face Hub using `huggingface-cli`. The uploaded model will remain fully compatible with the `MuQMuLan.from_pretrained()` interface.
-
----
-
-## Evaluation
-
-For evaluation, we recommend using the [sota-music-tagging-models](https://github.com/minzwon/sota-music-tagging-models/) toolkit. It supports various metrics and datasets widely used in music tagging and retrieval.
-

src/third_party/MuQ/src/recipes/contrastive_learning/config/accelerate/2gpu.yaml

@@ -1,17 +0,0 @@
-compute_environment: LOCAL_MACHINE
-debug: false
-distributed_type: MULTI_GPU
-downcast_bf16: 'no'
-enable_cpu_affinity: false
-gpu_ids: all
-machine_rank: 0
-main_training_function: main
-mixed_precision: 'no'
-num_machines: 1
-num_processes: 2
-rdzv_backend: static
-same_network: true
-tpu_env: []
-tpu_use_cluster: false
-tpu_use_sudo: false
-use_cpu: false

src/third_party/MuQ/src/recipes/contrastive_learning/config/accelerate/32gpu4node.yaml

@@ -1,19 +0,0 @@
-compute_environment: LOCAL_MACHINE
-debug: false
-distributed_type: MULTI_GPU
-downcast_bf16: 'no'
-enable_cpu_affinity: false
-gpu_ids: all
-machine_rank: 0
-main_process_ip: '11.214.123.26'
-main_process_port: 25545
-main_training_function: main
-mixed_precision: 'no'
-num_machines: 4
-num_processes: 32
-rdzv_backend: static
-same_network: true
-tpu_env: []
-tpu_use_cluster: false
-tpu_use_sudo: false
-use_cpu: false

src/third_party/MuQ/src/recipes/contrastive_learning/config/accelerate/32gpu4node_fp16.yaml

@@ -1,19 +0,0 @@
-compute_environment: LOCAL_MACHINE
-debug: false
-distributed_type: MULTI_GPU
-downcast_bf16: 'no'
-enable_cpu_affinity: false
-gpu_ids: all
-machine_rank: 0
-main_process_ip: $INDEX
-main_process_port: 25529
-main_training_function: main
-mixed_precision: fp16
-num_machines: 4
-num_processes: 32
-rdzv_backend: static
-same_network: true
-tpu_env: []
-tpu_use_cluster: false
-tpu_use_sudo: false
-use_cpu: false

src/third_party/MuQ/src/recipes/contrastive_learning/config/accelerate/32gpu4node_zero2.yaml

@@ -1,25 +0,0 @@
-compute_environment: LOCAL_MACHINE
-debug: false
-deepspeed_config:
-  deepspeed_multinode_launcher: standard
-  gradient_accumulation_steps: 1
-  offload_optimizer_device: none
-  offload_param_device: none
-  zero3_init_flag: false
-  zero_stage: 2
-distributed_type: DEEPSPEED
-downcast_bf16: 'no'
-enable_cpu_affinity: false
-machine_rank: 0
-main_process_ip: $CHIEF_IP
-main_process_port: 25520
-main_training_function: main
-mixed_precision: 'no'
-num_machines: 4
-num_processes: 32
-rdzv_backend: static
-same_network: true
-tpu_env: []
-tpu_use_cluster: false
-tpu_use_sudo: false
-use_cpu: false

src/third_party/MuQ/src/recipes/contrastive_learning/config/accelerate/8gpu_fp16.yaml

@@ -1,17 +0,0 @@
-compute_environment: LOCAL_MACHINE
-debug: false
-distributed_type: MULTI_GPU
-downcast_bf16: 'no'
-enable_cpu_affinity: false
-gpu_ids: all
-machine_rank: 0
-main_training_function: main
-mixed_precision: fp16
-num_machines: 1
-num_processes: 8
-rdzv_backend: static
-same_network: true
-tpu_env: []
-tpu_use_cluster: false
-tpu_use_sudo: false
-use_cpu: false

src/third_party/MuQ/src/recipes/contrastive_learning/config/accelerate/8gpu_fp16_zero2.yaml

@@ -1,22 +0,0 @@
-compute_environment: LOCAL_MACHINE
-debug: false
-deepspeed_config:
-  gradient_accumulation_steps: 1
-  offload_optimizer_device: none
-  offload_param_device: none
-  zero3_init_flag: false
-  zero_stage: 2
-distributed_type: DEEPSPEED
-downcast_bf16: 'no'
-enable_cpu_affinity: false
-machine_rank: 0
-main_training_function: main
-mixed_precision: fp16
-num_machines: 1
-num_processes: 8
-rdzv_backend: static
-same_network: true
-tpu_env: []
-tpu_use_cluster: false
-tpu_use_sudo: false
-use_cpu: false