...

comfy/ldm/higgsv2/cuda_graph_runner.py

@@ -1,59 +0,0 @@
-import torch
-import torch.nn as nn
-from typing import Optional, Dict
-import gc
-
-_NUM_WARMUP_ITERS = 2
-
-class CUDAGraphRunner(nn.Module):
-    def __init__(self, model):
-        super().__init__()
-        self.model = model
-
-        self.input_buffers: Dict[str, torch.Tensor] = {}
-        self.output_buffers: Dict[str, torch.Tensor] = {}
-
-        self._graph: Optional[torch.cuda.CUDAGraph] = None
-
-    @property
-    def graph(self):
-        assert self._graph is not None
-        return self._graph
-
-    def capture(self, *args, **kwargs):
-        assert self._graph is None
-
-        for _ in range(_NUM_WARMUP_ITERS):
-            self.model(*args, **kwargs)
-
-        torch.cuda.synchronize()
-
-        self._graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(self._graph, pool = kwargs.get("memory_pool", None), stream = kwargs.get("stream", None)):
-            last_hidden_states = self.model(*args, **kwargs)
-            gc.collect()
-
-        torch.cuda.synchronize()
-
-        self.input_buffers = {
-            "args": [arg for arg in args if isinstance(arg, torch.Tensor)],
-            "kwargs": {k: v for k, v in kwargs.items() if isinstance(v, torch.Tensor)},
-        }
-
-        self.output_buffers = {
-            "hidden_states": last_hidden_states
-        }
-
-    def forward(self, *args, **kwargs):
-
-        for i, arg in enumerate(args):
-            if isinstance(arg, torch.Tensor):
-                self.input_buffers["args"][i].copy_(arg, non_blocking=True)
-
-        for k, v in kwargs.items():
-            if k in self.input_buffers["kwargs"] and isinstance(v, torch.Tensor):
-                self.input_buffers["kwargs"][k].copy_(v, non_blocking=True)
-
-        self.graph.replay()
-
-        return self.output_buffers["hidden_states"]

comfy/ldm/higgsv2/loudness.py

@@ -1,330 +0,0 @@
-import copy
-import math
-import torch
-import scipy
-import torchaudio
-import numpy as np
-import torch.nn.functional as F
-from typing import Optional, List
-
-# defaulted to the new pytorch api
-def _new_rfft(x: torch.Tensor):
-    z = torch.fft.rfft(x, dim=-1)
-    return torch.view_as_real(z)
-
-def _new_irfft(x: torch.Tensor, length: int):
-    x = torch.view_as_complex(x)
-    return torch.fft.irfft(x, length, dim=-1)
-
-def _compl_mul_conjugate(a: torch.Tensor, b: torch.Tensor):
-    # changed this function to use the pytorch api
-    return torch.view_as_real(torch.view_as_complex(a) * torch.view_as_complex(b).conj())
-
-def unfold(input, kernel_size: int, stride: int):
-
-    shape = list(input.shape)
-    length = shape.pop(-1)
-
-    n_frames = math.ceil((max(length, kernel_size) - kernel_size) / stride) + 1
-    tgt_length = (n_frames - 1) * stride + kernel_size
-
-    padded = F.pad(input, (0, tgt_length - length)).contiguous()
-    strides: List[int] = []
-
-    for dim in range(padded.dim()):
-        strides.append(padded.stride(dim))
-
-    last_stride = strides.pop(-1)
-    assert last_stride == 1, 'data should be contiguous'
-
-    strides = strides + [stride, 1]
-    return padded.as_strided(shape + [n_frames, kernel_size], strides)
-
-# convert the signal and filter to frequency domain, multiply them, then inverse FFT to get back to time-domain
-# faster than a sliding window over time-domain.
-def fft_conv1d(
-        input: torch.Tensor, weight: torch.Tensor,
-        bias: Optional[torch.Tensor] = None, stride: int = 1, padding: int = 0,
-        block_ratio: float = 5):
-
-    input = F.pad(input, (padding, padding))
-    batch, _, length = input.shape
-    out_channels, _, kernel_size = weight.shape
-
-    _rfft = _new_rfft
-    _irfft = _new_irfft
-
-    if length < kernel_size:
-        raise RuntimeError(f"Input should be at least as large as the kernel size {kernel_size}, "
-                           f"but it is only {length} samples long.")
-    if block_ratio < 1:
-        raise RuntimeError("Block ratio must be greater than 1.")
-
-    # We are going to process the input blocks by blocks, as for some reason it is faster
-    # and less memory intensive (I think the culprit is `torch.einsum`.
-    block_size: int = min(int(kernel_size * block_ratio), length)
-    fold_stride = block_size - kernel_size + 1
-
-    # replaces to_pad
-    weight = F.pad(weight, (0, block_size - weight.shape[-1]), mode = "constant", value = 0)
-    weight_z = _rfft(weight)
-
-    # We pad the input and get the different frames, on which
-    frames = unfold(input, block_size, fold_stride)
-
-    frames_z = _rfft(frames)
-    out_z = _compl_mul_conjugate(frames_z, weight_z)
-    out = _irfft(out_z, block_size)
-    # The last bit is invalid, because FFT will do a circular convolution.
-    out = out[..., :-kernel_size + 1]
-    out = out.reshape(batch, out_channels, -1)
-    out = out[..., ::stride]
-    target_length = (length - kernel_size) // stride + 1
-    out = out[..., :target_length]
-    if bias is not None:
-        out += bias[:, None]
-    return out
-
-class IIRfilter(object):
-
-    def __init__(self, G, Q, fc, rate, filter_type, passband_gain=1.0):
-        self.G  = G
-        self.Q  = Q
-        self.fc = fc
-        self.rate = rate
-        self.filter_type = filter_type
-        self.passband_gain = passband_gain
-
-    def generate_coefficients(self):
-
-        A  = 10**(self.G/40.0)
-        w0 = 2.0 * np.pi * (self.fc / self.rate)
-        alpha = np.sin(w0) / (2.0 * self.Q)
-
-        if self.filter_type == 'high_shelf':
-            b0 =      A * ( (A+1) + (A-1) * np.cos(w0) + 2 * np.sqrt(A) * alpha )
-            b1 = -2 * A * ( (A-1) + (A+1) * np.cos(w0)                          )
-            b2 =      A * ( (A+1) + (A-1) * np.cos(w0) - 2 * np.sqrt(A) * alpha )
-            a0 =            (A+1) - (A-1) * np.cos(w0) + 2 * np.sqrt(A) * alpha
-            a1 =      2 * ( (A-1) - (A+1) * np.cos(w0)                          )
-            a2 =            (A+1) - (A-1) * np.cos(w0) - 2 * np.sqrt(A) * alpha
-
-        elif self.filter_type == 'high_pass':
-            b0 =  (1 + np.cos(w0))/2
-            b1 = -(1 + np.cos(w0))
-            b2 =  (1 + np.cos(w0))/2
-            a0 =   1 + alpha
-            a1 =  -2 * np.cos(w0)
-            a2 =   1 - alpha
-
-        return np.array([b0, b1, b2])/a0, np.array([a0, a1, a2])/a0
-
-    def apply_filter(self, data):
-        return self.passband_gain * scipy.signal.lfilter(self.b, self.a, data)
-
-    @property
-    def b_and_a(self):
-        return self.generate_coefficients()
-
-class Meter(torch.nn.Module):
-
-    def __init__(
-        self,
-        rate: int,
-        filter_class: str = "K-weighting",
-        block_size: float = 0.400,
-        zeros: int = 512,
-        use_fir: bool = False,
-    ):
-        super().__init__()
-
-        self.rate = rate
-        self.filter_class = filter_class
-        self.block_size = block_size
-        self.use_fir = use_fir
-
-        G = torch.from_numpy(np.array([1.0, 1.0, 1.0, 1.41, 1.41]))
-        self.register_buffer("G", G)
-
-        self._filters = {}
-        self._filters['high_shelf'] = IIRfilter(4.0, 1/np.sqrt(2), 1500.0, self.rate, 'high_shelf')
-        self._filters['high_pass'] = IIRfilter(0.0, 0.5, 38.0, self.rate, 'high_pass')
-
-        # Compute impulse responses so that filtering is fast via
-        # a convolution at runtime, on GPU, unlike lfilter.
-        impulse = np.zeros((zeros,))
-        impulse[..., 0] = 1.0
-
-        firs = np.zeros((len(self._filters), 1, zeros))
-        passband_gain = torch.tensor([filter.passband_gain for filter in self._filters.values()])
-
-        for i, (_, filter_stage) in enumerate(self._filters.items()):
-            b, a = filter_stage.b_and_a
-            firs[i] = scipy.signal.lfilter(b, a, impulse)
-
-        firs = torch.from_numpy(firs[..., ::-1].copy()).float()
-
-        self.register_buffer("firs", firs)
-        self.register_buffer("passband_gain", passband_gain)
-
-    def apply_filter_gpu(self, data: torch.Tensor):
-
-        # Data is of shape (nb, nch, nt)
-        # Reshape to (nb*nch, 1, nt)
-        nb, nt, nch = data.shape
-        data = data.permute(0, 2, 1)
-        data = data.reshape(nb * nch, 1, nt)
-
-        # Apply padding
-        pad_length = self.firs.shape[-1]
-
-        # Apply filtering in sequence
-        for i in range(self.firs.shape[0]):
-            data = F.pad(data, (pad_length, pad_length))
-            data = fft_conv1d(data, self.firs[i, None, ...])
-            data = self.passband_gain[i] * data
-            data = data[..., 1 : nt + 1]
-
-        data = data.permute(0, 2, 1)
-        data = data[:, :nt, :]
-        return data
-
-    def apply_filter_cpu(self, data: torch.Tensor):
-        for _, filter_stage in self._filters.items():
-            passband_gain = filter_stage.passband_gain
-            b, a = filter_stage.b_and_a
-
-            a_coeffs = torch.from_numpy(a).float().to(data.device)
-            b_coeffs = torch.from_numpy(b).float().to(data.device)
-
-            _data = data.permute(0, 2, 1)
-            filtered = torchaudio.functional.lfilter(
-                _data, a_coeffs, b_coeffs, clamp=False
-            )
-            data = passband_gain * filtered.permute(0, 2, 1)
-        return data
-
-    def apply_filter(self, data: torch.Tensor):
-        if data.is_cuda or self.use_fir:
-            data = self.apply_filter_gpu(data)
-        else:
-            data = self.apply_filter_cpu(data)
-        return data
-
-    def forward(self, data: torch.Tensor):
-        return self.integrated_loudness(data)
-
-    def _unfold(self, input_data):
-        T_g = self.block_size
-        overlap = 0.75  # overlap of 75% of the block duration
-        step = 1.0 - overlap  # step size by percentage
-
-        kernel_size = int(T_g * self.rate)
-        stride = int(T_g * self.rate * step)
-        unfolded = unfold(input_data.permute(0, 2, 1), kernel_size, stride)
-        unfolded = unfolded.transpose(-1, -2)
-
-        return unfolded
-
-    def integrated_loudness(self, data: torch.Tensor):
-
-        if not torch.is_tensor(data):
-            data = torch.from_numpy(data).float()
-        else:
-            data = data.float()
-
-        input_data = copy.copy(data)
-        # Data always has a batch and channel dimension.
-        # Is of shape (nb, nt, nch)
-        if input_data.ndim < 2:
-            input_data = input_data.unsqueeze(-1)
-        if input_data.ndim < 3:
-            input_data = input_data.unsqueeze(0)
-
-        nb, _, nch = input_data.shape
-
-        # Apply frequency weighting filters - account
-        # for the acoustic respose of the head and auditory system
-        input_data = self.apply_filter(input_data)
-
-        G = self.G  # channel gains
-        T_g = self.block_size  # 400 ms gating block standard
-        Gamma_a = -70.0  # -70 LKFS = absolute loudness threshold
-
-        unfolded = self._unfold(input_data)
-
-        z = (1.0 / (T_g * self.rate)) * unfolded.square().sum(2)
-        l = -0.691 + 10.0 * torch.log10((G[None, :nch, None] * z).sum(1, keepdim=True))
-        l = l.expand_as(z)
-
-        # find gating block indices above absolute threshold
-        z_avg_gated = z
-        z_avg_gated[l <= Gamma_a] = 0
-        masked = l > Gamma_a
-        z_avg_gated = z_avg_gated.sum(2) / masked.sum(2)
-
-        # calculate the relative threshold value (see eq. 6)
-        Gamma_r = (
-            -0.691 + 10.0 * torch.log10((z_avg_gated * G[None, :nch]).sum(-1)) - 10.0
-        )
-        Gamma_r = Gamma_r[:, None, None]
-        Gamma_r = Gamma_r.expand(nb, nch, l.shape[-1])
-
-        # find gating block indices above relative and absolute thresholds  (end of eq. 7)
-        z_avg_gated = z
-        z_avg_gated[l <= Gamma_a] = 0
-        z_avg_gated[l <= Gamma_r] = 0
-        masked = (l > Gamma_a) * (l > Gamma_r)
-        z_avg_gated = z_avg_gated.sum(2) / masked.sum(2)
-
-        # # Cannot use nan_to_num (pytorch 1.8 does not come with GCP-supported cuda version)
-        # z_avg_gated = torch.nan_to_num(z_avg_gated)
-        z_avg_gated = torch.where(
-            z_avg_gated.isnan(), torch.zeros_like(z_avg_gated), z_avg_gated
-        )
-        z_avg_gated[z_avg_gated == float("inf")] = float(np.finfo(np.float32).max)
-        z_avg_gated[z_avg_gated == -float("inf")] = float(np.finfo(np.float32).min)
-
-        LUFS = -0.691 + 10.0 * torch.log10((G[None, :nch] * z_avg_gated).sum(1))
-        return LUFS.float()
-
-
-def loudness(
-    audio_data, sample_rate: int, target_loudness: int, filter_class: str = "K-weighting", block_size: float = 0.400, **kwargs
-):
-    MIN_LOUDNESS = -70
-    device = audio_data.device
-
-    original_length = audio_data.shape[-1]
-    signal_duration = original_length / sample_rate
-
-    # Pad if too short
-    if signal_duration < 0.5:
-        pad_len = int((0.5 - signal_duration) * sample_rate)
-        audio_data = torch.nn.functional.pad(audio_data, (0, pad_len), mode="constant", value=0)
-
-    # create BS.1770 meter
-    meter = Meter(
-        sample_rate, filter_class=filter_class, block_size=block_size, **kwargs
-    )
-    meter = meter.to(audio_data.device)
-    # measure loudness
-    loudness = meter.integrated_loudness(audio_data.permute(0, 2, 1))
-    audio_data = audio_data[..., :original_length]
-    min_loudness = (
-        torch.ones_like(loudness, device=loudness.device) * MIN_LOUDNESS
-    )
-    _loudness = torch.maximum(loudness, min_loudness)
-
-    _loudness = _loudness.to(device)
-
-    delta_loudness = target_loudness - _loudness
-    gain = torch.pow(torch.tensor(10.0, device=device, dtype=audio_data.dtype), delta_loudness / 20.0)
-
-    output = gain * audio_data
-
-    if torch.max(torch.abs(output)) >= 1.0:
-        import warnings
-        warnings.warn("Possible clipped samples in output.")
-
-    return output

comfy/ldm/higgsv2/preprocess.py

@@ -1,727 +0,0 @@
-from typing import (
-    Dict, List, Optional, Union, Tuple, MutableMapping, Any, Mapping, Collection, get_type_hints, get_args, get_origin
-)
-
-from dataclasses import dataclass, fields, _FIELDS, _FIELD, _FIELD_INITVAR, is_dataclass
-import torch.nn.functional as F
-import numpy as np
-import torch
-import math
-from functools import lru_cache
-
-__MAX_SIZE = 2048
-
-def _ceil_to_nearest(n, round_to):
-    return (n + round_to - 1) // round_to * round_to
-
-@torch.jit.script
-def build_delay_pattern_mask(
-    input_ids: torch.Tensor,
-    bos_token_id: int,
-    pad_token_id: int,
-):
-    bsz, num_codebooks, seq_len = input_ids.shape
-
-    new_seq_len = seq_len + num_codebooks - 1
-    input_ids_with_gen_mask = torch.ones((bsz, num_codebooks, new_seq_len), dtype=torch.long, device=input_ids.device)
-    bos_mask = torch.tril(input_ids_with_gen_mask, -1) > 0
-    eos_mask = torch.triu(input_ids_with_gen_mask, seq_len) > 0
-    input_ids_with_gen_mask[bos_mask] = bos_token_id
-    input_ids_with_gen_mask[(~bos_mask) & (~eos_mask)] = input_ids.reshape(-1)
-    input_ids = input_ids_with_gen_mask.clone()
-    input_ids[eos_mask] = pad_token_id
-    input_ids_with_gen_mask[eos_mask] = -1
-    return input_ids, input_ids_with_gen_mask
-
-
-# implementation of dacite's from_dict with only necessary parts for ChatML
-@lru_cache(maxsize=None)
-def cache(function):
-    return lru_cache(maxsize=__MAX_SIZE, typed=True)(function)
-
-@cache
-def get_fields(data_class):
-    fields = getattr(data_class, _FIELDS)
-    return [f for f in fields.values() if f._field_type is _FIELD or f._field_type is _FIELD_INITVAR]
-
-def is_optional(type_) -> bool:
-    return get_origin(type_) is Union and type(None) in get_args(type_)
-
-def orig(data_class) -> Any:
-    if is_dataclass(data_class):
-        return data_class
-    return get_origin(data_class)
-
-@cache
-def extract_generic(type_, defaults: Tuple = ()) -> tuple:
-    try:
-        if getattr(type_, "_special", False):
-            return defaults
-        if type_.__args__ == ():
-            return (type_.__args__,)
-        return type_.__args__ or defaults  # type: ignore
-    except AttributeError:
-        return defaults
-
-def _build_value_for_collection(collection, data: Any) -> Any:
-    if isinstance(data, Mapping):
-        value_type = extract_generic(collection, defaults=(Any, Any))[1]
-        return {
-            key: _build_value(type_=value_type, data=value)
-            for key, value in data.items()
-        }
-
-    elif isinstance(data, Collection) and not isinstance(data, (str, bytes, Mapping)):
-        item_type = extract_generic(collection, defaults=(Any,))[0]
-        return [
-            _build_value(type_=item_type, data=item)
-            for item in data
-        ]
-
-    return data
-
-def _build_value(type_, data) -> Any:
-    if is_optional(type_) and data is None:
-        return data
-    if get_origin(type_) is Union:
-        data = _build_value_for_union(union=type_, data=data)
-    elif hasattr(type_, "__origin__"):
-        data = _build_value_for_collection(collection=type_, data=data)
-    elif cache(is_dataclass)(orig(type_)) and isinstance(data, Mapping):
-        data = from_dict(data_class=type_, data=data)
-    return data
-
-def _build_value_for_union(union: type, data: Any) -> Any:
-    for inner_type in get_args(union):
-        if data is None and inner_type is type(None):
-            return None
-        try:
-            return _build_value(inner_type, data)
-        except Exception:
-            continue
-    raise ValueError(f"Cannot match {data!r} to any type in {union}")
-
-def is_instance(value: Any, type_) -> bool:
-    if type_ is Any:
-        return True
-
-    origin = get_origin(type_)
-    args = get_args(type_)
-
-    if origin is Union:
-        return any(is_instance(value, arg) for arg in args)
-
-    if origin in (list, List):
-        if not isinstance(value, list):
-            return False
-        (elem_type,) = args or (Any,)
-        return all(is_instance(item, elem_type) for item in value)
-
-    if origin in (dict, Dict, Mapping):
-        if not isinstance(value, dict):
-            return False
-        key_type, val_type = args or (Any, Any)
-        return all(
-            is_instance(k, key_type) and is_instance(v, val_type)
-            for k, v in value.items()
-        )
-
-    try:
-        return isinstance(value, type_)
-    except TypeError:
-        return False
-
-def from_dict(data_class, data):
-
-    init_values: MutableMapping[str, Any] = {}
-    post_init_values: MutableMapping[str, Any] = {}
-
-    data_class_hints = get_type_hints(data_class)
-    data_class_fields = cache(get_fields)(data_class)
-
-    extra_fields = set(data.keys()) - {f.name for f in data_class_fields}
-    if extra_fields:
-        formatted_keys = ", ".join(f'"{key}"' for key in extra_fields)
-        raise ValueError(f"cannot match {formatted_keys} to any data class field")
-
-    for field in data_class_fields:
-        field_type = data_class_hints[field.name]
-        key = field.name
-
-        if key in data:
-            try:
-                value = _build_value(type_=field_type, data=data[key])
-            except Exception as error:
-                raise ValueError(error)
-
-            if not is_instance(value, field_type):
-                raise ValueError((
-                    f'wrong value type for field "{field.name}" - should be "{field_type}" '
-                    f'instead of value "{value}" of type "{type(value)}"'
-                ))
-
-            init_values[field.name] = value
-
-    instance = data_class(**init_values)
-
-    for key, value in post_init_values.items():
-        setattr(instance, key, value)
-
-    return instance
-
-def normalize_chinese_punctuation(text):
-    """
-    Convert Chinese (full-width) punctuation marks to English (half-width) equivalents.
-    """
-    # Mapping of Chinese punctuation to English punctuation
-    chinese_to_english_punct = {
-        "，": ", ",  # comma
-        "。": ".",  # period
-        "：": ":",  # colon
-        "；": ";",  # semicolon
-        "？": "?",  # question mark
-        "！": "!",  # exclamation mark
-        "（": "(",  # left parenthesis
-        "）": ")",  # right parenthesis
-        "【": "[",  # left square bracket
-        "】": "]",  # right square bracket
-        "《": "<",  # left angle quote
-        "》": ">",  # right angle quote
-        "“": '"',  # left double quotation
-        "”": '"',  # right double quotation
-        "‘": "'",  # left single quotation
-        "’": "'",  # right single quotation
-        "、": ",",  # enumeration comma
-        "—": "-",  # em dash
-        "…": "...",  # ellipsis
-        "·": ".",  # middle dot
-        "「": '"',  # left corner bracket
-        "」": '"',  # right corner bracket
-        "『": '"',  # left double corner bracket
-        "』": '"',  # right double corner bracket
-    }
-
-    # Replace each Chinese punctuation with its English counterpart
-    for zh_punct, en_punct in chinese_to_english_punct.items():
-        text = text.replace(zh_punct, en_punct)
-
-    return text
-
-def transcript_normalize(text: str):
-    transcript = normalize_chinese_punctuation(text)
-
-    transcript = transcript.replace("(", " ")
-    transcript = transcript.replace(")", " ")
-    transcript = transcript.replace("°F", " degrees Fahrenheit")
-    transcript = transcript.replace("°C", " degrees Celsius")
-
-    for tag, replacement in [
-        ("[laugh]", "<SE>[Laughter]</SE>"),
-        ("[humming start]", "<SE_s>[Humming]</SE_s>"),
-        ("[humming end]", "<SE_e>[Humming]</SE_e>"),
-        ("[music start]", "<SE_s>[Music]</SE_s>"),
-        ("[music end]", "<SE_e>[Music]</SE_e>"),
-        ("[music]", "<SE>[Music]</SE>"),
-        ("[sing start]", "<SE_s>[Singing]</SE_s>"),
-        ("[sing end]", "<SE_e>[Singing]</SE_e>"),
-        ("[applause]", "<SE>[Applause]</SE>"),
-        ("[cheering]", "<SE>[Cheering]</SE>"),
-        ("[cough]", "<SE>[Cough]</SE>"),
-    ]:
-        transcript = transcript.replace(tag, replacement)
-    lines = transcript.split("\n")
-    transcript = "\n".join([" ".join(line.split()) for line in lines if line.strip()])
-    transcript = transcript.strip()
-
-    if not any([transcript.endswith(c) for c in [".", "!", "?", ",", ";", '"', "'", "</SE_e>", "</SE>"]]):
-        transcript += "."
-
-    return transcript
-
-@dataclass
-class AudioContent:
-    audio_url: str
-    raw_audio: Optional[str] = None
-    offset: Optional[float] = None
-    duration: Optional[float] = None
-    row_id: Optional[int] = None
-    type: str = "audio"
-
-
-@dataclass
-class TextContent:
-    text: str
-    type: str = "text"
-
-
-@dataclass
-class Message:
-    role: str
-    content: Union[str, AudioContent, TextContent, List[Union[str, AudioContent, TextContent]]]
-    recipient: Optional[str] = None
-
-
-@dataclass
-class ChatMLSample:
-    messages: List[Message]
-    start_index: Optional[int] = None
-    misc: Optional[Dict] = None
-    speaker: Optional[str] = None
-
-def prepare_chatml_sample(sample: Union[ChatMLSample, Dict], tokenizer):
-
-    try:
-        if not isinstance(sample, ChatMLSample):
-
-            # replacing pd.isna
-            def is_nan(x):
-                if isinstance(x, float):
-                    return math.isnan(x)
-                if isinstance(x, np.generic):
-                    return np.isnan(x)
-                if isinstance(x, torch.Tensor) and x.numel() == 1:
-                    return torch.isnan(x).item()
-                return False
-
-            if "speaker" in sample and is_nan(sample["speaker"]):
-                sample["speaker"] = None
-            if "start_index" in sample and is_nan(sample["start_index"]):
-                sample["start_index"] = None
-            if "content" in sample and is_nan(sample["content"]):
-                sample["content"] = ""
-
-            def convert_nan_to_none(obj):
-
-                if isinstance(obj, np.ndarray):
-                    return obj.tolist()
-                elif isinstance(obj, float) and math.isnan(obj):
-                    return None
-                elif isinstance(obj, dict):
-                    return {k: convert_nan_to_none(v) for k, v in obj.items()}
-                elif isinstance(obj, (list, tuple)):
-                    return [convert_nan_to_none(item) for item in obj]
-                return obj
-
-            clean_sample = convert_nan_to_none(sample)
-
-            val_keys = []
-            for field in fields(ChatMLSample):
-                if field.name in clean_sample:
-                    val_keys.append(field.name)
-            clean_sample = {k: clean_sample[k] for k in val_keys}
-
-            sample = from_dict(
-                data_class=ChatMLSample, data=clean_sample,
-            )
-
-        input_tokens = []
-        audio_contents = []
-        speaker_id = None
-        if sample.speaker is not None:
-            speaker_id = sample.speaker
-        elif sample.misc is not None:
-            if "speaker" in sample.misc:
-                speaker_id = sample.misc["speaker"]
-
-        total_m = len(sample.messages)
-        for turn_id, message in enumerate(sample.messages):
-            role = message.role
-            recipient = message.recipient
-            content = message.content
-            content_l = []
-
-            if isinstance(content, str):
-                content_l.append(TextContent(text=content))
-            elif isinstance(content, TextContent):
-                content_l.append(content)
-            elif isinstance(content, AudioContent):
-                content_l.append(content)
-            elif isinstance(content, list):
-                for ele in content:
-                    if isinstance(ele, str):
-                        content_l.append(TextContent(text=ele))
-                    else:
-                        content_l.append(ele)
-            if turn_id == 0:
-                prefix = f"<|begin_of_text|><|start_header_id|>{role}<|end_header_id|>\n\n"
-            else:
-                prefix = f"<|start_header_id|>{role}<|end_header_id|>\n\n"
-            eot_postfix = "<|eot_id|>"
-            eom_postfix = "<|eom_id|>"
-
-            prefix_tokens = tokenizer.encode(prefix, add_special_tokens=False)
-            input_tokens.extend(prefix_tokens)
-
-            if recipient:
-                assert role == "assistant", "Recipient is only available for assistant role."
-                recipient_tokens = tokenizer.encode(f"{recipient}<|recipient|>", add_special_tokens=False)
-                input_tokens.extend(recipient_tokens)
-
-            for content in content_l:
-                if content.type == "text":
-                    text_tokens = tokenizer.encode(content.text, add_special_tokens=False)
-                    input_tokens.extend(text_tokens)
-
-                elif content.type == "audio":
-                    audio_contents.append(content)
-                    if role == "user" or role == "system":
-                        text_tokens = tokenizer.encode(
-                            "<|audio_bos|><|AUDIO|><|audio_eos|>",
-                            add_special_tokens=False,
-                        )
-                        input_tokens.extend(text_tokens)
-                    elif role == "assistant":
-                        text_tokens = tokenizer.encode(
-                            "<|audio_out_bos|><|AUDIO_OUT|><|audio_eos|>",
-                            add_special_tokens=False,
-                        )
-                        input_tokens.extend(text_tokens)
-
-            next_id = turn_id + 1
-            if role == "assistant" and next_id != total_m and sample.messages[next_id].role == "assistant":
-                postfix_tokens = tokenizer.encode(eom_postfix, add_special_tokens=False)
-                input_tokens.extend(postfix_tokens)
-            else:
-                postfix_tokens = tokenizer.encode(eot_postfix, add_special_tokens=False)
-                input_tokens.extend(postfix_tokens)
-
-        return input_tokens, audio_contents, speaker_id
-
-    except Exception:
-        return None, None, None
-
-@dataclass
-class HiggsAudioBatchInput:
-    input_ids: torch.LongTensor  # shape (bsz, seq_len).
-    attention_mask: torch.Tensor  # shape (bsz, seq_len).
-    audio_out_ids: Optional[torch.LongTensor]  # shape (num_codebooks, audio_out_total_length)
-    audio_out_ids_start: Optional[torch.LongTensor]  # shape (num_audio_out,)
-    audio_out_ids_start_group_loc: Optional[torch.LongTensor]  # shape (num_audio_out,), specify which a sample's group location in the batch
-    audio_in_ids: Optional[torch.LongTensor]  # shape (num_codebooks, audio_in_total_length)
-    audio_in_ids_start: Optional[torch.LongTensor]  # shape (num_audio_in,)
-    label_ids: Optional[torch.LongTensor]  # shape (bsz, seq_len)
-    label_audio_ids: Optional[torch.LongTensor]  # shape (num_codebooks, audio_out_total_length)
-    reward: Optional[float] = None
-
-@dataclass
-class ChatMLDatasetSample:
-
-    input_ids: torch.LongTensor  # (seq_len,) Input text tokens
-    label_ids: torch.LongTensor  # (seq_len,) Label IDs
-    audio_ids_concat: torch.LongTensor  # (num_codebooks, audio_seq_len) Concatenated audio tokens
-    audio_ids_start: torch.LongTensor  # (num_audios,) Start index of each audio token in `audio_ids_concat`
-    audio_waveforms_concat: torch.Tensor  # (total_wv_length,) Concatenated audio waveforms
-    audio_waveforms_start: torch.LongTensor  # (num_audios,) Start index of each waveform in `audio_waveforms_concat`
-    audio_sample_rate: torch.Tensor  # (num_audios,) Sampling rate per audio waveform
-    audio_speaker_indices: torch.LongTensor  # (num_audios,) Speaker indices per audio; -1 = unknown
-    audio_label_ids_concat: Optional[torch.LongTensor] = None  # (num_codebooks, audio_seq_len) Optional audio token labels
-    reward: Optional[float] = None  # Optional scalar reward
-
-
-    def num_audios(self):
-        return max(len(self.audio_waveforms_start), len(self.audio_ids_start))
-
-    def get_audio_codes(self, idx):
-        code_start = self.audio_ids_start[idx]
-        if idx < len(self.audio_ids_start) - 1:
-            code_end = self.audio_ids_start[idx + 1]
-        else:
-            code_end = self.audio_ids_concat.shape[-1]
-
-        return self.audio_ids_concat[:, code_start:code_end]
-
-    def get_audio_codes_labels(self, idx):
-        if self.audio_label_ids_concat is None:
-            return None
-        code_start = self.audio_ids_start[idx]
-        if idx < len(self.audio_ids_start) - 1:
-            code_end = self.audio_ids_start[idx + 1]
-        else:
-            code_end = self.audio_ids_concat.shape[-1]
-
-        return self.audio_label_ids_concat[:, code_start:code_end]
-
-    def get_wv(self, idx):
-        wv_start = self.audio_waveforms_start[idx]
-        sr = self.audio_sample_rate[idx]
-        if idx < len(self.audio_waveforms_start) - 1:
-            wv_end = self.audio_waveforms_start[idx + 1]
-        else:
-            wv_end = self.audio_waveforms_concat.shape[-1]
-        return self.audio_waveforms_concat[wv_start:wv_end], sr
-
-class HiggsAudioSampleCollator:
-
-    def __init__(
-        self,
-        audio_in_token_id,
-        audio_out_token_id,
-        pad_token_id,
-        audio_stream_bos_id,
-        audio_stream_eos_id,
-        round_to=8,
-        pad_left=False,
-        return_audio_in_tokens=True,
-        audio_num_codebooks=None,
-        use_delay_pattern=False,
-        disable_audio_codes_transform=False,
-        add_new_bos_eos_for_long_chunk=True,
-        mask_audio_out_token_label=True,
-    ):
-        self.round_to = round_to
-        self.pad_left = pad_left
-        self.audio_in_token_id = audio_in_token_id
-        self.audio_out_token_id = audio_out_token_id
-        self.audio_stream_bos_id = audio_stream_bos_id
-        self.audio_stream_eos_id = audio_stream_eos_id
-        self.pad_token_id = pad_token_id
-        self.return_audio_in_tokens = return_audio_in_tokens
-        self.audio_num_codebooks = audio_num_codebooks
-        self.use_delay_pattern = use_delay_pattern
-
-        self.disable_audio_codes_transform = disable_audio_codes_transform
-        self.add_new_bos_eos_for_long_chunk = add_new_bos_eos_for_long_chunk
-        self.mask_audio_out_token_label = mask_audio_out_token_label
-
-    def _process_and_duplicate_audio_tokens(
-        self, input_ids: torch.Tensor, audio_idx: int, wv: torch.Tensor, labels: Optional[torch.Tensor] = None
-    ) -> Tuple[torch.Tensor, torch.Tensor, int]:
-
-        total_samples = len(wv)
-        num_chunks = math.ceil(total_samples / self.chunk_size_samples)
-
-        if num_chunks <= 1:
-            return input_ids, labels, 1
-
-        audio_token_seq = input_ids[audio_idx - 1 : audio_idx + 2]
-        duplicated_sequence = audio_token_seq.repeat(num_chunks)
-
-        new_input_ids = torch.cat([input_ids[: audio_idx - 1], duplicated_sequence, input_ids[audio_idx + 2 :]])
-
-        new_labels = None
-        if labels is not None:
-            label_seq = labels[audio_idx - 1 : audio_idx + 2]
-            duplicated_labels = label_seq.repeat(num_chunks)
-            new_labels = torch.cat([labels[: audio_idx - 1], duplicated_labels, labels[audio_idx + 2 :]])
-
-        return new_input_ids, new_labels, num_chunks
-
-    def __call__(self, batch: List[ChatMLDatasetSample]):
-
-        label_ids = None
-        label_audio_ids = None
-        if all([ele.label_ids is None for ele in batch]):
-            return_labels = False
-        else:
-            return_labels = True
-
-        processed_batch = batch
-
-        # Get the max sequence length based on processed batch
-        max_seq_length = _ceil_to_nearest(max([len(sample.input_ids) for sample in processed_batch]), self.round_to)
-
-        # Get the ids for audio-in and audio-out for each batch
-        audio_in_ids_l = []
-        audio_out_ids_l = []
-        audio_out_ids_group_loc_l = []
-        audio_in_label_ids_l = None
-        audio_out_label_ids_l = None
-        reward_l = []
-
-        if return_labels:
-            audio_out_no_train_flag = []  # Whether the audio-out data should be trained on or not.
-
-        # Process the audio inputs and outputs
-        for i in range(len(processed_batch)):
-            audio_in_mask = processed_batch[i].input_ids == self.audio_in_token_id
-            audio_out_mask = processed_batch[i].input_ids == self.audio_out_token_id
-            audio_ids = torch.ones_like(processed_batch[i].input_ids)
-            audio_ids[audio_in_mask ^ audio_out_mask] = torch.cumsum(audio_ids[audio_in_mask ^ audio_out_mask], 0) - 1
-            audio_in_ids = audio_ids[audio_in_mask]
-            audio_out_ids = audio_ids[audio_out_mask]
-
-            if return_labels:
-                audio_out_no_train_flag.append(processed_batch[i].label_ids[audio_out_mask] < 0)
-                if self.mask_audio_out_token_label:
-                    processed_batch[i].label_ids[audio_out_mask] = -100
-
-            if self.return_audio_in_tokens:
-                audio_in_ids_l.extend(
-                    [processed_batch[i].get_audio_codes(idx)[: self.audio_num_codebooks, :] for idx in audio_in_ids]
-                )
-                if processed_batch[i].audio_label_ids_concat is not None:
-                    if audio_in_label_ids_l is None:
-                        audio_in_label_ids_l = []
-                    audio_in_label_ids_l.extend(
-                        [
-                            processed_batch[i].get_audio_codes_labels(idx)[: self.audio_num_codebooks, :]
-                            for idx in audio_in_ids
-                        ]
-                    )
-
-            audio_out_ids_l.extend(
-                [processed_batch[i].get_audio_codes(idx)[: self.audio_num_codebooks, :] for idx in audio_out_ids]
-            )
-            audio_out_ids_group_loc_l.append(i)
-            if processed_batch[i].reward is not None:
-                reward_l.append(processed_batch[i].reward)
-
-            if processed_batch[i].audio_label_ids_concat is not None:
-                if audio_out_label_ids_l is None:
-                    audio_out_label_ids_l = []
-                audio_out_label_ids_l.extend(
-                    [
-                        processed_batch[i].get_audio_codes_labels(idx)[: self.audio_num_codebooks, :]
-                        for idx in audio_out_ids
-                    ]
-                )
-
-        if return_labels:
-            audio_out_no_train_flag = torch.cat(audio_out_no_train_flag, dim=0)
-
-        if len(audio_in_ids_l) > 0:
-
-            # I tried to remove the for-loop in original implementation
-            # but to do batching with padding caused problem so I turned it into a list compre.
-            lengths = [seg.shape[1] for seg in audio_in_ids_l]
-            aug_lengths = [length + 2 for length in lengths]
-            audio_in_ids_start = torch.cumsum(
-                torch.tensor([0] + aug_lengths[:-1], dtype=torch.long), dim=0
-            )
-
-            if self.disable_audio_codes_transform:
-                audio_in_ids = torch.cat(audio_in_ids_l, dim=1).long()
-            else:
-                with_tokens = [
-                    torch.cat([
-                        torch.full((seg.shape[0], 1), self.audio_stream_bos_id, dtype=torch.long),
-                        seg,
-                        torch.full((seg.shape[0], 1), self.audio_stream_eos_id, dtype=torch.long),
-                    ], dim=1)
-                    for seg in audio_in_ids_l
-                ]
-
-                if self.use_delay_pattern:
-                    with_tokens = [
-                        build_delay_pattern_mask(
-                            tok.unsqueeze(0),
-                            bos_token_id=self.audio_stream_bos_id,
-                            pad_token_id=self.audio_stream_eos_id
-                        )[0]
-                        for tok in with_tokens
-                    ]
-                audio_in_ids = torch.cat(with_tokens, dim=1).long()
-        else:
-            audio_in_ids = torch.zeros((0, 0), dtype=torch.long)
-            audio_in_ids_start = torch.zeros(0, dtype=torch.long)
-
-        audio_out_ids_start_group_loc = None
-        if len(audio_out_ids_l) > 0:
-            new_audio_out_ids_l = []
-            label_audio_ids_l = []
-            for idx, ele in enumerate(audio_out_ids_l):
-                if self.disable_audio_codes_transform:
-
-                    audio_codes = ele
-                    if return_labels:
-                        label_audio_ids = audio_out_label_ids_l[idx]
-                else:
-                    audio_codes = torch.cat(
-                        [
-                            torch.full((ele.shape[0], 1), self.audio_stream_bos_id, dtype=torch.long),
-                            ele,
-                            torch.full((ele.shape[0], 1), self.audio_stream_eos_id, dtype=torch.long),
-                        ],
-                        dim=1,
-                    )
-                    if return_labels:
-                        label_audio_ids = torch.cat(
-                            [
-                                torch.full((ele.shape[0], 1), -100, dtype=torch.long),
-                                ele,
-                                torch.full((ele.shape[0], 1), self.audio_stream_eos_id, dtype=torch.long),
-                            ],
-                            dim=1,
-                        )
-                    if self.use_delay_pattern:
-                        audio_codes = build_delay_pattern_mask(
-                            audio_codes.unsqueeze(0),
-                            bos_token_id=self.audio_stream_bos_id,
-                            pad_token_id=self.audio_stream_eos_id,
-                        )[0].squeeze(0)
-                        if return_labels:
-                            label_audio_ids = build_delay_pattern_mask(
-                                label_audio_ids.unsqueeze(0),
-                                bos_token_id=-100,
-                                pad_token_id=-100,
-                            )[0].squeeze(0)
-                new_audio_out_ids_l.append(audio_codes)
-
-                if return_labels:
-                    if audio_out_no_train_flag[idx]:
-                        label_audio_ids[:] = -100
-                    label_audio_ids_l.append(label_audio_ids)
-
-            audio_out_ids = torch.cat(new_audio_out_ids_l, dim=1).long()
-            if return_labels:
-                label_audio_ids = torch.cat(label_audio_ids_l, dim=1).long()
-            audio_out_ids_start = torch.cumsum(
-                torch.tensor([0] + [audio_codes.shape[1] for audio_codes in new_audio_out_ids_l[:-1]]), dim=0
-            )
-            audio_out_ids_start_group_loc = torch.tensor(audio_out_ids_group_loc_l, dtype=torch.long)
-        else:
-            audio_out_ids = torch.zeros((0, 0), dtype=torch.long)
-            audio_out_ids_start = torch.zeros(0, dtype=torch.long)
-            if return_labels:
-                label_audio_ids = torch.zeros((0, 0), dtype=torch.long)
-
-        reward = torch.tensor(reward_l, dtype=torch.float32)
-
-        # cleaner and faster implementation
-        def pad_sequence(seq, length, pad_value):
-            if self.pad_left:
-                padding = (length - len(seq), 0)
-            else:
-                padding = (0, length - len(seq))
-            return F.pad(seq, padding, value=pad_value)
-
-        input_ids = torch.stack([
-            pad_sequence(ele.input_ids, max_seq_length, self.pad_token_id)
-            for ele in processed_batch
-        ])
-
-        if return_labels:
-            label_ids = torch.stack([
-                pad_sequence(ele.label_ids, max_seq_length, -100)
-                for ele in processed_batch
-            ])
-
-        attention_mask = torch.stack([
-            pad_sequence(torch.ones_like(ele.input_ids), max_seq_length, 0)
-            for ele in processed_batch
-        ])
-
-        if not self.return_audio_in_tokens:
-            audio_in_ids = None
-            audio_in_ids_start = None
-
-        if self.audio_num_codebooks is not None:
-            if audio_in_ids is not None:
-                audio_in_ids = audio_in_ids[: self.audio_num_codebooks]
-            if audio_out_ids is not None:
-                audio_out_ids = audio_out_ids[: self.audio_num_codebooks]
-            if label_audio_ids is not None:
-                label_audio_ids = label_audio_ids[: self.audio_num_codebooks]
-
-        return HiggsAudioBatchInput(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            audio_out_ids=audio_out_ids,
-            audio_out_ids_start=audio_out_ids_start,
-            audio_out_ids_start_group_loc=audio_out_ids_start_group_loc,
-            audio_in_ids=audio_in_ids,
-            audio_in_ids_start=audio_in_ids_start,
-            label_ids=label_ids,
-            label_audio_ids=label_audio_ids,
-            reward=reward,
-        )

comfy/ldm/higgsv2/tokenizer.py

@@ -1,871 +0,0 @@
-import math
-import torch
-import torch.nn as nn
-from typing import Optional
-import torch.nn.functional as F
-from torch.nn.utils.parametrizations import weight_norm
-import torchaudio
-import numpy as np
-from torch import vmap
-from transformers import AutoModel
-
-def WNConv1d(*args, device = None, dtype = None, operations = None, **kwargs):
-    return weight_norm(operations.Conv1d(*args, **kwargs, device = device, dtype = dtype))
-
-
-def WNConvTranspose1d(*args, device = None, dtype = None, operations = None, **kwargs):
-    return weight_norm(operations.ConvTranspose1d(*args, **kwargs, device = device, dtype = dtype))
-
-
-@torch.jit.script
-def snake(x, alpha):
-    shape = x.shape
-    x = x.reshape(shape[0], shape[1], -1)
-    x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
-    x = x.reshape(shape)
-    return x
-
-
-class Snake1d(nn.Module):
-    def __init__(self, channels, device = None, dtype = None):
-        super().__init__()
-        self.alpha = nn.Parameter(torch.ones(1, channels, 1, device = device, dtype = dtype))
-
-    def forward(self, x):
-        return snake(x, self.alpha)
-
-class DACResidualUnit(nn.Module):
-    def __init__(self, dim: int = 16, dilation: int = 1, device = None, dtype = None, operations = None):
-        super().__init__()
-        pad = ((7 - 1) * dilation) // 2
-        self.block = nn.Sequential(
-            Snake1d(dim, device = device, dtype = dtype),
-            WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad, device = device, dtype = dtype, operations = operations),
-            Snake1d(dim, device = device, dtype = dtype),
-            WNConv1d(dim, dim, kernel_size=1, device = device, dtype = dtype, operations = operations),
-        )
-
-    def forward(self, x):
-        y = self.block(x)
-        pad = (x.shape[-1] - y.shape[-1]) // 2
-        if pad > 0:
-            x = x[..., pad:-pad]
-        return x + y
-
-
-class DACEncoderBlock(nn.Module):
-    def __init__(self, dim: int = 16, stride: int = 1, device = None, dtype = None, operations = None):
-        super().__init__()
-        self.block = nn.Sequential(
-            DACResidualUnit(dim // 2, dilation=1, device = device, dtype = dtype, operations = operations),
-            DACResidualUnit(dim // 2, dilation=3, device = device, dtype = dtype, operations = operations),
-            DACResidualUnit(dim // 2, dilation=9, device = device, dtype = dtype, operations = operations),
-            Snake1d(dim // 2),
-            WNConv1d(
-                dim // 2,
-                dim,
-                kernel_size=2 * stride,
-                stride=stride,
-                padding=math.ceil(stride / 2),
-                device = device, dtype = dtype, operations = operations
-            ),
-        )
-
-    def forward(self, x):
-        return self.block(x)
-
-
-class DACEncoder(nn.Module):
-    def __init__(
-        self,
-        d_model: int = 64,
-        strides: list = [2, 4, 8, 8],
-        d_latent: int = 256,
-        device = None, dtype = None, operations = None
-    ):
-        super().__init__()
-        # Create first convolution
-        self.block = [WNConv1d(1, d_model, kernel_size=7, padding=3, device = device, dtype = dtype, operations = operations)]
-
-        # Create EncoderBlocks that double channels as they downsample by `stride`
-        for stride in strides:
-            d_model *= 2
-            self.block += [DACEncoderBlock(d_model, stride=stride, device = device, dtype = dtype, operations = operations)]
-
-        # Create last convolution
-        self.block += [
-            Snake1d(d_model),
-            WNConv1d(d_model, d_latent, kernel_size=3, padding=1, device = device, dtype = dtype, operations = operations),
-        ]
-
-        # Wrap black into nn.Sequential
-        self.block = nn.Sequential(*self.block)
-        self.enc_dim = d_model
-
-    def forward(self, x):
-        return self.block(x)
-
-
-class DACDecoderBlock(nn.Module):
-    def __init__(self, input_dim: int = 16, output_dim: int = 8, stride: int = 1, device = None, dtype = None, operations = None):
-        super().__init__()
-        self.block = nn.Sequential(
-            Snake1d(input_dim, device = device, dtype = dtype),
-            WNConvTranspose1d(
-                input_dim,
-                output_dim,
-                kernel_size=2 * stride,
-                stride=stride,
-                padding=math.ceil(stride / 2),
-                output_padding=stride % 2,  # out_pad,
-                device = device, dtype = dtype, operations = operations
-            ),
-            DACResidualUnit(output_dim, dilation=1, device = device, dtype = dtype, operations = operations),
-            DACResidualUnit(output_dim, dilation=3, device = device, dtype = dtype, operations = operations),
-            DACResidualUnit(output_dim, dilation=9, device = device, dtype = dtype, operations = operations),
-        )
-
-    def forward(self, x):
-        return self.block(x)
-
-
-class DACDecoder(nn.Module):
-    def __init__(
-        self,
-        input_channel,
-        channels,
-        rates,
-        d_out: int = 1,
-        device = None, dtype = None, operations = None
-    ):
-        super().__init__()
-
-        # Add first conv layer
-        layers = [WNConv1d(input_channel, channels, kernel_size=7, padding=3, device = device, dtype = dtype, operations = operations )]
-
-        # Add upsampling + MRF blocks
-        for i, stride in enumerate(rates):
-            input_dim = channels // 2**i
-            output_dim = channels // 2 ** (i + 1)
-            layers += [DACDecoderBlock(input_dim, output_dim, stride, device = device, dtype = dtype, operations = operations)]
-
-        # Add final conv layer
-        layers += [
-            Snake1d(output_dim, device = device, dtype = dtype),
-            WNConv1d(output_dim, d_out, kernel_size=7, padding=3, device = device, dtype = dtype, operations = operations),
-        ]
-
-        self.model = nn.Sequential(*layers)
-
-    def forward(self, x):
-        return self.model(x)
-
-class Conv1d1x1:
-    def __new__(cls, in_channels, out_channels, bias=True, device=None, dtype=None, operations=None):
-        operations = operations or nn
-        return operations.Conv1d(
-            in_channels, out_channels, kernel_size=1,
-            bias=bias, device=device, dtype=dtype
-        )
-
-class Conv1d(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        kernel_size: int,
-        stride: int = 1,
-        padding: int = -1,
-        dilation: int = 1,
-        groups: int = 1,
-        bias: bool = True,
-        device = None, dtype = None, operations = None
-    ):
-        super().__init__()
-
-        if padding < 0:
-            padding = (kernel_size - 1) // 2 * dilation
-
-        self.dilation = dilation
-        self.conv = operations.Conv1d(
-            in_channels=in_channels,
-            out_channels=out_channels,
-            kernel_size=kernel_size,
-            stride=stride,
-            padding=padding,
-            dilation=dilation,
-            groups=groups,
-            bias=bias,
-            device = device, dtype = dtype
-        )
-
-    def forward(self, x):
-        x = self.conv(x)
-        return x
-
-class ConvTranspose1d(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        kernel_size: int,
-        stride: int,
-        padding=-1,
-        output_padding=-1,
-        groups=1,
-        bias=True,
-        device = None, dtype = None, operations = None
-    ):
-        super().__init__()
-        if padding < 0:
-            padding = (stride + 1) // 2
-        if output_padding < 0:
-            output_padding = 1 if stride % 2 else 0
-        self.deconv = operations.ConvTranspose1d(
-            in_channels=in_channels,
-            out_channels=out_channels,
-            kernel_size=kernel_size,
-            stride=stride,
-            padding=padding,
-            output_padding=output_padding,
-            groups=groups,
-            bias=bias,
-            device = device, dtype = dtype
-        )
-
-    def forward(self, x):
-        x = self.deconv(x)
-        return x
-
-class ResidualUnit(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        kernel_size=3,
-        dilation=1,
-        bias=False,
-        nonlinear_activation="ELU",
-        nonlinear_activation_params={},
-        device = None, dtype = None, operations = None
-    ):
-        super().__init__()
-        self.activation = getattr(nn, nonlinear_activation)(**nonlinear_activation_params)
-        self.conv1 = Conv1d(
-            in_channels=in_channels,
-            out_channels=out_channels,
-            kernel_size=kernel_size,
-            stride=1,
-            dilation=dilation,
-            bias=bias,
-            device = device, dtype = dtype, operations = operations
-        )
-        self.conv2 = Conv1d1x1(out_channels, out_channels, bias, device = device, dtype = dtype, operations = operations)
-
-    def forward(self, x):
-        y = self.conv1(self.activation(x))
-        y = self.conv2(self.activation(y))
-        return x + y
-
-
-class EncoderBlock(nn.Module):
-    def __init__(
-        self, in_channels: int, out_channels: int, stride: int, dilations=(1, 1), unit_kernel_size=3, bias=True, device = None, dtype = None, operations = None
-    ):
-        super().__init__()
-        self.res_units = torch.nn.ModuleList()
-        for dilation in dilations:
-            self.res_units += [ResidualUnit(in_channels, in_channels, kernel_size=unit_kernel_size, dilation=dilation, device = device, dtype = dtype, operations = operations)]
-        self.num_res = len(self.res_units)
-
-        kernel_size=3 if stride == 1 else (2 * stride) # special case: stride=1, do not use kernel=2
-        self.conv = Conv1d(
-            in_channels=in_channels,
-            out_channels=out_channels,
-            kernel_size = kernel_size,
-            stride=stride,
-            bias=bias,
-            device = device, dtype = dtype, operations = operations
-        )
-
-    def forward(self, x):
-        for idx in range(self.num_res):
-            x = self.res_units[idx](x)
-        x = self.conv(x)
-        return x
-
-
-class Encoder(nn.Module):
-    def __init__(
-        self,
-        input_channels: int,
-        encode_channels: int,
-        channel_ratios=(1, 1),
-        strides=(1, 1),
-        kernel_size=3,
-        bias=True,
-        block_dilations=(1, 1),
-        unit_kernel_size=3,
-        device = None, dtype = None, operations = None
-    ):
-        super().__init__()
-        assert len(channel_ratios) == len(strides)
-        self.conv = Conv1d(
-            in_channels=input_channels, out_channels=encode_channels, kernel_size=kernel_size, stride=1, bias=False,
-            device = device, dtype = dtype, operations = operations
-        )
-        self.conv_blocks = torch.nn.ModuleList()
-        in_channels = encode_channels
-        for idx, stride in enumerate(strides):
-            out_channels = int(encode_channels * channel_ratios[idx])  # could be float
-            self.conv_blocks += [
-                EncoderBlock(
-                    in_channels,
-                    out_channels,
-                    stride,
-                    dilations=block_dilations,
-                    unit_kernel_size=unit_kernel_size,
-                    bias=bias,
-                    device = device, dtype = dtype, operations = operations
-                )
-            ]
-            in_channels = out_channels
-        self.num_blocks = len(self.conv_blocks)
-        self.out_channels = out_channels
-
-    def forward(self, x):
-        x = self.conv(x)
-        for i in range(self.num_blocks):
-            x = self.conv_blocks[i](x)
-        return x
-
-
-class DecoderBlock(nn.Module):
-    """Decoder block (no up-sampling)"""
-
-    def __init__(
-        self, in_channels: int, out_channels: int, stride: int, dilations=(1, 1), unit_kernel_size=3, bias=True, device = None, dtype = None, operations = None
-    ):
-        super().__init__()
-
-        if stride == 1:
-            self.conv = Conv1d(
-                in_channels=in_channels,
-                out_channels=out_channels,
-                kernel_size=3,  # fix kernel=3 when stride=1 for unchanged shape
-                stride=stride,
-                bias=bias,
-                device = device, dtype = dtype, operations = operations
-            )
-        else:
-            self.conv = ConvTranspose1d(
-                in_channels=in_channels,
-                out_channels=out_channels,
-                kernel_size=(2 * stride),
-                stride=stride,
-                bias=bias,
-                device = device, dtype = dtype, operations = operations
-            )
-
-        self.res_units = nn.ModuleList([
-            ResidualUnit(out_channels, out_channels, kernel_size=unit_kernel_size, dilation=d, device = device, dtype = dtype, operations = operations)
-            for d in dilations
-        ])
-
-        self.num_res = len(self.res_units)
-
-    def forward(self, x):
-        x = self.conv(x)
-        for idx in range(self.num_res):
-            x = self.res_units[idx](x)
-        return x
-
-
-class Decoder(nn.Module):
-    def __init__(
-        self,
-        code_dim: int,
-        output_channels: int,
-        decode_channels: int,
-        channel_ratios=(1, 1),
-        strides=(1, 1),
-        kernel_size=3,
-        bias=True,
-        block_dilations=(1, 1),
-        unit_kernel_size=3,
-        device = None, dtype = None, operations = None
-    ):
-        super().__init__()
-        assert len(channel_ratios) == len(strides)
-        self.conv1 = Conv1d(
-            in_channels=code_dim,
-            out_channels=int(decode_channels * channel_ratios[0]),
-            kernel_size=kernel_size,
-            stride=1,
-            bias=False,
-            device = device, dtype = dtype, operations = operations
-        )
-
-        self.conv_blocks = torch.nn.ModuleList()
-        for idx, stride in enumerate(strides):
-            in_channels = int(decode_channels * channel_ratios[idx])
-            if idx < (len(channel_ratios) - 1):
-                out_channels = int(decode_channels * channel_ratios[idx + 1])
-            else:
-                out_channels = decode_channels
-            self.conv_blocks += [
-                DecoderBlock(
-                    in_channels,
-                    out_channels,
-                    stride,
-                    dilations=block_dilations,
-                    unit_kernel_size=unit_kernel_size,
-                    bias=bias,
-                    device = device, dtype = dtype, operations = operations
-                )
-            ]
-        self.num_blocks = len(self.conv_blocks)
-
-        self.conv2 = Conv1d(out_channels, output_channels, kernel_size = 3, bias=False, device = device, dtype = dtype, operations = operations)
-
-    def forward(self, z):
-        x = self.conv1(z)
-        for i in range(self.num_blocks):
-            x = self.conv_blocks[i](x)
-        x = self.conv2(x)
-        return x
-
-class HiggsAudioFeatureExtractor(nn.Module):
-    def __init__(self, sampling_rate=16000):
-        super().__init__()
-        self.sampling_rate = sampling_rate
-
-    def forward(self, audio_signal):
-        audio_signal = audio_signal.unsqueeze(0)
-        if len(audio_signal.shape) < 3:
-            audio_signal = audio_signal.unsqueeze(0)
-        return {"input_values": audio_signal}
-
-def uniform_init(*shape: int, device = None, dtype = None):
-    t = torch.empty(shape, device = device, dtype = dtype)
-    nn.init.kaiming_uniform_(t)
-    return t
-
-class EuclideanCodebook(nn.Module):
-
-    def __init__(
-        self,
-        dim: int,
-        codebook_size: int,
-        kmeans_init: int = False,
-        kmeans_iters: int = 10,
-        decay: float = 0.99,
-        epsilon: float = 1e-5,
-        threshold_ema_dead_code: int = 2,
-        device = None, dtype = None
-    ):
-        super().__init__()
-        self.decay = decay
-        init_fn = uniform_init
-        embed = init_fn(codebook_size, dim, device = device, dtype = dtype)
-
-        self.codebook_size = codebook_size
-
-        self.kmeans_iters = kmeans_iters
-        self.epsilon = epsilon
-        self.threshold_ema_dead_code = threshold_ema_dead_code
-
-        # Flag variable to indicate whether the codebook is initialized
-        self.register_buffer("inited", torch.Tensor([not kmeans_init]))
-        # Runing EMA cluster size/count: N_i^t in eq. (6) in vqvae paper
-        self.register_buffer("cluster_size", torch.zeros(codebook_size))
-        # Codebook
-        self.register_buffer("embed", embed)
-        # EMA codebook: eq. (7) in vqvae paper
-        self.register_buffer("embed_avg", embed.clone())
-
-    def preprocess(self, x):
-        x = x.view(-1, x.shape[-1])
-        return x
-
-    def quantize(self, x):
-        embed = self.embed.t()
-        if x.dtype != embed.dtype:
-            x = x.to(embed.dtype)
-
-        dist = -(x.pow(2).sum(1, keepdim=True) - 2 * x @ embed + embed.pow(2).sum(0, keepdim=True))
-        embed_ind = dist.max(dim=-1).indices
-        return embed_ind
-
-    def postprocess_emb(self, embed_ind, shape):
-        return embed_ind.view(*shape[:-1])
-
-    def dequantize(self, embed_ind):
-        quantize = F.embedding(embed_ind, self.embed)
-        return quantize
-
-    def encode(self, x):
-        shape = x.shape
-        # pre-process
-        x = self.preprocess(x)  # [B, T, D] -> [B*T, D]
-        # quantize
-        embed_ind = self.quantize(x)
-        # post-process
-        embed_ind = self.postprocess_emb(embed_ind, shape)
-        return embed_ind
-
-    def decode(self, embed_ind):
-        quantize = self.dequantize(embed_ind)
-        return quantize
-
-    def forward(self, x):
-        orig_shape = x.shape  # [B, T, D]
-        flat = x.view(-1, x.shape[-1]) # [B*T, D]
-
-        embed_ind = self.quantize(flat)
-        embed_ind = self.postprocess_emb(embed_ind, orig_shape)
-        # now embed_ind has shape [B, T]
-
-        quantize = self.dequantize(embed_ind)
-        # quantize: [B, T, D]
-
-        return quantize, embed_ind
-
-class VectorQuantization(nn.Module):
-
-    def __init__(
-        self,
-        dim: int,
-        codebook_size: int,
-        codebook_dim: Optional[int] = None,
-        decay: float = 0.99,
-        epsilon: float = 1e-5,
-        kmeans_init: bool = True,
-        kmeans_iters: int = 50,
-        threshold_ema_dead_code: int = 2,
-        commitment_weight: float = 1.0,
-        device = None, dtype = None, operations = None
-    ):
-        super().__init__()
-        _codebook_dim: int = codebook_dim if codebook_dim is not None else dim
-
-        requires_projection = _codebook_dim != dim
-        self.project_in = operations.Linear(dim, _codebook_dim, device = device, dtype = dtype) if requires_projection else nn.Identity()
-        self.project_out = operations.Linear(_codebook_dim, dim, device = device, dtype = dtype) if requires_projection else nn.Identity()
-
-        self.epsilon = epsilon
-        self.commitment_weight = commitment_weight
-
-        self._codebook = EuclideanCodebook(
-            dim=_codebook_dim,
-            codebook_size=codebook_size,
-            kmeans_init=kmeans_init,
-            kmeans_iters=kmeans_iters,
-            decay=decay,
-            epsilon=epsilon,
-            threshold_ema_dead_code=threshold_ema_dead_code,
-            device = device, dtype = dtype
-        )
-        self.codebook_size = codebook_size
-
-    @property
-    def codebook(self):
-        return self._codebook.embed
-
-    def encode(self, x):
-        x = x.permute(0, 2, 1)
-        x = self.project_in(x)
-        embed_in = self._codebook.encode(x)
-        return embed_in
-
-    def decode(self, embed_ind):
-        quantize = self._codebook.decode(embed_ind)
-        quantize = self.project_out(quantize)
-        quantize = quantize.permute(0, 2, 1)
-        return quantize
-
-    def forward(self, x):
-        device = x.device
-        x = x.transpose(1, 2).contiguous()  # [b d n] -> [b n d]
-        x = self.project_in(x)
-
-        quantize, embed_ind = self._codebook(x)
-
-        loss = torch.tensor([0.0], device=device, requires_grad=self.training)
-
-        quantize = self.project_out(quantize)
-        quantize = quantize.transpose(1, 2).contiguous()  # [b n d] -> [b d n]
-        return quantize, embed_ind, loss
-
-
-class ResidualVectorQuantization(nn.Module):
-    def __init__(self, *, num_quantizers, device = None, dtype = None, operations = None, **kwargs):
-        super().__init__()
-        self.layers = nn.ModuleList([VectorQuantization(device = device, dtype = dtype, operations = operations, **kwargs) for _ in range(num_quantizers)])
-
-    def forward(self, x, n_q: Optional[int] = None):
-        quantized_out = 0.0
-        residual = x
-
-        all_losses = []
-        all_indices = []
-
-        n_q = n_q or len(self.layers)
-
-        for layer in self.layers[:n_q]:
-            quantized, indices, loss = layer(residual)
-            residual = residual - quantized
-            quantized_out = quantized_out + quantized
-
-            all_indices.append(indices)
-            all_losses.append(loss)
-
-        out_losses, out_indices = map(torch.stack, (all_losses, all_indices))
-        return quantized_out, out_indices, out_losses
-
-    def decode(self, q_indices: torch.Tensor) -> torch.Tensor:
-        """ Vectorized Implementation of dequantization | 2x faster than original impl """
-
-        biases = torch.stack([layer.project_out.bias for layer in self.layers])
-
-        codebook_device = self.layers[0]._codebook.embed.device
-        q_indices = q_indices.to(codebook_device)
-
-        def decode_one(codebook_weight, proj_weight, embed_id, proj_biases):
-            quantized = F.embedding(embed_id, codebook_weight).transpose(1, 2)  # (B, D, T)
-            quantized = F.linear(quantized.transpose(1, 2), proj_weight, proj_biases).transpose(1, 2)
-            return quantized
-
-        codebook_weights = torch.stack([q._codebook.embed for q in self.layers])       # (n_codebooks, vocab_size, D)
-        proj_weights = torch.stack([q.project_out.weight for q in self.layers])
-
-        quantized = vmap(decode_one)(codebook_weights, proj_weights, q_indices, biases)
-
-        return quantized.sum(0)
-
-class ResidualVectorQuantizer(nn.Module):
-
-    def __init__(
-        self,
-        dimension: int = 256,
-        codebook_dim: int = None,
-        n_q: int = 8,
-        bins: int = 1024,
-        decay: float = 0.99,
-        kmeans_init: bool = True,
-        kmeans_iters: int = 50,
-        threshold_ema_dead_code: int = 2,
-        device = None,
-        dtype = None,
-        operations = None
-    ):
-        super().__init__()
-        self.n_q = n_q
-        self.dimension = dimension
-        self.codebook_dim = codebook_dim
-        self.bins = bins
-        self.decay = decay
-        self.kmeans_init = kmeans_init
-        self.kmeans_iters = kmeans_iters
-        self.threshold_ema_dead_code = threshold_ema_dead_code
-        self.vq = ResidualVectorQuantization(
-            dim=self.dimension,
-            codebook_dim=self.codebook_dim,
-            codebook_size=self.bins,
-            num_quantizers=self.n_q,
-            decay=self.decay,
-            kmeans_init=self.kmeans_init,
-            kmeans_iters=self.kmeans_iters,
-            threshold_ema_dead_code=self.threshold_ema_dead_code,
-            device = device, dtype = dtype, operations = operations
-        )
-
-    def forward(self, x: torch.Tensor, sample_rate: int, bandwidth: Optional[float] = None):  # -> QuantizedResult:
-
-        bw_per_q = self.get_bandwidth_per_quantizer(sample_rate)
-        n_q = self.get_num_quantizers_for_bandwidth(sample_rate, bandwidth)
-        quantized, codes, commit_loss = self.vq(x, n_q=n_q)
-        bw = torch.tensor(n_q * bw_per_q).to(x)
-        return quantized, codes, bw, torch.mean(commit_loss)
-
-    def get_num_quantizers_for_bandwidth(self, sample_rate: int, bandwidth: Optional[float] = None) -> int:
-        """Return n_q based on specified target bandwidth."""
-        bw_per_q = self.get_bandwidth_per_quantizer(sample_rate)
-        n_q = self.n_q
-        if bandwidth and bandwidth > 0.0:
-            n_q = int(max(1, math.floor(bandwidth / bw_per_q)))
-        return n_q
-
-    def get_bandwidth_per_quantizer(self, sample_rate: int):
-        """Return bandwidth per quantizer for a given input sample rate."""
-        return math.log2(self.bins) * sample_rate / 1000
-
-    def decode(self, codes: torch.Tensor) -> torch.Tensor:
-        """Decode the given codes to the quantized representation."""
-        quantized = self.vq.decode(codes)
-        return quantized
-
-class HiggsAudioTokenizer(nn.Module):
-    def __init__(
-        self,
-        D: int = 256,
-        target_bandwidths= [0.5, 1, 1.5, 2, 4],
-        ratios = [8, 5, 4, 2, 3],  #  downsampling by 320
-        sample_rate: int = 24000,
-        bins: int = 1024,
-        n_q: int = 8,
-        codebook_dim: int = 64,
-        last_layer_semantic: bool = True,
-        downsample_mode: str = "step_down",
-        vq_scale: int = 1,
-        semantic_sample_rate: int = None,
-        device = None,
-        dtype = None,
-        operations = None,
-        **kwargs
-    ):
-        super().__init__()
-        operations = operations or nn
-        self.hop_length = np.prod(ratios)
-
-        self.frame_rate = math.ceil(sample_rate / np.prod(ratios))  # 50 Hz
-
-        self.target_bandwidths = target_bandwidths
-        self.n_q = n_q
-        self.sample_rate = sample_rate
-        self.encoder = DACEncoder(64, ratios, D, device = device, dtype = dtype, operations = operations)
-
-        self.decoder_2 = DACDecoder(D, 1024, ratios, device = device, dtype = dtype, operations = operations)
-        self.last_layer_semantic = last_layer_semantic
-        self.device = device
-
-        self.semantic_model = AutoModel.from_pretrained("bosonai/hubert_base", trust_remote_code=True)
-        self.semantic_sample_rate = 16000
-        self.semantic_dim = 768
-        self.encoder_semantic_dim = 768
-
-        # Overwrite semantic model sr to ensure semantic_downsample_factor is an integer
-        if semantic_sample_rate is not None:
-            self.semantic_sample_rate = semantic_sample_rate
-
-        self.semantic_model.eval()
-
-        # make the semantic model parameters do not need gradient
-        for param in self.semantic_model.parameters():
-            param.requires_grad = False
-
-        self.semantic_downsample_factor = int(self.hop_length / (self.sample_rate / self.semantic_sample_rate) / 320)
-
-        self.quantizer_dim = int((D + self.encoder_semantic_dim) // vq_scale)
-        self.encoder_semantic = Encoder(input_channels=self.semantic_dim, encode_channels=self.encoder_semantic_dim, device = device, dtype = dtype, operations = operations)
-        self.decoder_semantic = Decoder(
-            code_dim=self.encoder_semantic_dim, output_channels=self.semantic_dim, decode_channels=self.semantic_dim, device = device, dtype = dtype, operations = operations
-        )
-
-        self.quantizer = ResidualVectorQuantizer(
-            dimension=self.quantizer_dim, codebook_dim=codebook_dim, n_q=n_q, bins=bins, device = device, dtype = dtype, operations = operations
-        )
-
-        self.fc_prior = operations.Linear(D + self.encoder_semantic_dim, self.quantizer_dim, device = device, dtype = dtype)
-        self.fc_post1 = operations.Linear(self.quantizer_dim, self.encoder_semantic_dim, device = device, dtype = dtype)
-        self.fc_post2 = operations.Linear(self.quantizer_dim, D, device = device, dtype = dtype)
-
-        self.downsample_mode = downsample_mode
-
-        self.audio_tokenizer_feature_extractor = HiggsAudioFeatureExtractor(sampling_rate=self.sample_rate)
-
-    @property
-    def sampling_rate(self):
-        return self.sample_rate
-
-    @torch.no_grad()
-    def get_regress_target(self, x):
-        x = torchaudio.functional.resample(x, self.sample_rate, self.semantic_sample_rate)
-
-        x = x[:, 0, :]
-        x = F.pad(x, (160, 160))
-        target = self.semantic_model(x, output_hidden_states=True).hidden_states
-        target = torch.stack(target, dim=1)
-
-        target = target.mean(1)
-
-        if self.downsample_mode == "step_down":
-            if self.semantic_downsample_factor > 1:
-                target = target[:, :: self.semantic_downsample_factor, :]
-
-        return target
-
-    def forward(self):
-        pass
-
-    @property
-    def tps(self):
-        return self.frame_rate
-
-    def encode(self, wv, sr):
-
-        if sr != self.sampling_rate:
-            # best computed values to match librosa's resample
-            resampler_torch = torchaudio.transforms.Resample(
-                orig_freq=sr,
-                new_freq=self.sampling_rate,
-                resampling_method="sinc_interp_kaiser",
-                lowpass_filter_width = 121,
-                rolloff = 0.9568384289091556,
-                beta = 21.01531462440614
-            ).to(wv.device)
-
-            wv = resampler_torch(wv)
-
-        if self.audio_tokenizer_feature_extractor is not None:
-            inputs = self.audio_tokenizer_feature_extractor(wv)
-            input_values = inputs["input_values"].to(self.device)
-        else:
-            input_values = torch.from_numpy(wv).float().unsqueeze(0)
-        with torch.no_grad():
-            input_values = input_values.to(wv.device)
-            encoder_outputs = self._xcodec_encode(input_values)
-            vq_code = encoder_outputs[0]
-        return vq_code
-
-    def _xcodec_encode(self, x: torch.Tensor, target_bw: Optional[int] = None) -> torch.Tensor:
-        bw = target_bw
-
-        e_semantic_input = self.get_regress_target(x).detach()
-
-        e_semantic = self.encoder_semantic(e_semantic_input.transpose(1, 2))
-        e_acoustic = self.encoder(x)
-
-        if e_acoustic.shape[2] != e_semantic.shape[2]:
-            pad_size = 160 * self.semantic_downsample_factor
-            e_acoustic = self.encoder(F.pad(x[:, 0, :], (pad_size, pad_size)).unsqueeze(0))
-
-        if e_acoustic.shape[2] != e_semantic.shape[2]:
-            if e_acoustic.shape[2] > e_semantic.shape[2]:
-                e_acoustic = e_acoustic[:, :, : e_semantic.shape[2]]
-            else:
-                e_semantic = e_semantic[:, :, : e_acoustic.shape[2]]
-
-        e = torch.cat([e_acoustic, e_semantic], dim=1)
-
-        e = self.fc_prior(e.transpose(1, 2))
-
-        e = e.transpose(1, 2)
-        _, codes, _, _ = self.quantizer(e, self.frame_rate, bw)
-        codes = codes.permute(1, 0, 2)
-
-        return codes
-
-    def decode(self, vq_code: torch.Tensor) -> torch.Tensor:
-        vq_code = vq_code.to(self.device)
-
-        if vq_code.ndim < 3:
-            vq_code = vq_code.unsqueeze(0)
-
-        vq_code = vq_code.permute(1, 0, 2)
-        quantized = self.quantizer.decode(vq_code)
-        quantized = quantized.transpose(1, 2)
-        quantized_acoustic = self.fc_post2(quantized).transpose(1, 2)
-
-        o = self.decoder_2(quantized_acoustic)
-        return o.detach()

comfy/model_base.py

@@ -44,7 +44,6 @@ import comfy.ldm.hidream.model
 import comfy.ldm.chroma.model
 import comfy.ldm.ace.model
 import comfy.ldm.omnigen.omnigen2
-import comfy.ldm.higgsv2.model
 import comfy.ldm.qwen_image.model
 import comfy.ldm.hunyuan_foley.model
 
@@ -1386,10 +1385,6 @@ class Omnigen2(BaseModel):
             out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
         return out
 
-class Higgsv2(BaseModel):
-    def __init__(self, model_config, model_type=ModelType.EPS, device=None, unet_model=comfy.ldm.higgsv2.model.HiggsAudioModel):
-        super().__init__(model_config, model_type, device, unet_model)
-
 class QwenImage(BaseModel):
     def __init__(self, model_config, model_type=ModelType.FLUX, device=None):
         super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.qwen_image.model.QwenImageTransformer2DModel)

comfy/model_detection.py

@@ -417,73 +417,6 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
         dit_config["guidance_cond_proj_dim"] = None#f"{key_prefix}t_embedder.cond_proj.weight" in state_dict_keys
         return dit_config
 
-    if "{}layers.27.audio_post_attention_layernorm.weight".format(key_prefix) in state_dict_keys:
-
-        autoregressive_config = {}
-        autoregressive_config["image_model"] = "higgsv2"
-
-        autoregressive_config["audio_adapter_type"] = "dual_ffn_fast_forward"
-        autoregressive_config["audio_bos_token"] = "<|audio_bos|>"
-        autoregressive_config["audio_codebook_size"] = 1024
-        autoregressive_config["audio_num_codebooks"] = 8
-        autoregressive_config["audio_ffn_hidden_size"] = 3072
-        autoregressive_config["audio_ffn_intermediate_size"] = 8192
-        autoregressive_config["audio_in_token"] = "<|AUDIO|>"
-        autoregressive_config["audio_in_token_idx"] = 128015
-        autoregressive_config["audio_out_token"] = "<|AUDIO_OUT|>"
-        autoregressive_config["audio_out_token_idx"] = 128016
-        autoregressive_config["audio_out_bos_token"] = "<|audio_out_bos|>"
-        autoregressive_config["audio_out_bos_token_id"] = 128013
-        autoregressive_config["audio_eos_token"] = "<|audio_eos|>"
-        autoregressive_config["audio_eos_token_id"] = 128012
-        autoregressive_config["audio_stream_bos_id"] = 1024
-        autoregressive_config["audio_stream_eos_id"] = 1025
-        autoregressive_config["encode_audio_in_tokens"] = True
-
-        autoregressive_config["pad_token_id"] = 128001
-        autoregressive_config["padding_idx"] = 128001
-
-        autoregressive_config["hidden_size"] = 3072
-        autoregressive_config["use_delay_pattern"] = True
-
-        autoregressive_config["vocab_size"] = 128256
-        autoregressive_config["num_hidden_layers"] = 28
-        autoregressive_config["num_attention_heads"] = 24
-        autoregressive_config["num_key_value_heads"] = 8
-        autoregressive_config["max_seq_len"] = 131072
-        autoregressive_config["max_position_embeddings"] = 131072
-        autoregressive_config["bos_token_id"] = 128000
-        autoregressive_config["eos_token_id"] = 128001
-        autoregressive_config["use_cache"] = True
-
-        autoregressive_config["text_config"] = {
-            "model_type": "llama",
-            "vocab_size": 128256,
-            "max_position_embeddings": 131072,
-            "num_hidden_layers": 28,
-            "hidden_size": 3072,
-            "num_attention_heads": 24,
-            "num_key_value_heads": 8,
-            "initializer_range": 0.02,
-            "rms_norm_eps": 1e-05,
-            "pad_token_id": None,
-            "bos_token_id": 128000,
-            "eos_token_id": 128001,
-            "num_return_sequences": 1,
-            "head_dim": 128,
-            "mlp_bias": False,
-            "intermediate_size": 8192
-        }
-
-        autoregressive_config["use_kv_buckets"] = True
-        autoregressive_config["num_attention_heads"] = 24
-
-        autoregressive_config["audio_decoder_proj_num_layers"] = 0
-        autoregressive_config["audio_dual_ffn_layers"] = list(range(28))
-        autoregressive_config["output_vae"] = False
-
-        return autoregressive_config
-
     if '{}caption_projection.0.linear.weight'.format(key_prefix) in state_dict_keys:  # HiDream
         dit_config = {}
         dit_config["image_model"] = "hidream"

comfy/sd.py

@@ -1091,7 +1091,6 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c
 
     manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
     model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
-    output_vae = model_config.unet_config.get("output_vae", output_vae)
 
     if model_config.clip_vision_prefix is not None:
         if output_clipvision:

comfy/supported_models.py

@@ -20,7 +20,6 @@ import comfy.text_encoders.lumina2
 import comfy.text_encoders.wan
 import comfy.text_encoders.ace
 import comfy.text_encoders.omnigen2
-import comfy.text_encoders.higgsv2
 import comfy.text_encoders.qwen_image
 
 from . import supported_models_base
@@ -1312,22 +1311,6 @@ class QwenImage(supported_models_base.BASE):
         hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref))
         return supported_models_base.ClipTarget(comfy.text_encoders.qwen_image.QwenImageTokenizer, comfy.text_encoders.qwen_image.te(**hunyuan_detect))
 
-class Higgsv2(supported_models_base.BASE):
+models = [LotusD, Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, PixArtAlpha, PixArtSigma, HunyuanDiT, HunyuanDiT1, FluxInpaint, Flux, FluxSchnell, GenmoMochi, LTXV, HunyuanVideoSkyreelsI2V, HunyuanVideoI2V, HunyuanVideo, CosmosT2V, CosmosI2V, CosmosT2IPredict2, CosmosI2VPredict2, Lumina2, WAN22_T2V, WAN21_T2V, WAN21_I2V, WAN21_FunControl2V, WAN21_Vace, WAN21_Camera, WAN22_Camera, WAN22_S2V, Hunyuan3Dv2mini, Hunyuan3Dv2, Hunyuan3Dv2_1, HunyuanFoley, HiDream, Chroma, ACEStep, Omnigen2, QwenImage]
-    unet_config = {
-        "image_model": "higgsv2",
-    }
-
-    memory_usage_factor = 1.0
-    supported_inference_dtypes = [torch.float16, torch.bfloat16, torch.float32]
-    text_encoder_key_prefix = ["dac."]
-
-    def get_model(self, state_dict, prefix="", device=None):
-        out = model_base.Higgsv2(self, device=device)
-        return out
-
-    def clip_target(self, state_dict = {}):
-        return supported_models_base.ClipTarget(comfy.text_encoders.higgsv2.DummyTokenizer, comfy.text_encoders.higgsv2.HiggsTokenizer)
-
-models = [LotusD, Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, PixArtAlpha, PixArtSigma, HunyuanDiT, HunyuanDiT1, FluxInpaint, Flux, FluxSchnell, GenmoMochi, LTXV, HunyuanVideoSkyreelsI2V, HunyuanVideoI2V, HunyuanVideo, CosmosT2V, CosmosI2V, CosmosT2IPredict2, CosmosI2VPredict2, Lumina2, WAN22_T2V, WAN21_T2V, WAN21_I2V, WAN21_FunControl2V, WAN21_Vace, WAN21_Camera, WAN22_Camera, WAN22_S2V, Hunyuan3Dv2mini, Hunyuan3Dv2, Hunyuan3Dv2_1, HiDream, Chroma, ACEStep, Omnigen2, QwenImage, Higgsv2]
 
 models += [SVD_img2vid]

comfy/text_encoders/higgs_text_tokenizer/special_tokens_map.json

@@ -1,16 +0,0 @@
-{
-  "bos_token": {
-    "content": "<|begin_of_text|>",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
-  "eos_token": {
-    "content": "<|end_of_text|>",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  }
-}

comfy/text_encoders/higgsv2.py

@@ -1,81 +0,0 @@
-import os
-import torch
-import comfy.ops
-import torch.nn as nn
-from transformers import AutoTokenizer
-from comfy.ldm.higgsv2.tokenizer import HiggsAudioTokenizer
-from comfy.ldm.higgsv2.preprocess import HiggsAudioSampleCollator
-
-class DummyTokenizer:
-    def __init__(self, embedding_directory=None, tokenizer_data={}):
-        pass
-
-def revert_delay_pattern_vectorized(data: torch.Tensor) -> torch.Tensor:
-    num_codebooks, total_len = data.shape
-    seq_len = total_len - num_codebooks + 1
-
-    col_idx = torch.arange(seq_len, device=data.device)[None, :] \
-             + torch.arange(num_codebooks, device=data.device)[:, None]
-    out = data[torch.arange(num_codebooks)[:, None], col_idx]
-    return out
-
-class HiggsTokenizer(nn.Module):
-    def __init__(self, device, dtype, model_options={}, **kwargs):
-        super().__init__()
-
-        self.dtype = torch.float32
-        self.device = device
-        self.dtypes = [torch.float32]
-
-        here = os.path.dirname(__file__)
-        tokenizer_path = os.path.join(here, "higgs_text_tokenizer")
-
-        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
-        scaled_fp8 = model_options.get("scaled_fp8", None)
-
-        if scaled_fp8 is not None:
-            operations = comfy.ops.scaled_fp8_ops(fp8_matrix_mult=False, override_dtype=scaled_fp8)
-        else:
-            operations = comfy.ops.manual_cast
-
-        self.audio_codebook_size = 1024
-        self.audio_tokenizer = HiggsAudioTokenizer(device = device, dtype = dtype, operations = operations)
-
-        if scaled_fp8 is not None:
-            self.audio_tokenizer.scaled_fp8 = torch.nn.Parameter(torch.tensor([], dtype=scaled_fp8))
-
-        self.collator = HiggsAudioSampleCollator(
-            audio_in_token_id = 128015,
-            audio_out_token_id = 128016,
-            audio_stream_bos_id = 1024,
-            audio_stream_eos_id = 1025,
-            pad_token_id = 128001,
-            return_audio_in_tokens = False,
-            use_delay_pattern = True,
-            audio_num_codebooks = 8,
-            round_to = 1,
-        )
-
-        postfix = "<|start_header_id|>assistant<|end_header_id|>\n\n"
-        self.postfix = postfix + "<|audio_out_bos|>" # force audio generation
-
-    def decode_tokens(self, audio_tokens):
-        outputs = []
-
-        # due to instability issues, I had to convert the audio tokenizer to float32, avoiding outputing nans
-        self.audio_tokenizer = self.audio_tokenizer.to(self.dtype)
-        torch.cuda.synchronize()
-
-        for audio in audio_tokens:
-            vq_code = revert_delay_pattern_vectorized(audio).clip(0, self.audio_codebook_size - 1)[:, 1:-1]
-            wv_numpy = self.audio_tokenizer.decode(vq_code.unsqueeze(0))[0, 0]
-            outputs.append(wv_numpy)
-
-        # currently only supports one batch size
-        return (None, {"waveform": torch.cat(outputs, dim = 0).unsqueeze(0).unsqueeze(1), "sample_rate": self.audio_tokenizer.sample_rate}) # audio only
-
-    def load_state_dict(self, sd, strict = False):
-        return self.audio_tokenizer.load_state_dict(sd, strict = strict)
-
-    def state_dict(self):
-        return self.audio_tokenizer.state_dict()

comfy/text_encoders/llama.py

@@ -1,9 +1,6 @@
 import torch
 import torch.nn as nn
 from typing import Optional, Any
-from dataclasses import dataclass, field
-from transformers.cache_utils import Cache
-from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
 import math
 
 from comfy.ldm.modules.attention import optimized_attention_for_device
@@ -24,20 +21,9 @@ class Llama2Config:
     rms_norm_eps: float = 1e-5
     rope_theta: float = 500000.0
     transformer_type: str = "llama"
-    head_dim: int = 128
+    head_dim = 128
     rms_norm_add = False
     mlp_activation = "silu"
-    qkv_bias: bool = False
-    rope_type: str = "llama3"
-    rope_scaling: dict = field(
-        default_factory=lambda: {
-            "factor": 32.0,
-            "high_freq_factor": 4.0,
-            "low_freq_factor": 1.0,
-            "original_max_position_embeddings": 8192,
-            "rope_type": "llama3"
-        }
-    )
     qkv_bias = False
     rope_dims = None
 
@@ -145,67 +131,15 @@ def apply_rope(xq, xk, freqs_cis):
     sin = freqs_cis[1]
     q_embed = (xq * cos) + (rotate_half(xq) * sin)
     k_embed = (xk * cos) + (rotate_half(xk) * sin)
-    return q_embed.to(org_dtype), k_embed.to(org_dtype), sin, cos
+    return q_embed.to(org_dtype), k_embed.to(org_dtype)
-
-class LlamaRoPE(nn.Module):
-    def __init__(self, config, device = None, dtype = None):
-        super().__init__()
-
-        if config.rope_scaling is not None:
-            self.rope_type = config.rope_scaling.get("rope_type", config.rope_type)
-        else:
-            self.rope_type = "default"
-
-        self.config = config
-        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
-
-        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-
-    def _dynamic_frequency_update(self, position_ids, device):
-
-        seq_len = torch.max(position_ids) + 1
-        if seq_len > self.max_seq_len_cached:
-            inv_freq, self.attention_scaling = self.rope_init_fn(
-                self.config, device, seq_len=seq_len, **self.rope_kwargs
-            )
-            self.register_buffer("inv_freq", inv_freq, persistent=False)
-            self.max_seq_len_cached = seq_len
-
-        if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len:
-            self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
-            self.max_seq_len_cached = self.original_max_seq_len
-
-    @torch.no_grad()
-    def forward(self, x, position_ids):
-        if "dynamic" in self.rope_type:
-            self._dynamic_frequency_update(position_ids, device=x.device)
-
-        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
-        position_ids_expanded = position_ids[:, None, :].float()
-
-        device_type = x.device.type
-        device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
-        with torch.autocast(device_type=device_type, enabled=False):
-            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
-            emb = torch.cat((freqs, freqs), dim=-1)
-            cos = emb.cos()
-            sin = emb.sin()
-
-        cos = cos * self.attention_scaling
-        sin = sin * self.attention_scaling
-
-        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
 
 
 class Attention(nn.Module):
-    def __init__(self, config: Llama2Config, layer_idx: int = None, device=None, dtype=None, ops: Any = None):
+    def __init__(self, config: Llama2Config, device=None, dtype=None, ops: Any = None):
         super().__init__()
         self.num_heads = config.num_attention_heads
         self.num_kv_heads = config.num_key_value_heads
         self.hidden_size = config.hidden_size
-        self.layer_idx = layer_idx
 
         self.head_dim = config.head_dim
         self.inner_size = self.num_heads * self.head_dim
@@ -220,8 +154,6 @@ class Attention(nn.Module):
         self,
         hidden_states: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
-        past_key_value: Optional[Cache] = None,
-        cache_position: Optional[torch.LongTensor] = None,
         freqs_cis: Optional[torch.Tensor] = None,
         optimized_attention=None,
     ):
@@ -234,22 +166,13 @@ class Attention(nn.Module):
         xk = xk.view(batch_size, seq_length, self.num_kv_heads, self.head_dim).transpose(1, 2)
         xv = xv.view(batch_size, seq_length, self.num_kv_heads, self.head_dim).transpose(1, 2)
 
-        xq, xk, sin, cos = apply_rope(xq, xk, freqs_cis=freqs_cis)
+        xq, xk = apply_rope(xq, xk, freqs_cis=freqs_cis)
-
-        if past_key_value is not None:
-            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
-            xk, xv = past_key_value.update(xk, xv, self.layer_idx, cache_kwargs)
 
         xk = xk.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
         xv = xv.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
 
         output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True)
-        out = self.o_proj(output)
+        return self.o_proj(output)
-
-        if past_key_value is not None:
-            return out, past_key_value
-
-        return out
 
 class MLP(nn.Module):
     def __init__(self, config: Llama2Config, device=None, dtype=None, ops: Any = None):
@@ -478,18 +401,12 @@ class Qwen25_7BVLI(BaseLlama, torch.nn.Module):
         for e in embeds_info:
             if e.get("type") == "image":
                 grid = e.get("extra", None)
-                position_ids = torch.zeros((3, embeds.shape[1]), device=embeds.device)
                 start = e.get("index")
-                position_ids[:, :start] = torch.arange(0, start, device=embeds.device)
                 end = e.get("size") + start
                 len_max = int(grid.max()) // 2
                 start_next = len_max + start
-                position_ids[:, end:] = torch.arange(start_next, start_next + (embeds.shape[1] - end), device=embeds.device)
-                position_ids[0, start:end] = start
                 max_d = int(grid[0][1]) // 2
-                position_ids[1, start:end] = torch.arange(start, start + max_d, device=embeds.device).unsqueeze(1).repeat(1, math.ceil((end - start) / max_d)).flatten(0)[:end - start]
                 max_d = int(grid[0][2]) // 2
-                position_ids[2, start:end] = torch.arange(start, start + max_d, device=embeds.device).unsqueeze(0).repeat(math.ceil((end - start) / max_d), 1).flatten(0)[:end - start]
 
         if grid is None:
             position_ids = None

comfy_extras/nodes_audio.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import av
-import re
 import torchaudio
 import torch
 import comfy.model_management
@@ -11,257 +10,9 @@ import io
 import json
 import random
 import hashlib
-import numpy as np
 import node_helpers
 from comfy.cli_args import args
-from comfy.comfy_types import IO
 from comfy.comfy_types import FileLocator
-from dataclasses import  asdict
-from comfy.ldm.higgsv2.loudness import loudness
-from comfy.ldm.higgsv2.preprocess import (
-    prepare_chatml_sample, Message, ChatMLSample, ChatMLDatasetSample, AudioContent, transcript_normalize
-)
-
-
-MULTISPEAKER_DEFAULT_SYSTEM_MESSAGE = """You are an AI assistant designed to convert text into speech.
-If the user's message includes a [SPEAKER*] tag, do not read out the tag and generate speech for the following text, using the specified voice.
-If no speaker tag is present, select a suitable voice on your own."""
-
-class LoudnessNormalization:
-
-    CATEGORY = "audio"
-    RETURN_TYPES = ("AUDIO",)
-    FUNCTION = "normalize"
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"audio": ("AUDIO", ),
-                             "block_size": ("FLOAT", {"default": 0.400, "min": 0.1, "max": 1.0, "step": 0.05}),
-                             "loudness_threshold": ("FLOAT", {"default": -23.0, "min": -70.0, "max": 0.0, "step": 0.5,
-                                                              "tooltip": "Target loudness in LUFS. Common values are -23.0 (broadcast), -14.0 (streaming)."})}}
-
-    def normalize(self, audio, loudness_threshold, block_size):
-        sampling_rate = audio["sample_rate"]
-        waveform = audio["waveform"]
-        return {"waveform": loudness(waveform, sampling_rate, target_loudness = loudness_threshold, block_size = block_size), "sample_rate": sampling_rate}
-
-def prepare_chatml_input(
-    clip,
-    input_tokens,
-    audio_contents,
-    sampling_rate,
-    postfix_str: str = "",
-):
-    if hasattr(clip, "postfix"):
-        postfix_str = clip.postfix
-
-    if postfix_str:
-        postfix = clip.tokenizer.encode(postfix_str, add_special_tokens=False)
-        input_tokens.extend(postfix)
-
-    audio_ids_l = []
-    if audio_contents is not None:
-
-        if not hasattr(clip, "audio_tokenizer"):
-            raise ValueError("This model does not have an audio tokenizer. The chosen model may not support ChatML Format")
-
-        for audio_content in audio_contents:
-            audio_content.raw_audio = audio_content.raw_audio.squeeze(0)
-            if audio_content.raw_audio.shape[0] == 2:
-                audio_content.raw_audio = audio_content.raw_audio.mean(dim = 0, keepdim = True)
-
-            if audio_content.raw_audio.device != next(clip.audio_tokenizer.parameters()).device:
-                audio_content.raw_audio = audio_content.raw_audio.to(next(clip.audio_tokenizer.parameters()).device)
-
-            audio_ids = clip.audio_tokenizer.encode(audio_content.raw_audio, sampling_rate)
-            audio_ids_l.append(audio_ids.squeeze(0))
-
-    if len(audio_ids_l) > 0:
-        audio_ids_start = torch.tensor(
-            np.cumsum(np.array([0] + [audio_ids.shape[1] for audio_ids in audio_ids_l])),
-            dtype=torch.long,
-            device=audio_contents[0].raw_audio.device,
-        ).to("cpu")[0:-1]
-        audio_ids_concat = torch.cat(audio_ids_l, dim=1).to("cpu")
-    else:
-        audio_ids_start = None
-        audio_ids_concat = None
-
-    sample = ChatMLDatasetSample(
-        input_ids=torch.LongTensor(input_tokens),
-        label_ids=None,
-        audio_ids_concat=audio_ids_concat,
-        audio_ids_start=audio_ids_start,
-        audio_waveforms_concat=None,
-        audio_waveforms_start=None,
-        audio_sample_rate=None,
-        audio_speaker_indices=None,
-    )
-
-    if hasattr(clip, "collator"):
-        sample.input_ids = sample.input_ids.cpu()
-        sample = clip.collator([sample])
-
-    inputs = asdict(sample)
-    for k, v in inputs.items():
-        if isinstance(v, torch.Tensor):
-            inputs[k] = v.to(clip.device)
-
-    return inputs
-
-def postprocess_chatml(text: str) -> str:
-    speakers = set(re.findall(r'\[SPEAKER\d+\]', text))
-    skip_recon = True
-
-    if len(speakers) > 1:
-        parts = text.split('<|eot_id|>')
-
-        # keep the first <|eot_id|> and the last one
-        first_eot = parts[0] + '<|eot_id|>'
-        middle_parts = ''.join(parts[1:-1])
-        last_eot = '<|eot_id|>' + parts[-1]
-
-        text = first_eot + middle_parts + last_eot
-        skip_recon = False
-
-    return text, skip_recon
-
-class CreateChatMLSample:
-    def __init__(self):
-        self.device = comfy.model_management.intermediate_device()
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "text": (IO.STRING, {
-                    "default": "SYSTEM: " + MULTISPEAKER_DEFAULT_SYSTEM_MESSAGE + "\n\n<|scene_desc_start|>\nSPEAKER0:masculine\nSPEAKER1:feminine\n<|scene_desc_end|>",
-                    "multiline": True,
-                    "dynamicPrompts": True,
-                    "tooltip": (
-                        "The conversations to be encoded. "
-                        "To register a conversation start with SPEAKER-0: some text. "
-                        "To add a system prompt start with system:"
-                    ),
-                }),
-                "clip": (IO.CLIP, {"tooltip": "The CLIP model used for tokenizing the text."}),
-            },
-            "optional": {
-                "audio": (IO.AUDIO, {
-                    "tooltip": "An audio clip to be inserted into the conversation. To register add [audio]",
-                })
-            }
-        }
-
-    RETURN_TYPES = ("TOKENS",)
-    OUTPUT_TOOLTIPS = ("Turns text and audio into a ChatML Format.",)
-
-    FUNCTION = "convert_to_ml_format"
-    CATEGORY = "conditioning"
-
-    def convert_to_ml_format(self, clip, text, audio=None):
-
-        if audio is not None:
-            clip.load_model()
-
-        if hasattr(clip, "cond_stage_model"):
-            clip = clip.cond_stage_model
-
-        text = transcript_normalize(text)
-
-        messages = []
-        lines = text.splitlines()
-        sampling_rate = False
-        current_role = None
-        collecting_system = False
-        system_buffer = []
-
-        for line in lines:
-            line = line.strip()
-            if not line:
-                continue
-
-            # system start
-            if line.lower().startswith("system:"):
-                collecting_system = True
-                system_buffer.append(line[len("system:"):].strip())
-                continue
-
-            # while collecting system prompt
-            if collecting_system:
-                system_buffer.append(line)
-                if "<|scene_desc_end|>" in line or "SPEAKER-" in line:
-                    system_prompt = "\n".join(system_buffer)# + "\n<|scene_desc_end|>"
-                    messages.append(Message(role="system", content=system_prompt))
-                    system_buffer = []
-                    collecting_system = False
-                continue
-
-            # speaker lines SPEAKER-0: text
-            match = re.match(r"SPEAKER-(\d+):\s*(.*)", line, re.IGNORECASE)
-            if match:
-                speaker_id = match.group(1)
-                content = match.group(2)
-                current_role = f"[SPEAKER{speaker_id}] "
-                messages.append(Message(role = "user", content = current_role + content.strip()))
-            else:
-                # continuation line goes to last speaker or instruction
-                if current_role is not None and messages:
-                    messages[-1].content += "\n" + line
-
-        # return normal input_ids
-        if not (len(messages) >= 1):
-            return (clip.tokenizer(text),)
-
-        all_text = "".join(msg.content for msg in messages if msg.role == "user")
-
-        # postprocess to allow multi-user speech
-        all_text, skip_recon = postprocess_chatml(all_text)
-        if not skip_recon:
-            lines = all_text.splitlines()
-            messages = [messages[0]] if messages[0].role == "system" else []
-            current_role = None
-
-            for line in lines:
-                line = line.strip()
-                if not line:
-                    continue
-
-                match = re.match(r'(\[SPEAKER\d+\])\s*(.*)', line)
-                if match:
-                    current_role = match.group(1)
-                    content = match.group(2).strip()  # only take the text after the tag
-                    messages.append(Message(role="user", content=f"{current_role} {content}" if content else current_role))
-                else:
-                    if current_role and messages:
-                        messages[-1].content += "\n" + line
-
-            # dedepulicate the messages
-            for idx, m in enumerate(messages):
-                double_eot = "<|eot_id|><|eot_id|>"
-                if double_eot in m.content:
-                    cut_index = m.content.index(double_eot)
-                    messages[idx].content = m.content[:cut_index + (len(double_eot) // 2)]
-                    break
-
-        if audio is not None:
-            # for audio cloning, the first message is a transcript, second is the audio,
-            # third is the request of what the model should say
-            waveform = audio["waveform"]
-            sampling_rate = audio["sample_rate"]
-            messages.insert(1, Message(
-                role = "assistant",
-                content = AudioContent(raw_audio = waveform, audio_url = "placeholder")
-            ))
-        chat_ml_sample = ChatMLSample(messages)
-        input_tokens, audio_contents, _ = prepare_chatml_sample(
-            chat_ml_sample,
-            clip.tokenizer,
-        )
-
-        if audio is None:
-            audio_contents = None
-        out = prepare_chatml_input(clip, input_tokens, audio_contents, sampling_rate = sampling_rate)
-        return (out,)
 
 class EmptyLatentAudio:
     def __init__(self):
@@ -623,8 +374,6 @@ NODE_CLASS_MAPPINGS = {
     "LoadAudio": LoadAudio,
     "PreviewAudio": PreviewAudio,
     "ConditioningStableAudio": ConditioningStableAudio,
-    "LoudnessNormalization": LoudnessNormalization,
-    "CreateChatMLSample": CreateChatMLSample,
     "RecordAudio": RecordAudio,
 }
 
@@ -637,7 +386,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "SaveAudio": "Save Audio (FLAC)",
     "SaveAudioMP3": "Save Audio (MP3)",
     "SaveAudioOpus": "Save Audio (Opus)",
-    "LoudnessNormalization": "Loudness Normalization",
-    "CreateChatMLSample": "Create ChatML Sample",
     "RecordAudio": "Record Audio",
 }