Delete qhash

qhash/autoencoder.py

@@ -1,26 +0,0 @@
-import math
-
-import torch
-import torchaudio
-from transformers.models.dac import DacModel
-
-
-class DACAutoencoder:
-    def __init__(self):
-        super().__init__()
-        self.dac = DacModel.from_pretrained("Quantamhash/dac_44khz")
-        self.dac.eval().requires_grad_(False)
-        self.codebook_size = self.dac.config.codebook_size
-        self.num_codebooks = self.dac.quantizer.n_codebooks
-        self.sampling_rate = self.dac.config.sampling_rate
-
-    def preprocess(self, wav: torch.Tensor, sr: int) -> torch.Tensor:
-        wav = torchaudio.functional.resample(wav, sr, 44_100)
-        right_pad = math.ceil(wav.shape[-1] / 512) * 512 - wav.shape[-1]
-        return torch.nn.functional.pad(wav, (0, right_pad))
-
-    def encode(self, wav: torch.Tensor) -> torch.Tensor:
-        return self.dac.encode(wav).audio_codes
-
-    def decode(self, codes: torch.Tensor) -> torch.Tensor:
-        return self.dac.decode(audio_codes=codes).audio_values.unsqueeze(1)

qhash/backbone.py

@@ -1,50 +0,0 @@
-import torch
-import torch.nn as nn
-from mamba_ssm.models.mixer_seq_simple import create_block
-from mamba_ssm.ops.triton.layer_norm import layer_norm_fn
-from mamba_ssm.utils.generation import InferenceParams
-
-from qhash.config import BackboneConfig
-
-
-class ZonosBackbone(nn.Module):
-    def __init__(self, config: BackboneConfig):
-        super().__init__()
-        self.config = config
-
-        self.layers = nn.ModuleList(
-            [
-                create_block(
-                    d_model=config.d_model,
-                    d_intermediate=config.d_intermediate
-                    if (i not in config.attn_layer_idx)
-                    else config.attn_mlp_d_intermediate,
-                    ssm_cfg=config.ssm_cfg,
-                    layer_idx=i,
-                    attn_layer_idx=config.attn_layer_idx,
-                    attn_cfg=config.attn_cfg,
-                    norm_epsilon=config.norm_epsilon,
-                    residual_in_fp32=config.residual_in_fp32,
-                    fused_add_norm=True,
-                    rms_norm=config.rms_norm,
-                )
-                for i in range(config.n_layer)
-            ]
-        )
-
-        self.norm_f = nn.LayerNorm(config.d_model, eps=config.norm_epsilon)
-
-    def forward(self, hidden_states: torch.Tensor, inference_params: InferenceParams | None = None):
-        residual = None
-        for layer in self.layers:
-            hidden_states, residual = layer(hidden_states, residual, inference_params)
-
-        return layer_norm_fn(
-            hidden_states,
-            self.norm_f.weight,
-            self.norm_f.bias,
-            residual,
-            eps=self.norm_f.eps,
-            residual_in_fp32=self.config.residual_in_fp32,
-            is_rms_norm=self.config.rms_norm,
-        )

qhash/codebook_pattern.py

@@ -1,12 +0,0 @@
-import torch
-import torch.nn.functional as F
-
-
-def apply_delay_pattern(codes: torch.Tensor, mask_token: int):
-    codes = F.pad(codes, (0, codes.shape[1]), value=mask_token)
-    return torch.stack([codes[:, k].roll(k + 1) for k in range(codes.shape[1])], dim=1)
-
-
-def revert_delay_pattern(codes: torch.Tensor):
-    _, n_q, seq_len = codes.shape
-    return torch.stack([codes[:, k, k + 1 : seq_len - n_q + k + 1] for k in range(n_q)], dim=1)

qhash/conditioning.py

@@ -1,373 +0,0 @@
-from functools import cache
-from typing import Any, Literal, Iterable
-
-import torch
-import torch.nn as nn
-
-from qhash.config import PrefixConditionerConfig
-
-
-class Conditioner(nn.Module):
-    def __init__(
-        self,
-        output_dim: int,
-        name: str,
-        cond_dim: int | None = None,
-        projection: Literal["none", "linear", "mlp"] = "none",
-        uncond_type: Literal["learned", "none"] = "none",
-        **kwargs,
-    ):
-        super().__init__()
-        self.name = name
-        self.output_dim = output_dim
-        self.cond_dim = cond_dim = cond_dim or output_dim
-
-        if projection == "linear":
-            self.project = nn.Linear(cond_dim, output_dim)
-        elif projection == "mlp":
-            self.project = nn.Sequential(
-                nn.Linear(cond_dim, output_dim),
-                nn.SiLU(),
-                nn.Linear(output_dim, output_dim),
-            )
-        else:
-            self.project = nn.Identity()
-
-        self.uncond_vector = None
-        if uncond_type == "learned":
-            self.uncond_vector = nn.Parameter(torch.zeros(output_dim))
-
-    def apply_cond(self, *inputs: Any) -> torch.Tensor:
-        raise NotImplementedError()
-
-    def forward(self, inputs: tuple[Any, ...] | None) -> torch.Tensor:
-        if inputs is None:
-            assert self.uncond_vector is not None
-            return self.uncond_vector.data.view(1, 1, -1)
-
-        cond = self.apply_cond(*inputs)
-        cond = self.project(cond)
-        return cond
-
-
-# ------- ESPEAK CONTAINMENT ZONE ------------------------------------------------------------------------------------------------------------------------------------------------
-import re
-import unicodedata
-
-import inflect
-import torch
-import torch.nn as nn
-from kanjize import number2kanji
-from phonemizer.backend import EspeakBackend
-from sudachipy import Dictionary, SplitMode
-
-# --- Number normalization code from https://github.com/daniilrobnikov/vits2/blob/main/text/normalize_numbers.py ---
-
-_inflect = inflect.engine()
-_comma_number_re = re.compile(r"([0-9][0-9\,]+[0-9])")
-_decimal_number_re = re.compile(r"([0-9]+\.[0-9]+)")
-_pounds_re = re.compile(r"£([0-9\,]*[0-9]+)")
-_dollars_re = re.compile(r"\$([0-9\.\,]*[0-9]+)")
-_ordinal_re = re.compile(r"[0-9]+(st|nd|rd|th)")
-_number_re = re.compile(r"[0-9]+")
-
-
-def _remove_commas(m: re.Match) -> str:
-    return m.group(1).replace(",", "")
-
-
-def _expand_decimal_point(m: re.Match) -> str:
-    return m.group(1).replace(".", " point ")
-
-
-def _expand_dollars(m: re.Match) -> str:
-    match = m.group(1)
-    parts = match.split(".")
-    if len(parts) > 2:
-        return match + " dollars"  # Unexpected format
-    dollars = int(parts[0]) if parts[0] else 0
-    cents = int(parts[1]) if len(parts) > 1 and parts[1] else 0
-    if dollars and cents:
-        dollar_unit = "dollar" if dollars == 1 else "dollars"
-        cent_unit = "cent" if cents == 1 else "cents"
-        return "%s %s, %s %s" % (dollars, dollar_unit, cents, cent_unit)
-    elif dollars:
-        dollar_unit = "dollar" if dollars == 1 else "dollars"
-        return "%s %s" % (dollars, dollar_unit)
-    elif cents:
-        cent_unit = "cent" if cents == 1 else "cents"
-        return "%s %s" % (cents, cent_unit)
-    else:
-        return "zero dollars"
-
-
-def _expand_ordinal(m: re.Match) -> str:
-    return _inflect.number_to_words(m.group(0))
-
-
-def _expand_number(m: re.Match) -> str:
-    num = int(m.group(0))
-    if num > 1000 and num < 3000:
-        if num == 2000:
-            return "two thousand"
-        elif num > 2000 and num < 2010:
-            return "two thousand " + _inflect.number_to_words(num % 100)
-        elif num % 100 == 0:
-            return _inflect.number_to_words(num // 100) + " hundred"
-        else:
-            return _inflect.number_to_words(num, andword="", zero="oh", group=2).replace(", ", " ")
-    else:
-        return _inflect.number_to_words(num, andword="")
-
-
-def normalize_numbers(text: str) -> str:
-    text = re.sub(_comma_number_re, _remove_commas, text)
-    text = re.sub(_pounds_re, r"\1 pounds", text)
-    text = re.sub(_dollars_re, _expand_dollars, text)
-    text = re.sub(_decimal_number_re, _expand_decimal_point, text)
-    text = re.sub(_ordinal_re, _expand_ordinal, text)
-    text = re.sub(_number_re, _expand_number, text)
-    return text
-
-
-# --- Number normalization code end ---
-
-
-PAD_ID, UNK_ID, BOS_ID, EOS_ID = 0, 1, 2, 3
-SPECIAL_TOKEN_IDS = [PAD_ID, UNK_ID, BOS_ID, EOS_ID]
-
-_punctuation = ';:,.!?¡¿—…"«»“”() *~-/\\&'
-_letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
-_letters_ipa = (
-    "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
-)
-
-symbols = [*_punctuation, *_letters, *_letters_ipa]
-_symbol_to_id = {s: i for i, s in enumerate(symbols, start=len(SPECIAL_TOKEN_IDS))}
-
-
-def _get_symbol_id(s: str) -> int:
-    return _symbol_to_id.get(s, 1)
-
-
-def get_symbol_ids(text: str) -> list[int]:
-    return list(map(_get_symbol_id, text))
-
-
-def tokenize_phonemes(phonemes: list[str]) -> tuple[torch.Tensor, list[int]]:
-    phoneme_ids = [[BOS_ID, *get_symbol_ids(phonemes), EOS_ID] for phonemes in phonemes]
-    lengths = list(map(len, phoneme_ids))
-    longest = max(lengths)
-    phoneme_ids = [[PAD_ID] * (longest - len(ids)) + ids for ids in phoneme_ids]
-    return torch.tensor(phoneme_ids), lengths
-
-
-def normalize_jp_text(text: str, tokenizer=Dictionary(dict="full").create()) -> str:
-    text = unicodedata.normalize("NFKC", text)
-    text = re.sub(r"\d+", lambda m: number2kanji(int(m[0])), text)
-    final_text = " ".join([x.reading_form() for x in tokenizer.tokenize(text, SplitMode.A)])
-    return final_text
-
-
-def clean(texts: list[str], languages: list[str]) -> list[str]:
-    texts_out = []
-    for text, language in zip(texts, languages):
-        if "ja" in language:
-            text = normalize_jp_text(text)
-        else:
-            text = normalize_numbers(text)
-        texts_out.append(text)
-    return texts_out
-
-
-@cache
-def get_backend(language: str) -> "EspeakBackend":
-    import logging
-
-    from phonemizer.backend import EspeakBackend
-
-    logger = logging.getLogger("phonemizer")
-    backend = EspeakBackend(
-        language,
-        preserve_punctuation=True,
-        with_stress=True,
-        punctuation_marks=_punctuation,
-        logger=logger,
-    )
-    logger.setLevel(logging.ERROR)
-    return backend
-
-
-def phonemize(texts: list[str], languages: list[str]) -> list[str]:
-    texts = clean(texts, languages)
-
-    batch_phonemes = []
-    for text, language in zip(texts, languages):
-        backend = get_backend(language)
-        phonemes = backend.phonemize([text], strip=True)
-        batch_phonemes.append(phonemes[0])
-
-    return batch_phonemes
-
-
-class EspeakPhonemeConditioner(Conditioner):
-    def __init__(self, output_dim: int, **kwargs):
-        super().__init__(output_dim, **kwargs)
-        self.phoneme_embedder = nn.Embedding(len(SPECIAL_TOKEN_IDS) + len(symbols), output_dim)
-
-    def apply_cond(self, texts: list[str], languages: list[str]) -> torch.Tensor:
-        """
-        Args:
-            texts: list of texts to convert to phonemes
-            languages: ISO 639-1 -or otherwise eSpeak compatible- language code
-        """
-        device = self.phoneme_embedder.weight.device
-
-        phonemes = phonemize(texts, languages)
-        phoneme_ids, _ = tokenize_phonemes(phonemes)
-        phoneme_embeds = self.phoneme_embedder(phoneme_ids.to(device))
-
-        return phoneme_embeds
-
-
-# ------- ESPEAK CONTAINMENT ZONE ------------------------------------------------------------------------------------------------------------------------------------------------
-
-
-class FourierConditioner(Conditioner):
-    def __init__(
-        self,
-        output_dim: int,
-        input_dim: int = 1,
-        std: float = 1.0,
-        min_val: float = 0.0,
-        max_val: float = 1.0,
-        **kwargs,
-    ):
-        assert output_dim % 2 == 0
-        super().__init__(output_dim, **kwargs)
-        self.register_buffer("weight", torch.randn([output_dim // 2, input_dim]) * std)
-        self.input_dim, self.min_val, self.max_val = input_dim, min_val, max_val
-
-    def apply_cond(self, x: torch.Tensor) -> torch.Tensor:
-        assert x.shape[-1] == self.input_dim
-        x = (x - self.min_val) / (self.max_val - self.min_val)  # [batch_size, seq_len, input_dim]
-        f = 2 * torch.pi * x.to(self.weight.dtype) @ self.weight.T  # [batch_size, seq_len, output_dim // 2]
-        return torch.cat([f.cos(), f.sin()], dim=-1)  # [batch_size, seq_len, output_dim]
-
-
-class IntegerConditioner(Conditioner):
-    def __init__(self, output_dim: int, min_val: int = 0, max_val: int = 512, **kwargs):
-        super().__init__(output_dim, **kwargs)
-        self.min_val = min_val
-        self.max_val = max_val
-        self.int_embedder = nn.Embedding(max_val - min_val + 1, output_dim)
-
-    def apply_cond(self, x: torch.Tensor) -> torch.Tensor:
-        assert x.shape[-1] == 1
-        return self.int_embedder(x.squeeze(-1) - self.min_val)  # [batch_size, seq_len, output_dim]
-
-
-class PassthroughConditioner(Conditioner):
-    def __init__(self, output_dim: int, **kwargs):
-        super().__init__(output_dim, **kwargs)
-
-    def apply_cond(self, x: torch.Tensor) -> torch.Tensor:
-        assert x.shape[-1] == self.cond_dim
-        return x
-
-
-_cond_cls_map = {
-    "PassthroughConditioner": PassthroughConditioner,
-    "EspeakPhonemeConditioner": EspeakPhonemeConditioner,
-    "FourierConditioner": FourierConditioner,
-    "IntegerConditioner": IntegerConditioner,
-}
-
-
-def build_conditioners(conditioners: list[dict], output_dim: int) -> list[Conditioner]:
-    return [_cond_cls_map[config["type"]](output_dim, **config) for config in conditioners]
-
-
-class PrefixConditioner(Conditioner):
-    def __init__(self, config: PrefixConditionerConfig, output_dim: int):
-        super().__init__(output_dim, "prefix", projection=config.projection)
-        self.conditioners = nn.ModuleList(build_conditioners(config.conditioners, output_dim))
-        self.norm = nn.LayerNorm(output_dim)
-        self.required_keys = {c.name for c in self.conditioners if c.uncond_vector is None}
-
-    def forward(self, cond_dict: dict) -> torch.Tensor:
-        if not set(cond_dict).issuperset(self.required_keys):
-            raise ValueError(f"Missing required keys: {self.required_keys - set(cond_dict)}")
-        conds = []
-        for conditioner in self.conditioners:
-            conds.append(conditioner(cond_dict.get(conditioner.name)))
-        max_bsz = max(map(len, conds))
-        assert all(c.shape[0] in (max_bsz, 1) for c in conds)
-        conds = [c.expand(max_bsz, -1, -1) for c in conds]
-        return self.norm(self.project(torch.cat(conds, dim=-2)))
-
-
-supported_language_codes = [
-    'af', 'am', 'an', 'ar', 'as', 'az', 'ba', 'bg', 'bn', 'bpy', 'bs', 'ca', 'cmn',
-    'cs', 'cy', 'da', 'de', 'el', 'en-029', 'en-gb', 'en-gb-scotland', 'en-gb-x-gbclan',
-    'en-gb-x-gbcwmd', 'en-gb-x-rp', 'en-us', 'eo', 'es', 'es-419', 'et', 'eu', 'fa',
-    'fa-latn', 'fi', 'fr-be', 'fr-ch', 'fr-fr', 'ga', 'gd', 'gn', 'grc', 'gu', 'hak',
-    'hi', 'hr', 'ht', 'hu', 'hy', 'hyw', 'ia', 'id', 'is', 'it', 'ja', 'jbo', 'ka',
-    'kk', 'kl', 'kn', 'ko', 'kok', 'ku', 'ky', 'la', 'lfn', 'lt', 'lv', 'mi', 'mk',
-    'ml', 'mr', 'ms', 'mt', 'my', 'nb', 'nci', 'ne', 'nl', 'om', 'or', 'pa', 'pap',
-    'pl', 'pt', 'pt-br', 'py', 'quc', 'ro', 'ru', 'ru-lv', 'sd', 'shn', 'si', 'sk',
-    'sl', 'sq', 'sr', 'sv', 'sw', 'ta', 'te', 'tn', 'tr', 'tt', 'ur', 'uz', 'vi',
-    'vi-vn-x-central', 'vi-vn-x-south', 'yue'
-]  # fmt: off
-
-
-def make_cond_dict(
-    text: str = "It would be nice to have time for testing, indeed.",
-    language: str = "en-us",
-    speaker: torch.Tensor | None = None,
-    emotion: list[float] = [0.3077, 0.0256, 0.0256, 0.0256, 0.0256, 0.0256, 0.2564, 0.3077],
-    fmax: float = 22050.0,
-    pitch_std: float = 20.0,
-    speaking_rate: float = 15.0,
-    vqscore_8: list[float] = [0.78] * 8,
-    ctc_loss: float = 0.0,
-    dnsmos_ovrl: float = 4.0,
-    speaker_noised: bool = False,
-    unconditional_keys: Iterable[str] = {"vqscore_8", "dnsmos_ovrl"},
-    device: str = "cuda",
-) -> dict:
-    """
-    A helper to build the 'cond_dict' that the model expects.
-    By default, it will generate a random speaker embedding
-    """
-    assert language.lower() in supported_language_codes, "Please pick a supported language"
-
-    language_code_to_id = {lang: i for i, lang in enumerate(supported_language_codes)}
-
-    cond_dict = {
-        "espeak": ([text], [language]),
-        "speaker": speaker,
-        "emotion": emotion,
-        "fmax": fmax,
-        "pitch_std": pitch_std,
-        "speaking_rate": speaking_rate,
-        "language_id": language_code_to_id[language],
-        "vqscore_8": vqscore_8,
-        "ctc_loss": ctc_loss,
-        "dnsmos_ovrl": dnsmos_ovrl,
-        "speaker_noised": int(speaker_noised),
-    }
-
-    for k in unconditional_keys:
-        cond_dict.pop(k, None)
-
-    for k, v in cond_dict.items():
-        if isinstance(v, (float, int, list)):
-            v = torch.tensor(v)
-        if isinstance(v, torch.Tensor):
-            cond_dict[k] = v.view(1, 1, -1).to(device)
-
-        if k == "emotion":
-            cond_dict[k] /= cond_dict[k].sum(dim=-1)
-
-    return cond_dict

qhash/config.py

@@ -1,38 +0,0 @@
-from dataclasses import dataclass, field
-from typing import Literal
-
-
-@dataclass
-class BackboneConfig:
-    d_model: int = 1024
-    d_intermediate: int = 0
-    attn_mlp_d_intermediate: int = 0
-    n_layer: int = 16
-    ssm_cfg: dict = field(default_factory=dict)
-    attn_layer_idx: list = field(default_factory=list)
-    attn_cfg: dict = field(default_factory=dict)
-    rms_norm: bool = False
-    residual_in_fp32: bool = False
-    norm_epsilon: float = 1e-5
-
-
-@dataclass
-class PrefixConditionerConfig:
-    conditioners: list[dict]
-    projection: Literal["none", "linear", "mlp"]
-
-
-@dataclass
-class ZonosConfig:
-    backbone: BackboneConfig
-    prefix_conditioner: PrefixConditionerConfig
-    eos_token_id: int = 1024
-    masked_token_id: int = 1025
-
-    @classmethod
-    def from_dict(cls, d: dict) -> "ZonosConfig":
-        d = d.copy()
-        backbone_config = BackboneConfig(**d.pop("backbone"))
-        prefix_conditioner_config = PrefixConditionerConfig(**d.pop("prefix_conditioner"))
-        config = cls(backbone_config, prefix_conditioner_config, **d)
-        return config

qhash/model.py

@@ -1,270 +0,0 @@
-import json
-from typing import Callable
-
-import safetensors
-import torch
-import torch.nn as nn
-from huggingface_hub import hf_hub_download
-from mamba_ssm.utils.generation import InferenceParams
-from tqdm import tqdm
-
-from qhash.backbone import ZonosBackbone
-from qhash.autoencoder import DACAutoencoder
-from qhash.codebook_pattern import apply_delay_pattern, revert_delay_pattern
-from qhash.conditioning import PrefixConditioner
-from qhash.config import ZonosConfig
-from qhash.sampling import sample_from_logits
-from qhash.speaker_cloning import SpeakerEmbeddingLDA
-
-
-class Zonos(nn.Module):
-    def __init__(self, config: ZonosConfig):
-        super().__init__()
-        self.config = config
-        dim = config.backbone.d_model
-        self.eos_token_id = config.eos_token_id
-        self.masked_token_id = config.masked_token_id
-
-        self.autoencoder = DACAutoencoder()
-        self.backbone = ZonosBackbone(config.backbone)
-        self.prefix_conditioner = PrefixConditioner(config.prefix_conditioner, dim)
-        self.spk_clone_model = None
-
-        # TODO: pad to multiple of at least 8
-        self.embeddings = nn.ModuleList([nn.Embedding(1026, dim) for _ in range(self.autoencoder.num_codebooks)])
-        self.heads = nn.ModuleList([nn.Linear(dim, 1025, bias=False) for _ in range(self.autoencoder.num_codebooks)])
-
-        self._cg_graph = None
-        self._cg_batch_size = None
-        self._cg_input_ids = None
-        self._cg_logits = None
-        self._cg_inference_params = None
-        self._cg_scale = None
-
-    @classmethod
-    def from_pretrained(cls, repo_id: str, revision: str | None = None, device: str = "cuda") -> "Zonos":
-        config_path = hf_hub_download(repo_id=repo_id, filename="config.json", revision=revision)
-        model_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors", revision=revision)
-        return cls.from_local(config_path, model_path, device)
-
-    @classmethod
-    def from_local(cls, config_path: str, model_path: str, device: str = "cuda") -> "Zonos":
-        config = ZonosConfig.from_dict(json.load(open(config_path)))
-        model = cls(config).to(device, torch.bfloat16)
-        model.autoencoder.dac.to(device)
-
-        sd = model.state_dict()
-        with safetensors.safe_open(model_path, framework="pt") as f:
-            for k in f.keys():
-                sd[k] = f.get_tensor(k)
-        model.load_state_dict(sd)
-
-        return model
-
-    def make_speaker_embedding(self, wav: torch.Tensor, sr: int) -> torch.Tensor:
-        """Generate a speaker embedding from an audio clip."""
-        if self.spk_clone_model is None:
-            self.spk_clone_model = SpeakerEmbeddingLDA()
-        _, spk_embedding = self.spk_clone_model(wav.to(self.spk_clone_model.device), sr)
-        return spk_embedding.unsqueeze(0).bfloat16()
-
-    def embed_codes(self, codes: torch.Tensor) -> torch.Tensor:
-        return sum(emb(codes[:, i]) for i, emb in enumerate(self.embeddings))
-
-    def apply_heads(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        return torch.stack([head(hidden_states) for head in self.heads], dim=1)
-
-    def _compute_logits(
-        self, hidden_states: torch.Tensor, inference_params: InferenceParams, cfg_scale: float
-    ) -> torch.Tensor:
-        """
-        Pass `hidden_states` into `backbone` and `multi_head`, applying
-        classifier-free guidance if `cfg_scale != 1.0`.
-        """
-        last_hidden_states = self.backbone(hidden_states, inference_params)[:, -1, :].unsqueeze(1)
-        logits = self.apply_heads(last_hidden_states).squeeze(2).float()
-        if cfg_scale != 1.0:
-            cond_logits, uncond_logits = logits.chunk(2)
-            logits = uncond_logits + (cond_logits - uncond_logits) * cfg_scale
-        return logits
-
-    def _decode_one_token(
-        self,
-        input_ids: torch.Tensor,
-        inference_params: InferenceParams,
-        cfg_scale: float,
-    ) -> torch.Tensor:
-        """
-        Single-step decode. Prepares the hidden states, possibly replicates them
-        for CFG, and then delegates to `_compute_logits`.
-
-        Below we wrap this function with a simple CUDA Graph capturing mechanism,
-        doing 3 warmup steps if needed and then capturing or replaying the graph.
-        We only recapture if the batch size changes.
-        """
-        # TODO: support cfg_scale==1
-        if cfg_scale == 1.0:
-            hidden_states = self.embed_codes(input_ids)
-            return self._compute_logits(hidden_states, inference_params, cfg_scale)
-
-        bsz = input_ids.size(0)
-
-        need_capture = (self._cg_graph is None) or (self._cg_batch_size != bsz)
-
-        if need_capture:
-            self._cg_graph = None
-
-            self._cg_batch_size = bsz
-            self._cg_inference_params = inference_params
-            self._cg_scale = cfg_scale
-
-            for _ in range(3):
-                hidden_states = self.embed_codes(input_ids)
-                hidden_states = hidden_states.repeat(2, 1, 1)  # because cfg != 1.0
-                logits = self._compute_logits(hidden_states, inference_params, cfg_scale)
-
-            self._cg_input_ids = input_ids.clone()
-            self._cg_logits = torch.empty_like(logits)
-
-            g = torch.cuda.CUDAGraph()
-
-            def capture_region():
-                hidden_states_local = self.embed_codes(self._cg_input_ids)
-                hidden_states_local = hidden_states_local.repeat(2, 1, 1)
-                self._cg_logits = self._compute_logits(hidden_states_local, self._cg_inference_params, self._cg_scale)
-
-            with torch.cuda.graph(g):
-                capture_region()
-
-            self._cg_graph = g
-
-        else:
-            self._cg_input_ids.copy_(input_ids)
-
-        self._cg_graph.replay()
-
-        return self._cg_logits
-
-    def _prefill(
-        self,
-        prefix_hidden_states: torch.Tensor,
-        input_ids: torch.Tensor,
-        inference_params: InferenceParams,
-        cfg_scale: float,
-    ) -> torch.Tensor:
-        """
-        "Prefill" mode: we already have `prefix_hidden_states`, and we want
-        to append new embeddings, then compute the logits.
-        """
-        # Replicate input_ids if CFG is enabled
-        if cfg_scale != 1.0:
-            input_ids = input_ids.expand(prefix_hidden_states.shape[0], -1, -1)
-        hidden_states = torch.cat([prefix_hidden_states, self.embed_codes(input_ids)], dim=1)
-        return self._compute_logits(hidden_states, inference_params, cfg_scale)
-
-    def setup_cache(self, batch_size: int, max_seqlen: int, dtype: torch.dtype = torch.bfloat16) -> InferenceParams:
-        key_value_memory_dict = {
-            i: layer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype)
-            for i, layer in enumerate(self.backbone.layers)
-        }
-        lengths_per_sample = torch.full((batch_size,), 0, dtype=torch.int32, device="cuda")
-        return InferenceParams(max_seqlen, batch_size, 0, 0, key_value_memory_dict, lengths_per_sample)
-
-    def prepare_conditioning(self, cond_dict: dict, uncond_dict: dict | None = None) -> torch.Tensor:
-        if uncond_dict is None:
-            uncond_dict = {k: cond_dict[k] for k in self.prefix_conditioner.required_keys}
-        return torch.cat(
-            [
-                self.prefix_conditioner(cond_dict),
-                self.prefix_conditioner(uncond_dict),
-            ]
-        )
-
-    @torch.inference_mode()
-    def generate(
-        self,
-        prefix_conditioning: torch.Tensor,  # [bsz, cond_seq_len, d_model]
-        audio_prefix_codes: torch.Tensor | None = None,  # [bsz, 9, prefix_audio_seq_len]
-        max_new_tokens: int = 86 * 30,
-        cfg_scale: float = 2.0,
-        batch_size: int = 1,
-        sampling_params: dict = dict(min_p=0.1),
-        progress_bar: bool = True,
-        callback: Callable[[torch.Tensor, int, int], bool] | None = None,
-    ):
-        assert cfg_scale != 1, "TODO: add support for cfg_scale=1"
-        prefix_audio_len = 0 if audio_prefix_codes is None else audio_prefix_codes.shape[2]
-
-        unknown_token = -1
-        audio_seq_len = prefix_audio_len + max_new_tokens
-        seq_len = prefix_conditioning.shape[1] + audio_seq_len
-
-        inference_params = self.setup_cache(batch_size=batch_size * 2, max_seqlen=seq_len)
-
-        codes = torch.full((batch_size, 9, audio_seq_len), unknown_token, device="cuda")
-        if audio_prefix_codes is not None:
-            codes[..., :prefix_audio_len] = audio_prefix_codes
-
-        delayed_codes = apply_delay_pattern(codes, self.masked_token_id)
-
-        delayed_prefix_audio_codes = delayed_codes[..., : prefix_audio_len + 1]
-
-        logits = self._prefill(prefix_conditioning, delayed_prefix_audio_codes, inference_params, cfg_scale)
-        next_token = sample_from_logits(logits, **sampling_params)
-
-        offset = delayed_prefix_audio_codes.shape[2]
-        frame = delayed_codes[..., offset : offset + 1]
-        frame.masked_scatter_(frame == unknown_token, next_token)
-
-        prefix_length = prefix_conditioning.shape[1] + prefix_audio_len + 1
-        inference_params.seqlen_offset += prefix_length
-        inference_params.lengths_per_sample[:] += prefix_length
-
-        logit_bias = torch.zeros_like(logits)
-        logit_bias[:, 1:, self.eos_token_id] = -torch.inf  # only allow codebook 0 to predict EOS
-
-        stopping = torch.zeros(batch_size, dtype=torch.bool, device="cuda")
-        max_steps = delayed_codes.shape[2] - offset
-        remaining_steps = torch.full((batch_size,), max_steps, device="cuda")
-        progress = tqdm(total=max_steps, desc="Generating", disable=not progress_bar)
-
-        step = 0
-        while torch.max(remaining_steps) > 0:
-            offset += 1
-            input_ids = delayed_codes[..., offset - 1 : offset]
-            logits = self._decode_one_token(input_ids, inference_params, cfg_scale)
-
-            next_token = sample_from_logits(logits, generated_tokens=delayed_codes[..., :offset], **sampling_params)
-            eos_in_cb0 = next_token[:, 0] == self.eos_token_id
-
-            remaining_steps[eos_in_cb0[:, 0]] = torch.minimum(remaining_steps[eos_in_cb0[:, 0]], torch.tensor(9))
-            stopping |= eos_in_cb0[:, 0]
-
-            eos_codebook_idx = 9 - remaining_steps
-            eos_codebook_idx = torch.clamp(eos_codebook_idx, max=9 - 1)
-            for i in range(next_token.shape[0]):
-                if stopping[i]:
-                    idx = eos_codebook_idx[i].item()
-                    next_token[i, :idx] = self.masked_token_id
-                    next_token[i, idx] = self.eos_token_id
-
-            frame = delayed_codes[..., offset : offset + 1]
-            frame.masked_scatter_(frame == unknown_token, next_token)
-            inference_params.seqlen_offset += 1
-            inference_params.lengths_per_sample[:] += 1
-
-            remaining_steps -= 1
-
-            progress.update()
-            step += 1
-
-            if callback is not None and not callback(frame, step, max_steps):
-                break
-
-        out_codes = revert_delay_pattern(delayed_codes)
-        out_codes.masked_fill_(out_codes >= 1024, 0)
-        out_codes = out_codes[..., : offset - 9]
-
-        self._cg_graph = None  # reset cuda graph to avoid cache changes
-
-        return out_codes

qhash/sampling.py

@@ -1,141 +0,0 @@
-import torch
-
-
-def multinomial(input: torch.Tensor, num_samples: int, replacement=False, *, generator=None):
-    """torch.multinomial with arbitrary number of dimensions, and number of candidates on the last dimension.
-
-    Args:
-        input (torch.Tensor): The input tensor containing probabilities.
-        num_samples (int): Number of samples to draw.
-        replacement (bool): Whether to draw with replacement or not.
-    Keywords args:
-        generator (torch.Generator): A pseudorandom number generator for sampling.
-    Returns:
-        torch.Tensor: Last dimension contains num_samples indices
-            sampled from the multinomial probability distribution
-            located in the last dimension of tensor input.
-    """
-
-    if num_samples == 1:
-        q = torch.empty_like(input).exponential_(1, generator=generator)
-        return torch.argmax(input / q, dim=-1, keepdim=True).to(torch.int64)
-
-    input_ = input.reshape(-1, input.shape[-1])
-    output_ = torch.multinomial(input_, num_samples=num_samples, replacement=replacement, generator=generator)
-    output = output_.reshape(*list(input.shape[:-1]), -1)
-    return output
-
-
-def apply_top_k(
-    probs: torch.Tensor,
-    k: int,
-) -> torch.Tensor:
-    """Sample next token from top K values along the last dimension of the input probs tensor.
-
-    Args:
-        probs (torch.Tensor): Input probabilities with token candidates on the last dimension.
-        k (int): The k in “top-k”.
-    Returns:
-        torch.Tensor: Sampled tokens.
-    """
-    v, _ = torch.topk(probs, min(k, probs.size(-1)))
-    pivot = v.select(-1, -1).unsqueeze(-1)
-    probs = torch.where(probs < pivot, 0.0, probs)
-    probs.div_(probs.sum(dim=-1, keepdim=True))
-    return probs
-
-
-def apply_top_p(probs: torch.Tensor, p: float) -> torch.Tensor:
-    """Sample next token from top P probabilities along the last dimension of the input probs tensor.
-
-    Args:
-        probs (torch.Tensor): Input probabilities with token candidates on the last dimension.
-        p (int): The p in “top-p”.
-    Returns:
-        torch.Tensor: Sampled tokens.
-    """
-    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
-    probs_sum = torch.cumsum(probs_sort, dim=-1)
-    mask = probs_sum - probs_sort > p
-    probs_sort *= (~mask).float()
-    probs = probs.scatter(-1, probs_idx, probs_sort)
-    probs.div_(probs.sum(dim=-1, keepdim=True))
-    return probs
-
-
-def apply_min_p(probs: torch.Tensor, min_p: float) -> torch.Tensor:
-    """Sample next token using min-p sampling.
-
-    Args:
-        scores (torch.FloatTensor): Input logits with token candidates on the last dimension.
-        min_p (float): Minimum token probability, scaled by the probability of the most likely token.
-                       Must be between 0 and 1. Typical values are in the 0.01-0.2 range.
-    Returns:
-        torch.Tensor: Sampled tokens.
-    """
-    top_probs, _ = probs.max(dim=-1, keepdim=True)
-    tokens_to_remove = probs < (min_p * top_probs)
-    probs = probs.masked_fill(tokens_to_remove, 0.0)
-    probs.div_(probs.sum(dim=-1, keepdim=True))
-    return probs
-
-
-def modify_logit_for_repetition_penalty(
-    logits: torch.Tensor,
-    generated_tokens: torch.Tensor,
-    repetition_penalty: float,
-    repetition_penalty_window: int,
-):
-    """See https://arxiv.org/abs/1909.05858
-    Apply repetition penalty over a sliding window of the last `repetition_penalty_window` tokens.
-    logits: (batch_size, n_codebooks, vocab_size)
-    generated_tokens: (batch_size, n_codebooks, seq_len)
-    """
-    generated_tokens = generated_tokens[..., -repetition_penalty_window:]
-    generated_tokens = generated_tokens.clamp_max(logits.shape[-1] - 1).to(torch.int64)
-    rp = torch.full_like(logits, repetition_penalty)
-    factors = torch.ones_like(logits).scatter_reduce(2, generated_tokens, rp, reduce="prod")
-    return torch.where(logits <= 0, logits * factors, logits / factors)
-
-
-def sample_from_logits(
-    logits: torch.Tensor,
-    temperature: float = 1.0,
-    top_p: float = 0.0,
-    top_k: int = 0,
-    min_p: float = 0.0,
-    generated_tokens: torch.Tensor | None = None,
-    repetition_penalty: float = 3.0,
-    repetition_penalty_window: float = 2,
-) -> torch.Tensor:
-    """Sample next token from logits using temperature, top-p, top-k, or min-p sampling.
-
-    Args:
-        logits (torch.Tensor): Input logits with token candidates on the last dimension.
-        temperature (float): Sampling temperature. Lower temperature results in more deterministic samples.
-        top_p (float): The p in “top-p”.
-        top_k (int): The k in “top-k”.
-        min_p (float): Minimum token probability, scaled by the probability of the most likely token.
-                       Must be between 0 and 1. Typical values are in the 0.01-0.2 range.
-
-    Returns:
-        torch.Tensor: Sampled tokens.
-    """
-    if repetition_penalty != 1.0 and generated_tokens is not None:
-        logits = modify_logit_for_repetition_penalty(logits, generated_tokens, repetition_penalty, repetition_penalty_window)
-
-    if temperature > 0:
-        probs = torch.softmax(logits / temperature, dim=-1)
-
-        if top_p > 0:
-            probs = apply_top_p(probs, top_p)
-        if top_k > 0:
-            probs = apply_top_k(probs, top_k)
-        if min_p > 0:
-            probs = apply_min_p(probs, min_p)
-
-        next_token = multinomial(probs, num_samples=1)
-    else:
-        next_token = torch.argmax(logits, dim=-1, keepdim=True)
-
-    return next_token  # [batch_size, num_codebooks, 1]

qhash/speaker_cloning.py

@@ -1,406 +0,0 @@
-import math
-from functools import cache
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torchaudio
-from huggingface_hub import hf_hub_download
-import os
-
-
-class logFbankCal(nn.Module):
-    def __init__(
-        self,
-        sample_rate: int = 16_000,
-        n_fft: int = 512,
-        win_length: float = 0.025,
-        hop_length: float = 0.01,
-        n_mels: int = 80,
-    ):
-        super().__init__()
-        self.fbankCal = torchaudio.transforms.MelSpectrogram(
-            sample_rate=sample_rate,
-            n_fft=n_fft,
-            win_length=int(win_length * sample_rate),
-            hop_length=int(hop_length * sample_rate),
-            n_mels=n_mels,
-        )
-
-    def forward(self, x):
-        out = self.fbankCal(x)
-        out = torch.log(out + 1e-6)
-        out = out - out.mean(axis=2).unsqueeze(dim=2)
-        return out
-
-
-class ASP(nn.Module):
-    # Attentive statistics pooling
-    def __init__(self, in_planes, acoustic_dim):
-        super(ASP, self).__init__()
-        outmap_size = int(acoustic_dim / 8)
-        self.out_dim = in_planes * 8 * outmap_size * 2
-
-        self.attention = nn.Sequential(
-            nn.Conv1d(in_planes * 8 * outmap_size, 128, kernel_size=1),
-            nn.ReLU(),
-            nn.BatchNorm1d(128),
-            nn.Conv1d(128, in_planes * 8 * outmap_size, kernel_size=1),
-            nn.Softmax(dim=2),
-        )
-
-    def forward(self, x):
-        x = x.reshape(x.size()[0], -1, x.size()[-1])
-        w = self.attention(x)
-        mu = torch.sum(x * w, dim=2)
-        sg = torch.sqrt((torch.sum((x**2) * w, dim=2) - mu**2).clamp(min=1e-5))
-        x = torch.cat((mu, sg), 1)
-
-        x = x.view(x.size()[0], -1)
-        return x
-
-
-class SimAMBasicBlock(nn.Module):
-    expansion = 1
-
-    def __init__(self, ConvLayer, NormLayer, in_planes, planes, stride=1, block_id=1):
-        super(SimAMBasicBlock, self).__init__()
-        self.conv1 = ConvLayer(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
-        self.bn1 = NormLayer(planes)
-        self.conv2 = ConvLayer(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
-        self.bn2 = NormLayer(planes)
-        self.relu = nn.ReLU(inplace=True)
-        self.sigmoid = nn.Sigmoid()
-
-        self.downsample = nn.Sequential()
-        if stride != 1 or in_planes != self.expansion * planes:
-            self.downsample = nn.Sequential(
-                ConvLayer(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
-                NormLayer(self.expansion * planes),
-            )
-
-    def forward(self, x):
-        out = self.relu(self.bn1(self.conv1(x)))
-        out = self.bn2(self.conv2(out))
-        out = self.SimAM(out)
-        out += self.downsample(x)
-        out = self.relu(out)
-        return out
-
-    def SimAM(self, X, lambda_p=1e-4):
-        n = X.shape[2] * X.shape[3] - 1
-        d = (X - X.mean(dim=[2, 3], keepdim=True)).pow(2)
-        v = d.sum(dim=[2, 3], keepdim=True) / n
-        E_inv = d / (4 * (v + lambda_p)) + 0.5
-        return X * self.sigmoid(E_inv)
-
-
-class BasicBlock(nn.Module):
-    expansion = 1
-
-    def __init__(self, ConvLayer, NormLayer, in_planes, planes, stride=1, block_id=1):
-        super(BasicBlock, self).__init__()
-        self.conv1 = ConvLayer(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
-        self.bn1 = NormLayer(planes)
-        self.conv2 = ConvLayer(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
-        self.bn2 = NormLayer(planes)
-        self.relu = nn.ReLU(inplace=True)
-
-        self.downsample = nn.Sequential()
-        if stride != 1 or in_planes != self.expansion * planes:
-            self.downsample = nn.Sequential(
-                ConvLayer(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
-                NormLayer(self.expansion * planes),
-            )
-
-    def forward(self, x):
-        out = self.relu(self.bn1(self.conv1(x)))
-        out = self.bn2(self.conv2(out))
-        out += self.downsample(x)
-        out = self.relu(out)
-        return out
-
-
-class Bottleneck(nn.Module):
-    expansion = 4
-
-    def __init__(self, ConvLayer, NormLayer, in_planes, planes, stride=1, block_id=1):
-        super(Bottleneck, self).__init__()
-        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
-        self.bn1 = nn.BatchNorm2d(planes)
-        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
-        self.bn2 = nn.BatchNorm2d(planes)
-        self.conv3 = nn.Conv2d(planes, self.expansion * planes, kernel_size=1, bias=False)
-        self.bn3 = nn.BatchNorm2d(self.expansion * planes)
-
-        self.shortcut = nn.Sequential()
-        if stride != 1 or in_planes != self.expansion * planes:
-            self.shortcut = nn.Sequential(
-                nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
-                nn.BatchNorm2d(self.expansion * planes),
-            )
-
-    def forward(self, x):
-        out = F.relu(self.bn1(self.conv1(x)))
-        out = F.relu(self.bn2(self.conv2(out)))
-        out = self.bn3(self.conv3(out))
-        out += self.shortcut(x)
-        out = F.relu(out)
-        return out
-
-
-class ResNet(nn.Module):
-    def __init__(self, in_planes, block, num_blocks, in_ch=1, feat_dim="2d", **kwargs):
-        super(ResNet, self).__init__()
-        if feat_dim == "1d":
-            self.NormLayer = nn.BatchNorm1d
-            self.ConvLayer = nn.Conv1d
-        elif feat_dim == "2d":
-            self.NormLayer = nn.BatchNorm2d
-            self.ConvLayer = nn.Conv2d
-        elif feat_dim == "3d":
-            self.NormLayer = nn.BatchNorm3d
-            self.ConvLayer = nn.Conv3d
-        else:
-            print("error")
-
-        self.in_planes = in_planes
-
-        self.conv1 = self.ConvLayer(in_ch, in_planes, kernel_size=3, stride=1, padding=1, bias=False)
-        self.bn1 = self.NormLayer(in_planes)
-        self.relu = nn.ReLU(inplace=True)
-        self.layer1 = self._make_layer(block, in_planes, num_blocks[0], stride=1, block_id=1)
-        self.layer2 = self._make_layer(block, in_planes * 2, num_blocks[1], stride=2, block_id=2)
-        self.layer3 = self._make_layer(block, in_planes * 4, num_blocks[2], stride=2, block_id=3)
-        self.layer4 = self._make_layer(block, in_planes * 8, num_blocks[3], stride=2, block_id=4)
-
-    def _make_layer(self, block, planes, num_blocks, stride, block_id=1):
-        strides = [stride] + [1] * (num_blocks - 1)
-        layers = []
-        for stride in strides:
-            layers.append(block(self.ConvLayer, self.NormLayer, self.in_planes, planes, stride, block_id))
-            self.in_planes = planes * block.expansion
-        return nn.Sequential(*layers)
-
-    def forward(self, x):
-        x = self.relu(self.bn1(self.conv1(x)))
-        x = self.layer1(x)
-        x = self.layer2(x)
-        x = self.layer3(x)
-        x = self.layer4(x)
-        return x
-
-
-def ResNet293(in_planes: int, **kwargs):
-    return ResNet(in_planes, SimAMBasicBlock, [10, 20, 64, 3], **kwargs)
-
-
-class ResNet293_based(nn.Module):
-    def __init__(
-        self,
-        in_planes: int = 64,
-        embd_dim: int = 256,
-        acoustic_dim: int = 80,
-        featCal=None,
-        dropout: float = 0,
-        **kwargs,
-    ):
-        super(ResNet293_based, self).__init__()
-        self.featCal = featCal
-        self.front = ResNet293(in_planes)
-        block_expansion = SimAMBasicBlock.expansion
-        self.pooling = ASP(in_planes * block_expansion, acoustic_dim)
-        self.bottleneck = nn.Linear(self.pooling.out_dim, embd_dim)
-        self.drop = nn.Dropout(dropout) if dropout else None
-
-    def forward(self, x):
-        x = self.featCal(x)
-        x = self.front(x.unsqueeze(dim=1))
-        x = self.pooling(x)
-        if self.drop:
-            x = self.drop(x)
-        x = self.bottleneck(x)
-        return x
-
-
-class SEModule(nn.Module):
-    def __init__(self, channels, bottleneck=128):
-        super(SEModule, self).__init__()
-        self.se = nn.Sequential(
-            nn.AdaptiveAvgPool1d(1),
-            nn.Conv1d(channels, bottleneck, kernel_size=1, padding=0),
-            nn.ReLU(),
-            # nn.BatchNorm1d(bottleneck), # Removed
-            nn.Conv1d(bottleneck, channels, kernel_size=1, padding=0),
-            nn.Sigmoid(),
-        )
-
-    def forward(self, input):
-        x = self.se(input)
-        return input * x
-
-
-class Bottle2neck(nn.Module):
-    def __init__(self, inplanes, planes, kernel_size=None, dilation=None, scale=8):
-        super(Bottle2neck, self).__init__()
-        width = int(math.floor(planes / scale))
-        self.conv1 = nn.Conv1d(inplanes, width * scale, kernel_size=1)
-        self.bn1 = nn.BatchNorm1d(width * scale)
-        self.nums = scale - 1
-        convs = []
-        bns = []
-        num_pad = math.floor(kernel_size / 2) * dilation
-        for i in range(self.nums):
-            convs.append(nn.Conv1d(width, width, kernel_size=kernel_size, dilation=dilation, padding=num_pad))
-            bns.append(nn.BatchNorm1d(width))
-        self.convs = nn.ModuleList(convs)
-        self.bns = nn.ModuleList(bns)
-        self.conv3 = nn.Conv1d(width * scale, planes, kernel_size=1)
-        self.bn3 = nn.BatchNorm1d(planes)
-        self.relu = nn.ReLU()
-        self.width = width
-        self.se = SEModule(planes)
-
-    def forward(self, x):
-        residual = x
-        out = self.conv1(x)
-        out = self.relu(out)
-        out = self.bn1(out)
-
-        spx = torch.split(out, self.width, 1)
-        for i in range(self.nums):
-            if i == 0:
-                sp = spx[i]
-            else:
-                sp = sp + spx[i]
-            sp = self.convs[i](sp)
-            sp = self.relu(sp)
-            sp = self.bns[i](sp)
-            if i == 0:
-                out = sp
-            else:
-                out = torch.cat((out, sp), 1)
-        out = torch.cat((out, spx[self.nums]), 1)
-
-        out = self.conv3(out)
-        out = self.relu(out)
-        out = self.bn3(out)
-
-        out = self.se(out)
-        out += residual
-        return out
-
-
-class ECAPA_TDNN(nn.Module):
-    def __init__(self, C, featCal):
-        super(ECAPA_TDNN, self).__init__()
-        self.featCal = featCal
-        self.conv1 = nn.Conv1d(80, C, kernel_size=5, stride=1, padding=2)
-        self.relu = nn.ReLU()
-        self.bn1 = nn.BatchNorm1d(C)
-        self.layer1 = Bottle2neck(C, C, kernel_size=3, dilation=2, scale=8)
-        self.layer2 = Bottle2neck(C, C, kernel_size=3, dilation=3, scale=8)
-        self.layer3 = Bottle2neck(C, C, kernel_size=3, dilation=4, scale=8)
-        # I fixed the shape of the output from MFA layer, that is close to the setting from ECAPA paper.
-        self.layer4 = nn.Conv1d(3 * C, 1536, kernel_size=1)
-        self.attention = nn.Sequential(
-            nn.Conv1d(4608, 256, kernel_size=1),
-            nn.ReLU(),
-            nn.BatchNorm1d(256),
-            nn.Tanh(),  # Added
-            nn.Conv1d(256, 1536, kernel_size=1),
-            nn.Softmax(dim=2),
-        )
-        self.bn5 = nn.BatchNorm1d(3072)
-        self.fc6 = nn.Linear(3072, 192)
-        self.bn6 = nn.BatchNorm1d(192)
-
-    def forward(self, x):
-        x = self.featCal(x)
-        x = self.conv1(x)
-        x = self.relu(x)
-        x = self.bn1(x)
-
-        x1 = self.layer1(x)
-        x2 = self.layer2(x + x1)
-        x3 = self.layer3(x + x1 + x2)
-
-        x = self.layer4(torch.cat((x1, x2, x3), dim=1))
-        x = self.relu(x)
-
-        t = x.size()[-1]
-
-        global_x = torch.cat(
-            (
-                x,
-                torch.mean(x, dim=2, keepdim=True).repeat(1, 1, t),
-                torch.sqrt(torch.var(x, dim=2, keepdim=True).clamp(min=1e-4)).repeat(1, 1, t),
-            ),
-            dim=1,
-        )
-
-        w = self.attention(global_x)
-
-        mu = torch.sum(x * w, dim=2)
-        sg = torch.sqrt((torch.sum((x**2) * w, dim=2) - mu**2).clamp(min=1e-4))
-
-        x = torch.cat((mu, sg), 1)
-        x = self.bn5(x)
-        x = self.fc6(x)
-        x = self.bn6(x)
-
-        return x
-
-
-class SpeakerEmbedding(nn.Module):
-    def __init__(self, ckpt_path: str = "ResNet293_SimAM_ASP_base.pt", device: str = "cuda"):
-        super().__init__()
-        self.device = device
-        with torch.device(device):
-            self.model = ResNet293_based()
-            self.model.load_state_dict(torch.load(ckpt_path, weights_only=True, mmap=True))
-            self.model.featCal = logFbankCal()
-
-        self.requires_grad_(False).eval()
-
-    @property
-    def dtype(self):
-        return next(self.parameters()).dtype
-
-    @cache
-    def _get_resampler(self, orig_sample_rate: int):
-        return torchaudio.transforms.Resample(orig_sample_rate, 16_000).to(self.device)
-
-    def prepare_input(self, wav: torch.Tensor, sample_rate: int) -> torch.Tensor:
-        assert wav.ndim < 3
-        if wav.ndim == 2:
-            wav = wav.mean(0, keepdim=True)
-        wav = self._get_resampler(sample_rate)(wav)
-        return wav
-
-    def forward(self, wav: torch.Tensor, sample_rate: int):
-        wav = self.prepare_input(wav, sample_rate).to(self.device, self.dtype)
-        return self.model(wav).to(wav.device)
-
-class SpeakerEmbeddingLDA(nn.Module):
-    def __init__(
-        self,
-        device: str = "cuda",
-    ):
-        super().__init__()
-        spk_model_path = hf_hub_download(repo_id="Quantamhash/Qhash-v0.1-speaker-embedding", filename="ResNet293_SimAM_ASP_base.pt")
-        lda_spk_model_path = hf_hub_download(repo_id="Quantamhash/Qhash-v0.1-speaker-embedding", filename="ResNet293_SimAM_ASP_base_LDA-128.pt")
-
-        self.device = device
-        with torch.device(device):
-            self.model = SpeakerEmbedding(spk_model_path, device)
-            lda_sd = torch.load(lda_spk_model_path, weights_only=True)
-            out_features, in_features = lda_sd["weight"].shape
-            self.lda = nn.Linear(in_features, out_features, bias=True, dtype=torch.float32)
-            self.lda.load_state_dict(lda_sd)
-
-        self.requires_grad_(False).eval()
-
-    def forward(self, wav: torch.Tensor, sample_rate: int):
-        emb = self.model(wav, sample_rate).to(torch.float32)
-        return emb, self.lda(emb)