ice-012-audio / modeling_iceAudio.py
ice-012-audio AI TTS Model By: DarkPs
2afa665 about 2 months ago
92 kB
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-6):
        super().__init__(); self.weight=nn.Parameter(torch.ones(d)); self.eps=eps
    def forward(self,x):
        return x * torch.rsqrt(x.float().pow(2).mean(-1,keepdim=True)+self.eps).to(x.dtype) * self.weight

class HybridPosition(nn.Module):
    def __init__(self, head_dim, theta=10000.0):
        super().__init__(); self.head_dim=head_dim; self.theta=theta
        inv=1.0/(theta**(torch.arange(0,head_dim,2).float()/head_dim)); self.register_buffer('inv_freq',inv)
        self.relative_scale=nn.Parameter(torch.tensor(0.1))
    def forward(self,q,k):
        s=q.size(-2); pos=torch.arange(s,device=q.device).float(); freqs=torch.einsum('s,d->sd',pos,self.inv_freq.to(q.device)); emb=torch.cat([freqs,freqs],-1); c,sn=emb.cos(),emb.sin()
        def rot(x): return x*c[None,None,:,:]+torch.cat([-x[...,x.size(-1)//2:],x[...,:x.size(-1)//2]],-1)*sn[None,None,:,:]
        return rot(q),rot(k)

class MixedAttention(nn.Module):
    def __init__(self,d,heads,kv_heads,layer_idx,local_window=256,theta=2500000.0):
        super().__init__()
        self.d=d; self.h=heads; self.kv=kv_heads; self.hd=128; self.groups=heads//kv_heads; self.local_window=local_window
        self.q=nn.Linear(d,heads*self.hd,bias=False); self.k=nn.Linear(d,kv_heads*self.hd,bias=False); self.v=nn.Linear(d,kv_heads*self.hd,bias=False); self.o=nn.Linear(heads*self.hd,d,bias=False)
        self.qn=RMSNorm(self.hd); self.kn=RMSNorm(self.hd); self.pos=HybridPosition(self.hd, theta=float(theta))
        self.local_gate=nn.Parameter(torch.tensor(0.5)); self.global_gate=nn.Parameter(torch.tensor(0.5)); self.layer_idx=layer_idx
    def forward(self,x,mask=None):
        b,s,_=x.shape; q=self.q(x).view(b,s,self.h,self.hd).transpose(1,2); k=self.k(x).view(b,s,self.kv,self.hd).transpose(1,2); v=self.v(x).view(b,s,self.kv,self.hd).transpose(1,2)
        q=self.qn(q); k=self.kn(k); q,k=self.pos(q,k); k=k.repeat_interleave(self.groups,1); v=v.repeat_interleave(self.groups,1)
        score=torch.matmul(q.float(),k.float().transpose(-1,-2))/math.sqrt(self.hd)
        idx=torch.arange(s,device=x.device); causal=idx[None,:]>idx[:,None]; dist=(idx[None,:]-idx[:,None]).abs(); local=dist>self.local_window
        score=score.masked_fill(causal[None,None],-torch.inf)
        score=score + 0.0*torch.where(local[None,None],self.global_gate.tanh(),self.local_gate.tanh())
        if mask is not None:
            if mask.ndim==2: score=score.masked_fill(~mask[:,None,None,:],-torch.inf)
            elif mask.ndim==4: score=score.masked_fill(~mask,-torch.inf) if mask.dtype==torch.bool else score+mask
        p=F.softmax(score,dim=-1); y=torch.matmul(p,v.float()).to(x.dtype).transpose(1,2).reshape(b,s,self.h*self.hd); return self.o(y)

class SwiGLU(nn.Module):
    def __init__(self,d,hidden):
        super().__init__(); self.g=nn.Linear(d,hidden,bias=False); self.u=nn.Linear(d,hidden,bias=False); self.o=nn.Linear(hidden,d,bias=False); self.scale=nn.Parameter(torch.tensor(1.0))
    def forward(self,x): return self.o(F.silu(self.g(x))*self.u(x))*self.scale

class IndependentBlock(nn.Module):
    def __init__(self,config,layer_idx):
        super().__init__(); d=config.hidden_size; self.n1=RMSNorm(d); self.attn=MixedAttention(d,config.num_attention_heads,config.num_key_value_heads,layer_idx,theta=float(getattr(config,'rope_theta',2500000.0))); self.n2=RMSNorm(d); self.ff=SwiGLU(d,config.intermediate_size); self.attn_gate=nn.Parameter(torch.zeros(1)); self.ff_gate=nn.Parameter(torch.zeros(1))
    def forward(self,x,mask=None):
        x=x+self.attn_gate.tanh()*self.attn(self.n1(x),mask); return x+self.ff_gate.tanh()*self.ff(self.n2(x))

class IndependentBackbone(nn.Module):
    def __init__(self,config):
        super().__init__(); self.config=config; d=config.hidden_size; self.embed=nn.Embedding(config.vocab_size,d); self.blocks=nn.ModuleList([IndependentBlock(config,i) for i in range(config.num_hidden_layers)]); self.norm=RMSNorm(d)
    def get_input_embeddings(self): return self.embed
    def forward(self,input_ids=None,inputs_embeds=None,attention_mask=None,position_ids=None,return_dict=True,**kwargs):
        x=self.embed(input_ids) if inputs_embeds is None else inputs_embeds
        for block in self.blocks: x=block(x,attention_mask)
        x=self.norm(x); return type('BackboneOutput',(),{'__getitem__':lambda self,i:x,'last_hidden_state':x})()

import torch
import torch.nn as nn


class AcousticProsodyAdapter(nn.Module):
    def __init__(self, hidden_size=1024, ff_size=4096, num_layers=7, dropout=0.0):
        super().__init__()
        self.hidden_size = hidden_size
        self.audio_norm = nn.LayerNorm(hidden_size)
        self.speaker_encoder = nn.Sequential(
            nn.Linear(hidden_size, 2048), nn.SiLU(), nn.Linear(2048, hidden_size)
        )
        self.prosody_encoder = nn.Sequential(
            nn.Linear(hidden_size, 4096), nn.SiLU(), nn.Linear(4096, hidden_size)
        )
        layer = nn.TransformerEncoderLayer(
            d_model=hidden_size,
            nhead=16,
            dim_feedforward=ff_size,
            dropout=dropout,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.temporal_stack = nn.TransformerEncoder(layer, num_layers=num_layers)
        self.output_norm = nn.LayerNorm(hidden_size)
        self.output_projection = nn.Linear(hidden_size, hidden_size)
        self.residual_gate = nn.Parameter(torch.tensor([0.01]))

    @staticmethod
    def _mask_2d(audio_mask):
        mask = audio_mask.to(dtype=torch.bool)
        if mask.ndim == 3:
            mask = mask[:, 0]
        return mask

    def _route(self, hidden_states, audio_mask):
        mask = self._mask_2d(audio_mask)
        x = self.audio_norm(hidden_states)
        pooled = (x * mask.unsqueeze(-1)).sum(dim=1)
        denom = mask.sum(dim=1, keepdim=True).clamp_min(1).to(x.dtype)
        pooled = pooled / denom
        speaker = self.speaker_encoder(pooled).unsqueeze(1)
        prosody = self.prosody_encoder(pooled).unsqueeze(1)
        y = self.temporal_stack(x + speaker + prosody)
        y = self.output_projection(self.output_norm(y))
        return y

    def condition_audio_embeddings(self, audio_embeddings, audio_mask):
        return audio_embeddings + self.residual_gate.tanh() * self._route(
            audio_embeddings, audio_mask
        )

    def refine_hidden_states(self, hidden_states, audio_mask):
        return hidden_states + self.residual_gate.tanh() * self._route(
            hidden_states, audio_mask
        )

    def forward(self, hidden_states, audio_mask):
        return self.refine_hidden_states(hidden_states, audio_mask)

    def parameter_count(self):
        return sum(p.numel() for p in self.parameters())
        
import difflib
import logging
import math
import os
import re
from dataclasses import dataclass, fields
from functools import partial
from typing import Any, List, Optional, Union

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
from transformers.models.qwen3.modeling_qwen3 import Qwen3Model
import torchaudio

try:
    from torch.nn.attention.flex_attention import create_block_mask

    _flex_attention_available = True
except ImportError:
    _flex_attention_available = False
from transformers import (
    AutoFeatureExtractor,
    AutoModel,
    AutoTokenizer,
    HiggsAudioV2TokenizerModel,
    PretrainedConfig,
    PreTrainedModel,
)
from transformers.modeling_outputs import ModelOutput, BaseModelOutputWithPast
from transformers.models.auto import CONFIG_MAPPING, AutoConfig, AutoModelForCausalLM
from .configuration_iceAudio import iceAudioConfig, iceAudioLanguageConfig, iceReengineeredAudioConfig

import io
import logging

import numpy as np
import soundfile as sf
import torch
import torchaudio
from pydub import AudioSegment
from pydub.silence import detect_leading_silence, detect_nonsilent, split_on_silence

logger = logging.getLogger(__name__)

def load_waveform(audio_path: str):
    try:
        data, sr = sf.read(audio_path, dtype="float32", always_2d=True)
        return data.T, sr
    except Exception:
        import librosa

        data, sr = librosa.load(audio_path, sr=None, mono=False)
        if data.ndim == 1:
            data = data[np.newaxis, :]
        return data, sr


def load_audio(audio_path: str, sampling_rate: int) -> np.ndarray:
    data, sr = load_waveform(audio_path)

    if data.shape[0] > 1:
        data = np.mean(data, axis=0, keepdims=True)
    if sr != sampling_rate:
        data = torchaudio.functional.resample(
            torch.from_numpy(data), orig_freq=sr, new_freq=sampling_rate
        ).numpy()

    return data


def load_audio_bytes(raw: bytes, sampling_rate: int) -> np.ndarray:
    buf = io.BytesIO(raw)

    try:
        data, sr = sf.read(buf, dtype="float32", always_2d=True)
        data = data.T
    except Exception:
        import librosa

        buf.seek(0)
        data, sr = librosa.load(buf, sr=None, mono=False)
        if data.ndim == 1:
            data = data[np.newaxis, :]

    if data.shape[0] > 1:
        data = np.mean(data, axis=0, keepdims=True)
    if sr != sampling_rate:
        data = torchaudio.functional.resample(
            torch.from_numpy(data), orig_freq=sr, new_freq=sampling_rate
        ).numpy()

    return data

def numpy_to_audiosegment(audio: np.ndarray, sample_rate: int) -> AudioSegment:
    audio_int = (audio * 32768.0).clip(-32768, 32767).astype(np.int16)
    if audio_int.shape[0] > 1:
        audio_int = audio_int.T.flatten()
    return AudioSegment(
        data=audio_int.tobytes(),
        sample_width=2,
        frame_rate=sample_rate,
        channels=audio.shape[0],
    )


def audiosegment_to_numpy(aseg: AudioSegment) -> np.ndarray:
    data = np.array(aseg.get_array_of_samples()).astype(np.float32) / 32768.0
    if aseg.channels == 1:
        return data[np.newaxis, :]
    return data.reshape(-1, aseg.channels).T


def remove_silence(
    audio: np.ndarray,
    sampling_rate: int,
    mid_sil: int = 300,
    lead_sil: int = 100,
    trail_sil: int = 300,
) -> np.ndarray:
    wave = numpy_to_audiosegment(audio, sampling_rate)

    if mid_sil > 0:
        non_silent_segs = split_on_silence(
            wave,
            min_silence_len=mid_sil,
            silence_thresh=-50,
            keep_silence=mid_sil,
            seek_step=10,
        )
        wave = AudioSegment.silent(duration=0)
        for seg in non_silent_segs:
            wave += seg

    wave = remove_silence_edges(wave, lead_sil, trail_sil, -50)

    return audiosegment_to_numpy(wave)


def remove_silence_edges(
    audio: AudioSegment,
    lead_sil: int = 100,
    trail_sil: int = 300,
    silence_threshold: float = -50,
) -> AudioSegment:
    start_idx = detect_leading_silence(audio, silence_threshold=silence_threshold)
    start_idx = max(0, start_idx - lead_sil)
    audio = audio[start_idx:]

    audio = audio.reverse()
    start_idx = detect_leading_silence(audio, silence_threshold=silence_threshold)
    start_idx = max(0, start_idx - trail_sil)
    audio = audio[start_idx:]
    audio = audio.reverse()

    return audio


def fade_and_pad_audio(
    audio: np.ndarray,
    pad_duration: float = 0.1,
    fade_duration: float = 0.1,
    sample_rate: int = 24000,
) -> np.ndarray:
    if audio.shape[-1] == 0:
        return audio

    fade_samples = int(fade_duration * sample_rate)
    pad_samples = int(pad_duration * sample_rate)

    processed = audio.copy()

    if fade_samples > 0:
        k = min(fade_samples, processed.shape[-1] // 2)
        if k > 0:
            fade_in = np.linspace(0, 1, k, dtype=np.float32)[np.newaxis, :]
            processed[..., :k] *= fade_in

            fade_out = np.linspace(1, 0, k, dtype=np.float32)[np.newaxis, :]
            processed[..., -k:] *= fade_out

    if pad_samples > 0:
        silence = np.zeros(
            (processed.shape[0], pad_samples),
            dtype=processed.dtype,
        )
        processed = np.concatenate([silence, processed, silence], axis=-1)

    return processed


def trim_long_audio(
    audio: np.ndarray,
    sampling_rate: int,
    max_duration: float = 15.0,
    min_duration: float = 3.0,
    trim_threshold: float = 20.0,
) -> np.ndarray:
    duration = audio.shape[-1] / sampling_rate
    if duration <= trim_threshold:
        return audio

    seg = numpy_to_audiosegment(audio, sampling_rate)
    nonsilent = detect_nonsilent(
        seg, min_silence_len=100, silence_thresh=-40, seek_step=10
    )
    if not nonsilent:
        return audio

    max_ms = int(max_duration * 1000)
    min_ms = int(min_duration * 1000)

    best_split = 0
    for start, end in nonsilent:
        if start > best_split and start <= max_ms:
            best_split = start
        if end > max_ms:
            break

    if best_split < min_ms:
        best_split = min(max_ms, len(seg))

    trimmed = seg[:best_split]
    return audiosegment_to_numpy(trimmed)


def cross_fade_chunks(
    chunks: list[np.ndarray],
    sample_rate: int,
    silence_duration: float = 0.3,
) -> np.ndarray:
    if len(chunks) == 1:
        return chunks[0]

    total_n = int(silence_duration * sample_rate)
    fade_n = total_n // 3
    silence_n = fade_n
    merged = chunks[0].copy()

    for chunk in chunks[1:]:
        parts = [merged]

        fout_n = min(fade_n, merged.shape[-1])
        if fout_n > 0:
            w_out = np.linspace(1, 0, fout_n, dtype=np.float32)[np.newaxis, :]
            parts[-1][..., -fout_n:] *= w_out

        parts.append(np.zeros((chunks[0].shape[0], silence_n), dtype=np.float32))

        fade_in = chunk.copy()
        fin_n = min(fade_n, fade_in.shape[-1])
        if fin_n > 0:
            w_in = np.linspace(0, 1, fin_n, dtype=np.float32)[np.newaxis, :]
            fade_in[..., :fin_n] *= w_in

        parts.append(fade_in)
        merged = np.concatenate(parts, axis=-1)

    return merged

import bisect
import unicodedata
from functools import lru_cache
from typing import Optional


class RuleDurationEstimator:
    def __init__(self):
        self.weights = {
            "cjk": 3.0,
            "hangul": 2.5,
            "kana": 2.2,
            "ethiopic": 3.0,
            "yi": 3.0,
            "indic": 1.8,
            "thai_lao": 1.5,
            "khmer_myanmar": 1.8,
            "arabic": 1.5,
            "hebrew": 1.5,
            "latin": 1.0,
            "cyrillic": 1.0,
            "greek": 1.0,
            "armenian": 1.0,
            "georgian": 1.0,
            "punctuation": 0.5,
            "space": 0.2,
            "digit": 3.5,
            "mark": 0.0,
            "default": 1.0,
        }
        self.ranges = [
            (0x02AF, "latin"),
            (0x03FF, "greek"),
            (0x052F, "cyrillic"),
            (0x058F, "armenian"),
            (0x05FF, "hebrew"),
            (0x077F, "arabic"),
            (0x089F, "arabic"),
            (0x08FF, "arabic"),
            (0x097F, "indic"),
            (0x09FF, "indic"),
            (0x0A7F, "indic"),
            (0x0AFF, "indic"),
            (0x0B7F, "indic"),
            (0x0BFF, "indic"),
            (0x0C7F, "indic"),
            (0x0CFF, "indic"),
            (0x0D7F, "indic"),
            (0x0DFF, "indic"),
            (0x0EFF, "thai_lao"),
            (0x0FFF, "indic"),
            (0x109F, "khmer_myanmar"),
            (0x10FF, "georgian"),
            (0x11FF, "hangul"),
            (0x137F, "ethiopic"),
            (0x139F, "ethiopic"),
            (0x13FF, "default"),
            (0x167F, "default"),
            (0x169F, "default"),
            (0x16FF, "default"),
            (0x171F, "default"),
            (0x173F, "default"),
            (0x175F, "default"),
            (0x177F, "default"),
            (0x17FF, "khmer_myanmar"),
            (0x18AF, "default"),
            (0x18FF, "default"),
            (0x194F, "indic"),
            (0x19DF, "indic"),
            (0x19FF, "khmer_myanmar"),
            (0x1A1F, "indic"),
            (0x1AAF, "indic"),
            (0x1B7F, "indic"),
            (0x1BBF, "indic"),
            (0x1BFF, "indic"),
            (0x1C4F, "indic"),
            (0x1C7F, "indic"),
            (0x1C8F, "cyrillic"),
            (0x1CBF, "georgian"),
            (0x1CCF, "indic"),
            (0x1CFF, "indic"),
            (0x1D7F, "latin"),
            (0x1DBF, "latin"),
            (0x1DFF, "default"),
            (0x1EFF, "latin"),
            (0x309F, "kana"),
            (0x30FF, "kana"),
            (0x312F, "cjk"),
            (0x318F, "hangul"),
            (0x9FFF, "cjk"),
            (0xA4CF, "yi"),
            (0xA4FF, "default"),
            (0xA63F, "default"),
            (0xA69F, "cyrillic"),
            (0xA6FF, "default"),
            (0xA7FF, "latin"),
            (0xA82F, "indic"),
            (0xA87F, "default"),
            (0xA8DF, "indic"),
            (0xA8FF, "indic"),
            (0xA92F, "indic"),
            (0xA95F, "indic"),
            (0xA97F, "hangul"),
            (0xA9DF, "indic"),
            (0xA9FF, "khmer_myanmar"),
            (0xAA5F, "indic"),
            (0xAA7F, "khmer_myanmar"),
            (0xAADF, "indic"),
            (0xAAFF, "indic"),
            (0xAB2F, "ethiopic"),
            (0xAB6F, "latin"),
            (0xABBF, "default"),
            (0xABFF, "indic"),
            (0xD7AF, "hangul"),
            (0xFAFF, "cjk"),
            (0xFDFF, "arabic"),
            (0xFE6F, "default"),
            (0xFEFF, "arabic"),
            (0xFFEF, "latin"),
        ]
        self.breakpoints = [r[0] for r in self.ranges]

    @lru_cache(maxsize=4096)
    def _get_char_weight(self, char):
        code = ord(char)
        if (65 <= code <= 90) or (97 <= code <= 122):
            return self.weights["latin"]
        if code == 32:
            return self.weights["space"]
   
        if code == 0x0640:
            return self.weights["mark"]

        category = unicodedata.category(char)

        if category.startswith("M"):
            return self.weights["mark"]

        if category.startswith("P") or category.startswith("S"):
            return self.weights["punctuation"]

        if category.startswith("Z"):
            return self.weights["space"]

        if category.startswith("N"):
            return self.weights["digit"]

        idx = bisect.bisect_left(self.breakpoints, code)
        if idx < len(self.ranges):
            script_type = self.ranges[idx][1]
            return self.weights.get(script_type, self.weights["default"])

        if code > 0x20000:
            return self.weights["cjk"]

        return self.weights["default"]

    def calculate_total_weight(self, text):

        return sum(self._get_char_weight(c) for c in text)

    def estimate_duration(
        self,
        target_text: str,
        ref_text: str,
        ref_duration: float,
        low_threshold: Optional[float] = 50,
        boost_strength: float = 3,
    ) -> float:
        if ref_duration <= 0 or not ref_text:
            return 0.0

        ref_weight = self.calculate_total_weight(ref_text)
        if ref_weight == 0:
            return 0.0

        speed_factor = ref_weight / ref_duration
        target_weight = self.calculate_total_weight(target_text)

        estimated_duration = target_weight / speed_factor
        if low_threshold is not None and estimated_duration < low_threshold:
            alpha = 1.0 / boost_strength
            return low_threshold * (estimated_duration / low_threshold) ** alpha
        else:
            return estimated_duration

LANG_NAME_TO_ID = {'egyptian lahgtna': 'eg',
 'saudi lahgtna': 'sa',
 'moroccan lahgtna': 'ma',
 'bahraini lahgtna': 'bh',
 'sudanese lahgtna': 'sd',
 'iraqi lahgtna': 'iq',
 'lebanese lahgtna': 'lb',
 'syrian lahgtna': 'sy',
 'libyan lahgtna': 'ly',
 'palestinian lahgtna': 'ps',
 'tunisian lahgtna': 'tn',
 'algerian lahgtna': 'dz',
 'yemeni lahgtna': 'ye',
 'abkhazian': 'ab',
 'afrikaans': 'af',
 'albanian': 'sq',
 'algerian arabic': 'arq',
 'algerian saharan arabic': 'aao',
 'amharic': 'am',
 'armenian': 'hy',
 'assamese': 'as',
 'asturian': 'ast',
 'azerbaijani': 'az',
 'baharna arabic': 'abv',
 'banjar': 'bjn',
 'bashkir': 'ba',
 'basque': 'eu',
 'batak mandailing': 'btm',
 'belarusian': 'be',
 'bengali': 'bn',
 'betawi': 'bew',
 'bhojpuri': 'bho',
 'bosnian': 'bs',
 'breton': 'br',
 'buginese': 'bug',
 'bukharic': 'bhh',
 'bulgarian': 'bg',
 'bulu (cameroon)': 'bum',
 'bundeli': 'bns',
 'burmese': 'my',
 'burushaski': 'bsk',
 'cajatambo north lima quechua': 'qvl',
 'cameroon pidgin': 'wes',
 'campidanese sardinian': 'sro',
 'cantonese': 'yue',
 'catalan': 'ca',
 'cebuano': 'ceb',
 'central kurdish': 'ckb',
 'central nahuatl': 'nhn',
 'central pashto': 'pst',
 'central puebla nahuatl': 'ncx',
 'central yupik': 'esu',
 'central-eastern niger fulfulde': 'fuq',
 'chadian arabic': 'shu',
 'chichewa': 'ny',
 'chichicapan zapotec': 'zpv',
 'chiga': 'cgg',
 'chimalapa zoque': 'zoh',
 'chimborazo highland quichua': 'qug',
 'chinese': 'zh',
 'chiquián ancash quechua': 'qxa',
 'chitwania tharu': 'the',
 'chokwe': 'cjk',
 'chuvash': 'cv',
 'cibak': 'ckl',
 'coastal konjo': 'kjc',
 'copainalá zoque': 'zoc',
 'cornish': 'kw',
 'corongo ancash quechua': 'qwa',
 'croatian': 'hr',
 'cross river mbembe': 'mfn',
 'cuyamecalco mixtec': 'xtu',
 'czech': 'cs',
 'dadiya': 'dbd',
 'dagbani': 'dag',
 'dameli': 'dml',
 'danish': 'da',
 'dargwa': 'dar',
 'dazaga': 'dzg',
 'deccan': 'dcc',
 'degema': 'deg',
 'dera (nigeria)': 'kna',
 'dghwede': 'dgh',
 'dhatki': 'mki',
 'dhivehi': 'dv',
 'dhofari arabic': 'adf',
 'dijim-bwilim': 'cfa',
 'dogri': 'dgo',
 'domaaki': 'dmk',
 'dotyali': 'dty',
 'duala': 'dua',
 'dutch': 'nl',
 'dũya': 'ldb',
 'dyula': 'dyu',
 'eastern balochi': 'bgp',
 'eastern bolivian guaraní': 'gui',
 'eastern egyptian bedawi arabic': 'avl',
 'eastern krahn': 'kqo',
 'eastern mari': 'mhr',
 'eastern yiddish': 'ydd',
 'ebrié': 'ebr',
 'eggon': 'ego',
 'egyptian arabic': 'arz',
 'ejagham': 'etu',
 'eleme': 'elm',
 'eloyi': 'afo',
 'embu': 'ebu',
 'english': 'en',
 'erzya': 'myv',
 'esan': 'ish',
 'esperanto': 'eo',
 'estonian': 'et',
 'eton (cameroon)': 'eto',
 'ewondo': 'ewo',
 'extremaduran': 'ext',
 'fang (equatorial guinea)': 'fan',
 'fanti': 'fat',
 'farefare': 'gur',
 "fe'fe'": 'fmp',
 'filipino': 'fil',
 'filomena mata-coahuitlán totonac': 'tlp',
 'finnish': 'fi',
 'fipa': 'fip',
 'french': 'fr',
 'fulah': 'ff',
 'galician': 'gl',
 'gambian wolof': 'wof',
 'ganda': 'lg',
 'garhwali': 'gbm',
 'gawar-bati': 'gwt',
 'gawri': 'gwc',
 'gbagyi': 'gbr',
 'gbari': 'gby',
 'geji': 'gyz',
 'gen': 'gej',
 'georgian': 'ka',
 'german': 'de',
 'geser-gorom': 'ges',
 'gheg albanian': 'aln',
 "ghomálá'": 'bbj',
 'gidar': 'gid',
 'glavda': 'glw',
 'goan konkani': 'gom',
 'goaria': 'gig',
 'goemai': 'ank',
 'gola': 'gol',
 'greek': 'el',
 'guarani': 'gn',
 'guduf-gava': 'gdf',
 'guerrero amuzgo': 'amu',
 'gujarati': 'gu',
 'gujari': 'gju',
 'gulf arabic': 'afb',
 'gurgula': 'ggg',
 'gusii': 'guz',
 'gusilay': 'gsl',
 'gweno': 'gwe',
 'güilá zapotec': 'ztu',
 'hadothi': 'hoj',
 'hahon': 'hah',
 'haitian': 'ht',
 'hakha chin': 'cnh',
 'hakö': 'hao',
 'halia': 'hla',
 'hausa': 'ha',
 'hawaiian': 'haw',
 'hazaragi': 'haz',
 'hebrew': 'he',
 'hemba': 'hem',
 'herero': 'hz',
 'highland konjo': 'kjk',
 'hijazi arabic': 'acw',
 'hindi': 'hi',
 'huarijio': 'var',
 'huautla mazatec': 'mau',
 'huaxcaleca nahuatl': 'nhq',
 'huba': 'hbb',
 'huitepec mixtec': 'mxs',
 'hula': 'hul',
 'hungarian': 'hu',
 'hunjara-kaina ke': 'hkk',
 'hwana': 'hwo',
 'ibibio': 'ibb',
 'icelandic': 'is',
 'idakho-isukha-tiriki': 'ida',
 'idoma': 'idu',
 'igbo': 'ig',
 'igo': 'ahl',
 'ikposo': 'kpo',
 'ikwere': 'ikw',
 'imbabura highland quichua': 'qvi',
 'indonesian': 'id',
 'indus kohistani': 'mvy',
 'interlingua (international auxiliary language association)': 'ia',
 'inupiaq': 'ik',
 'irish': 'ga',
 'iron ossetic': 'os',
 'isekiri': 'its',
 'isoko': 'iso',
 'italian': 'it',
 'ito': 'itw',
 'itzá': 'itz',
 'ixtayutla mixtec': 'vmj',
 'izon': 'ijc',
 'jambi malay': 'jax',
 'japanese': 'ja',
 'jaqaru': 'jqr',
 'jauja wanca quechua': 'qxw',
 'jaunsari': 'jns',
 'javanese': 'jv',
 'jiba': 'juo',
 'jju': 'kaj',
 'judeo-moroccan arabic': 'aju',
 'juxtlahuaca mixtec': 'vmc',
 'kabardian': 'kbd',
 'kabras': 'lkb',
 'kabuverdianu': 'kea',
 'kabyle': 'kab',
 'kachi koli': 'gjk',
 'kairak': 'ckr',
 'kalabari': 'ijn',
 'kalasha': 'kls',
 'kalenjin': 'kln',
 'kalkoti': 'xka',
 'kamba': 'kam',
 'kamo': 'kcq',
 'kanauji': 'bjj',
 'kanembu': 'kbl',
 'kannada': 'kn',
 'karekare': 'kai',
 'kashmiri': 'ks',
 'kathoriya tharu': 'tkt',
 'kati': 'bsh',
 'kazakh': 'kk',
 'keiyo': 'eyo',
 'khams tibetan': 'khg',
 'khana': 'ogo',
 'khetrani': 'xhe',
 'khmer': 'km',
 'khowar': 'khw',
 'kinga': 'zga',
 'kinnauri': 'kfk',
 'kinyarwanda': 'rw',
 'kirghiz': 'ky',
 'kirya-konzəl': 'fkk',
 'kochila tharu': 'thq',
 'kohistani shina': 'plk',
 'kohumono': 'bcs',
 'kok borok': 'trp',
 'kol (papua new guinea)': 'kol',
 'kom (cameroon)': 'bkm',
 'koma': 'kmy',
 'konkani': 'knn',
 'konzo': 'koo',
 'korean': 'ko',
 'korwa': 'kfp',
 'kota (india)': 'kfe',
 'koti': 'eko',
 'kuanua': 'ksd',
 'kuanyama': 'kj',
 'kui (india)': 'uki',
 'kulung (nigeria)': 'bbu',
 'kuot': 'kto',
 'kushi': 'kuh',
 'kwambi': 'kwm',
 'kwasio': 'nmg',
 'lala-roba': 'lla',
 'lamang': 'hia',
 'lao': 'lo',
 'larike-wakasihu': 'alo',
 'lasi': 'lss',
 'latgalian': 'ltg',
 'latvian': 'lv',
 'levantine arabic': 'apc',
 'liana-seti': 'ste',
 'liberia kpelle': 'xpe',
 'liberian english': 'lir',
 'libyan arabic': 'ayl',
 'ligurian': 'lij',
 'lijili': 'mgi',
 'lingala': 'ln',
 'lithuanian': 'lt',
 'loarki': 'lrk',
 'logooli': 'rag',
 'logudorese sardinian': 'src',
 'loja highland quichua': 'qvj',
 'loloda': 'loa',
 'longuda': 'lnu',
 'loxicha zapotec': 'ztp',
 'luba-lulua': 'lua',
 'luo': 'luo',
 'lushai': 'lus',
 'luxembourgish': 'lb',
 'maasina fulfulde': 'ffm',
 'maba (chad)': 'mde',
 'macedo-romanian': 'rup',
 'macedonian': 'mk',
 'mada (cameroon)': 'mxu',
 'mafa': 'maf',
 'maithili': 'mai',
 'malay': 'ms',
 'malayalam': 'ml',
 'mali': 'gcc',
 "malinaltepec me'phaa": 'tcf',
 'maltese': 'mt',
 'mandara': 'tbf',
 'mandjak': 'mfv',
 'manggarai': 'mqy',
 'manipuri': 'mni',
 'mansoanka': 'msw',
 'manx': 'gv',
 'maori': 'mi',
 'marathi': 'mr',
 'marghi central': 'mrt',
 'marghi south': 'mfm',
 'maria (india)': 'mrr',
 'marwari (pakistan)': 'mve',
 'masana': 'mcn',
 'masikoro malagasy': 'msh',
 'matsés': 'mcf',
 'mazaltepec zapotec': 'zpy',
 'mazatlán mazatec': 'vmz',
 'mazatlán mixe': 'mzl',
 'mbe': 'mfo',
 'mbo (cameroon)': 'mbo',
 'mbum': 'mdd',
 'medumba': 'byv',
 'mekeo': 'mek',
 'meru': 'mer',
 'mesopotamian arabic': 'acm',
 'mewari': 'mtr',
 'min nan chinese': 'nan',
 'mingrelian': 'xmf',
 'mitlatongo mixtec': 'vmm',
 'miya': 'mkf',
 'mokpwe': 'bri',
 'moksha': 'mdf',
 'mom jango': 'ver',
 'mongolian': 'mn',
 'moroccan arabic': 'ary',
 'motu': 'meu',
 'mpiemo': 'mcx',
 'mpumpong': 'mgg',
 'mundang': 'mua',
 'mungaka': 'mhk',
 'musey': 'mse',
 'musgu': 'mug',
 'musi': 'mui',
 'naba': 'mne',
 'najdi arabic': 'ars',
 'nalik': 'nal',
 'nawdm': 'nmz',
 'ndonga': 'ng',
 'neapolitan': 'nap',
 'nepali': 'npi',
 'ngamo': 'nbh',
 'ngas': 'anc',
 'ngiemboon': 'nnh',
 'ngizim': 'ngi',
 'ngomba': 'jgo',
 'ngombale': 'nla',
 'nigerian fulfulde': 'fuv',
 'nigerian pidgin': 'pcm',
 'nimadi': 'noe',
 'nobiin': 'fia',
 'north mesopotamian arabic': 'ayp',
 'north moluccan malay': 'max',
 'northern betsimisaraka malagasy': 'bmm',
 'northern hindko': 'hno',
 'northern kurdish': 'kmr',
 'northern pame': 'pmq',
 'northern pashto': 'pbu',
 'northern uzbek': 'uzn',
 'northwest gbaya': 'gya',
 'norwegian': 'no',
 'norwegian bokmål': 'nb',
 'norwegian nynorsk': 'nn',
 'notsi': 'ncf',
 'nyankpa': 'yes',
 'nyungwe': 'nyu',
 'nzanyi': 'nja',
 'nüpode huitoto': 'hux',
 'occitan': 'oc',
 'od': 'odk',
 'odia': 'ory',
 'odual': 'odu',
 'omani arabic': 'acx',
 'orizaba nahuatl': 'nlv',
 'orma': 'orc',
 'ormuri': 'oru',
 'oromo': 'om',
 'pahari-potwari': 'phr',
 'paiwan': 'pwn',
 'panjabi': 'pa',
 'papuan malay': 'pmy',
 'parkari koli': 'kvx',
 'pedi': 'nso',
 'pero': 'pip',
 'persian': 'fa',
 'petats': 'pex',
 'phalura': 'phl',
 'piemontese': 'pms',
 'piya-kwonci': 'piy',
 'plateau malagasy': 'plt',
 'polish': 'pl',
 'poqomam': 'poc',
 'portuguese': 'pt',
 'pulaar': 'fuc',
 'pular': 'fuf',
 'puno quechua': 'qxp',
 'pushto': 'ps',
 'pökoot': 'pko',
 'qaqet': 'byx',
 'quiotepec chinantec': 'chq',
 'rana tharu': 'thr',
 'rangi': 'lag',
 'rapoisi': 'kyx',
 'ratahan': 'rth',
 'rayón zoque': 'zor',
 'romanian': 'ro',
 'romansh': 'rm',
 'rombo': 'rof',
 'rotokas': 'roo',
 'rukai': 'dru',
 'russian': 'ru',
 'sacapulteco': 'quv',
 'saidi arabic': 'aec',
 'sakalava malagasy': 'skg',
 'sakizaya': 'szy',
 'saleman': 'sau',
 'samba daka': 'ccg',
 'samba leko': 'ndi',
 'san felipe otlaltepec popoloca': 'pow',
 'san francisco del mar huave': 'hue',
 'san juan atzingo popoloca': 'poe',
 'san martín itunyoso triqui': 'trq',
 'san miguel el grande mixtec': 'mig',
 'sansi': 'ssi',
 'sanskrit': 'sa',
 'santa ana de tusi pasco quechua': 'qxt',
 'santa catarina albarradas zapotec': 'ztn',
 'santali': 'sat',
 'santiago del estero quichua': 'qus',
 'saposa': 'sps',
 'saraiki': 'skr',
 'sardinian': 'sc',
 'saya': 'say',
 'sediq': 'trv',
 'serbian': 'sr',
 'seri': 'sei',
 'shina': 'scl',
 'shona': 'sn',
 'siar-lak': 'sjr',
 'sibe': 'nco',
 'sicilian': 'scn',
 'sihuas ancash quechua': 'qws',
 'sikkimese': 'sip',
 'sinaugoro': 'snc',
 'sindhi': 'sd',
 'sindhi bhil': 'sbn',
 'sinhala': 'si',
 'sinicahua mixtec': 'xti',
 'sipacapense': 'qum',
 'siwai': 'siw',
 'slovak': 'sk',
 'slovenian': 'sl',
 'solos': 'sol',
 'somali': 'so',
 'soninke': 'snk',
 'south giziga': 'giz',
 'south ucayali ashéninka': 'cpy',
 'southeastern nochixtlán mixtec': 'mxy',
 'southern betsimisaraka malagasy': 'bzc',
 'southern pashto': 'pbt',
 'southern pastaza quechua': 'qup',
 'soyaltepec mazatec': 'vmp',
 'spanish': 'es',
 'standard arabic': 'arb',
 'standard moroccan tamazight': 'zgh',
 'sudanese arabic': 'apd',
 'sulka': 'sua',
 'svan': 'sva',
 'swahili': 'sw',
 'swedish': 'sv',
 "tae'": 'rob',
 'tahaggart tamahaq': 'thv',
 'taita': 'dav',
 'tajik': 'tg',
 'tamil': 'ta',
 'tandroy-mahafaly malagasy': 'tdx',
 'tangale': 'tan',
 'tanosy malagasy': 'txy',
 'tarok': 'yer',
 'tatar': 'tt',
 'tedaga': 'tuq',
 'telugu': 'te',
 'tem': 'kdh',
 'teop': 'tio',
 'tepeuxila cuicatec': 'cux',
 'tepinapa chinantec': 'cte',
 'tera': 'ttr',
 'terei': 'buo',
 'termanu': 'twu',
 'tesaka malagasy': 'tkg',
 'tetelcingo nahuatl': 'nhg',
 'teutila cuicatec': 'cut',
 'thai': 'th',
 'tibetan': 'bo',
 'tidaá mixtec': 'mtx',
 'tidore': 'tvo',
 'tigak': 'tgc',
 'tigre': 'tig',
 'tigrinya': 'ti',
 'tilquiapan zapotec': 'zts',
 'tinputz': 'tpz',
 "tlacoapa me'phaa": 'tpl',
 'tlacoatzintepec chinantec': 'ctl',
 'tlingit': 'tli',
 'toki pona': 'tok',
 'tomoip': 'tqp',
 'tondano': 'tdn',
 'tonsea': 'txs',
 'tooro': 'ttj',
 'torau': 'ttu',
 'torwali': 'trw',
 'tsimihety malagasy': 'xmw',
 'tsotso': 'lto',
 'tswana': 'tn',
 'tugen': 'tuy',
 'tuki': 'bag',
 'tula': 'tul',
 'tulu': 'tcy',
 'tunen': 'tvu',
 'tungag': 'lcm',
 'tunisian arabic': 'aeb',
 'tupuri': 'tui',
 'turkana': 'tuv',
 'turkish': 'tr',
 'turkmen': 'tk',
 'tututepec mixtec': 'mtu',
 'twi': 'tw',
 'ubaghara': 'byc',
 'uighur': 'ug',
 'ukrainian': 'uk',
 'umbundu': 'umb',
 'upper sorbian': 'hsb',
 'urdu': 'ur',
 'ushojo': 'ush',
 'uzbek': 'uz',
 'vai': 'vai',
 'vietnamese': 'vi',
 'votic': 'vot',
 'võro': 'vro',
 'waci gbe': 'wci',
 'wadiyara koli': 'kxp',
 'waja': 'wja',
 'wakhi': 'wbl',
 'wanga': 'lwg',
 'wapan': 'juk',
 'warji': 'wji',
 'welsh': 'cy',
 'wemale': 'weo',
 'western frisian': 'fy',
 'western highland purepecha': 'pua',
 'western juxtlahuaca mixtec': 'jmx',
 'western maninkakan': 'mlq',
 'western mari': 'mrj',
 'western niger fulfulde': 'fuh',
 'western panjabi': 'pnb',
 'wolof': 'wo',
 'wuzlam': 'udl',
 'xanaguía zapotec': 'ztg',
 'xhosa': 'xh',
 'yace': 'ekr',
 'yakut': 'sah',
 'yalahatan': 'jal',
 'yanahuanca pasco quechua': 'qur',
 'yangben': 'yav',
 'yaqui': 'yaq',
 'yauyos quechua': 'qux',
 'yekhee': 'ets',
 'yiddish': 'yi',
 'yidgha': 'ydg',
 'yoruba': 'yo',
 'yutanduchi mixtec': 'mab',
 'zacatlán-ahuacatlán-tepetzintla nahuatl': 'nhi',
 'zarma': 'dje',
 'zaza': 'zza',
 'zulu': 'zu',
 'ömie': 'aom'}

LANG_NAMES = set(LANG_NAME_TO_ID.keys())
LANG_IDS = set(LANG_NAME_TO_ID.values())

_TITLE_EXCEPTIONS = {
    "fe'fe'": "Fe'fe'",
    "dũya": "Dũya",
    "santiago del estero quichua": "Santiago del Estero Quichua",
    "santa ana de tusi pasco quechua": "Santa Ana de Tusi Pasco Quechua",
    "malinaltepec me'phaa": "Malinaltepec Me'phaa",
    "tlacoapa me'phaa": "Tlacoapa Me'phaa",
}


def lang_display_name(name: str) -> str:
    return _TITLE_EXCEPTIONS.get(name, name.title())

from typing import List, Optional


SPLIT_PUNCTUATION = set(".,;:!?。，；：！？")
CLOSING_MARKS = set("\"'""'）]》》>」】")

END_PUNCTUATION = {
    ";",
    ":",
    ",",
    ".",
    "!",
    "?",
    "…",
    ")",
    "]",
    "}",
    '"',
    "'",
    """,
    "'",
    "；",
    "：",
    "，",
    "。",
    "！",
    "？",
    "、",
    "……",
    "）",
    "】",
    """,
    "'",
}


ABBREVIATIONS = {
    "Mr.",
    "Mrs.",
    "Ms.",
    "Dr.",
    "Prof.",
    "Sr.",
    "Jr.",
    "Rev.",
    "Fr.",
    "Hon.",
    "Pres.",
    "Gov.",
    "Capt.",
    "Gen.",
    "Sen.",
    "Rep.",
    "Col.",
    "Maj.",
    "Lt.",
    "Cmdr.",
    "Sgt.",
    "Cpl.",
    "Co.",
    "Corp.",
    "Inc.",
    "Ltd.",
    "Est.",
    "Dept.",
    "St.",
    "Ave.",
    "Blvd.",
    "Rd.",
    "Mt.",
    "Ft.",
    "No.",
    "Jan.",
    "Feb.",
    "Mar.",
    "Apr.",
    "Aug.",
    "Sep.",
    "Sept.",
    "Oct.",
    "Nov.",
    "Dec.",
    "i.e.",
    "e.g.",
    "vs.",
    "Vs.",
    "Etc.",
    "approx.",
    "fig.",
    "def.",
}


def chunk_text_punctuation(
    text: str,
    chunk_len: int,
    min_chunk_len: Optional[int] = None,
) -> List[str]:

    sentences = []
    current_sentence = []

    tokens_list = list(text)

    for token in tokens_list:

        if (
            len(current_sentence) == 0
            and len(sentences) != 0
            and (token in SPLIT_PUNCTUATION or token in CLOSING_MARKS)
        ):
            sentences[-1].append(token)

        else:
            current_sentence.append(token)

            if token in SPLIT_PUNCTUATION:
                is_abbreviation = False

                if token == ".":
                    temp_str = "".join(current_sentence).strip()
                    if temp_str:
                        last_word = temp_str.split()[-1]
                        if last_word in ABBREVIATIONS:
                            is_abbreviation = True

                if not is_abbreviation:
                    sentences.append(current_sentence)
                    current_sentence = []

    if len(current_sentence) != 0:
        sentences.append(current_sentence)

    merged_chunks = []
    current_chunk = []
    for sentence in sentences:
        if len(current_chunk) + len(sentence) <= chunk_len:
            current_chunk.extend(sentence)
        else:
            if len(current_chunk) > 0:
                merged_chunks.append(current_chunk)
            current_chunk = sentence

    if len(current_chunk) > 0:
        merged_chunks.append(current_chunk)

    if min_chunk_len is not None:
        first_chunk_short_flag = (
            len(merged_chunks) > 0 and len(merged_chunks[0]) < min_chunk_len
        )
        final_chunks = []
        for i, chunk in enumerate(merged_chunks):
            if i == 1 and first_chunk_short_flag:
                final_chunks[-1].extend(chunk)
            else:
                if len(chunk) >= min_chunk_len:
                    final_chunks.append(chunk)
                else:
                    if len(final_chunks) == 0:
                        final_chunks.append(chunk)
                    else:
                        final_chunks[-1].extend(chunk)
    else:
        final_chunks = merged_chunks

    chunk_strings = [
        "".join(chunk).strip() for chunk in final_chunks if "".join(chunk).strip()
    ]
    return chunk_strings


def add_punctuation(text: str):

    text = text.strip()

    if not text:
        return text

    if text[-1] not in END_PUNCTUATION:
        is_chinese = any("\u4e00" <= char <= "\u9fff" for char in text)

        text += "。" if is_chinese else "."

    return text

import re

_ZH_RE = re.compile(r'[\u4e00-\u9fff]')

_INSTRUCT_CATEGORIES = [
    {"male": "男", "female": "女"},
    {"child": "儿童", "teenager": "少年", "young adult": "青年",
     "middle-aged": "中年", "elderly": "老年"},
    {"very low pitch": "极低音调", "low pitch": "低音调",
     "moderate pitch": "中音调", "high pitch": "高音调",
     "very high pitch": "极高音调"},
    {"whisper": "耳语"},
    {"american accent", "british accent", "australian accent",
     "chinese accent", "canadian accent", "indian accent",
     "korean accent", "portuguese accent", "russian accent", "japanese accent"},
    {"河南话", "陕西话", "四川话", "贵州话", "云南话", "桂林话",
     "济南话", "石家庄话", "甘肃话", "宁夏话", "青岛话", "东北话"},
]

_INSTRUCT_EN_TO_ZH = {}
_INSTRUCT_ZH_TO_EN = {}
_INSTRUCT_MUTUALLY_EXCLUSIVE = []
for _cat in _INSTRUCT_CATEGORIES:
    if isinstance(_cat, dict):
        _INSTRUCT_EN_TO_ZH.update(_cat)
        _INSTRUCT_ZH_TO_EN.update({v: k for k, v in _cat.items()})
        _INSTRUCT_MUTUALLY_EXCLUSIVE.append(set(_cat) | set(_cat.values()))
    else:
        _INSTRUCT_MUTUALLY_EXCLUSIVE.append(set(_cat))

_INSTRUCT_ALL_VALID = (
    set(_INSTRUCT_EN_TO_ZH) | set(_INSTRUCT_ZH_TO_EN)
    | _INSTRUCT_MUTUALLY_EXCLUSIVE[-2]
    | _INSTRUCT_MUTUALLY_EXCLUSIVE[-1]
)

_INSTRUCT_VALID_EN = frozenset(i for i in _INSTRUCT_ALL_VALID if not _ZH_RE.search(i))
_INSTRUCT_VALID_ZH = frozenset(i for i in _INSTRUCT_ALL_VALID if _ZH_RE.search(i))


logger = logging.getLogger(__name__)

@dataclass
class VoiceClonePrompt:
    ref_audio_tokens: torch.Tensor
    ref_text: str
    ref_rms: float


@dataclass
class iceAudioGenerationConfig:
    num_step: int = 32
    guidance_scale: float = 2.0
    t_shift: float = 0.1
    layer_penalty_factor: float = 5.0
    position_temperature: float = 5.0
    class_temperature: float = 0.0
    denoise: bool = True
    preprocess_prompt: bool = True
    postprocess_output: bool = True
    audio_chunk_duration: float = 15.0
    audio_chunk_threshold: float = 30.0

    @classmethod
    def from_dict(cls, kwargs_dict):
        valid_keys = {f.name for f in fields(cls)}
        filtered = {k: v for k, v in kwargs_dict.items() if k in valid_keys}
        return cls(**filtered)


@dataclass
class GenerationTask:
    batch_size: int
    texts: List[str]
    target_lens: List[int]
    langs: List[Optional[str]]
    instructs: List[Optional[str]]
    ref_texts: List[Optional[str]]
    ref_audio_tokens: List[Optional[torch.Tensor]]
    ref_rms: List[Optional[float]]
    speed: Optional[List[float]] = None

    def get_indices(self, config: iceAudioGenerationConfig, frame_rate: int):
        threshold = int(config.audio_chunk_threshold * frame_rate)
        short_idx = [i for i, l in enumerate(self.target_lens) if l <= threshold]
        long_idx = [i for i, l in enumerate(self.target_lens) if l > threshold]
        return short_idx, long_idx

    def slice_task(self, indices: List[int]):
        if not indices:
            return None
        return GenerationTask(
            batch_size=len(indices),
            texts=[self.texts[i] for i in indices],
            target_lens=[self.target_lens[i] for i in indices],
            langs=[self.langs[i] for i in indices],
            instructs=[self.instructs[i] for i in indices],
            ref_texts=[self.ref_texts[i] for i in indices],
            ref_audio_tokens=[self.ref_audio_tokens[i] for i in indices],
            ref_rms=[self.ref_rms[i] for i in indices],
            speed=[self.speed[i] for i in indices] if self.speed else None,
        )


@dataclass
class iceAudioModelOutput(ModelOutput):
    loss: Optional[torch.Tensor] = None
    logits: Optional[torch.Tensor] = None

def _resolve_model_path(name_or_path: str) -> str:
    if os.path.isdir(name_or_path):
        return name_or_path
    from huggingface_hub import snapshot_download

    return snapshot_download(name_or_path)

class IceAudioRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = float(eps)

    def forward(self, x):
        variance = x.float().pow(2).mean(-1, keepdim=True)
        return (x.float() * torch.rsqrt(variance + self.eps) * self.weight).to(x.dtype)


def _rotate_half(x):
    half = x.shape[-1] // 2
    return torch.cat((-x[..., half:], x[..., :half]), dim=-1)


class IceAudioRotaryEmbedding(nn.Module):
    def __init__(self, dim, base=1000000.0):
        super().__init__()
        self.dim = int(dim)
        self.base = float(base)
        inv = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.float32) / self.dim))
        self.register_buffer('inv_freq', inv, persistent=True)

    def forward(self, position_ids, dtype):
        inv = self.inv_freq.to(position_ids.device)
        freqs = torch.einsum('bs,d->bsd', position_ids.float(), inv)
        emb = torch.cat((freqs, freqs), dim=-1)
        return emb.cos().to(dtype), emb.sin().to(dtype)


class IceAudioMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.gate = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        self.act = getattr(config, 'hidden_act', 'silu')

    def forward(self, x):
        if self.act == 'silu':
            h = F.silu(self.gate(x))
        elif self.act == 'gelu':
            h = F.gelu(self.gate(x))
        else:
            h = F.silu(self.gate(x))
        return self.down(h * self.up(x))


class IceAudioAttention(nn.Module):
    def __init__(self, config, layer_idx=0):
        super().__init__()
        self.layer_idx = layer_idx
        self.num_heads = int(config.num_attention_heads)
        self.num_kv_heads = int(getattr(config, 'num_key_value_heads', self.num_heads))
        self.head_dim = int(getattr(config, 'head_dim', config.hidden_size // self.num_heads))
        self.groups = self.num_heads // self.num_kv_heads
        bias = bool(getattr(config, 'attention_bias', False))
        self.query = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=bias)
        self.key = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=bias)
        self.value = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=bias)
        self.q_norm = IceAudioRMSNorm(self.head_dim, getattr(config, 'rms_norm_eps', 1e-6))
        self.k_norm = IceAudioRMSNorm(self.head_dim, getattr(config, 'rms_norm_eps', 1e-6))
        self.output = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
        self.rope = IceAudioRotaryEmbedding(self.head_dim, getattr(config, 'rope_theta', 1000000.0))

    def _repeat_kv(self, x):
        if self.groups == 1:
            return x
        b, h, s, d = x.shape
        return x[:, :, None, :, :].expand(b, h, self.groups, s, d).reshape(b, h * self.groups, s, d)

    def forward(self, hidden_states, attention_mask=None, position_ids=None):
        bsz, q_len, _ = hidden_states.shape
        q = self.query(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.key(hidden_states).view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.value(hidden_states).view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
        q = self.q_norm(q)
        k = self.k_norm(k)
        if position_ids is None:
            position_ids = torch.arange(q_len, device=hidden_states.device).unsqueeze(0).expand(bsz, -1)
        cos, sin = self.rope(position_ids, q.dtype)
        cos, sin = cos[:, None, :, :], sin[:, None, :, :]
        q = q * cos + _rotate_half(q) * sin
        k = k * cos + _rotate_half(k) * sin
        k = self._repeat_kv(k)
        v = self._repeat_kv(v)
        scores = torch.matmul(q.float(), k.float().transpose(-2, -1)) / math.sqrt(self.head_dim)
        causal = torch.triu(torch.ones(q_len, q_len, device=hidden_states.device, dtype=torch.bool), diagonal=1)
        scores = scores.masked_fill(causal[None, None], torch.finfo(scores.dtype).min)
        if attention_mask is not None:
            mask_value = torch.finfo(scores.dtype).min
            if attention_mask.ndim == 2:
                if attention_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~attention_mask[:, None, None, :], mask_value)
                else:
                    scores = scores + (1.0 - attention_mask[:, None, None, :]).float() * mask_value
            elif attention_mask.ndim == 4:
                if attention_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~attention_mask, mask_value)
                else:
                    scores = scores + attention_mask.float()
        probs = F.softmax(scores, dim=-1, dtype=torch.float32)
        out = torch.matmul(probs, v.float()).to(q.dtype)
        out = out.transpose(1, 2).contiguous().view(bsz, q_len, self.num_heads * self.head_dim)
        return self.output(out)


class IceAudioBlock(nn.Module):
    def __init__(self, config, layer_idx=0):
        super().__init__()
        self.pre_norm = IceAudioRMSNorm(config.hidden_size, getattr(config, 'rms_norm_eps', 1e-6))
        self.attention = IceAudioAttention(config, layer_idx=layer_idx)
        self.post_norm = IceAudioRMSNorm(config.hidden_size, getattr(config, 'rms_norm_eps', 1e-6))
        self.mlp = IceAudioMLP(config)

    def forward(self, x, attention_mask=None, position_ids=None):
        x = x + self.attention(self.pre_norm(x), attention_mask, position_ids)
        x = x + self.mlp(self.post_norm(x))
        return x


class IceAudioLanguageModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        data = config.to_dict() if hasattr(config, 'to_dict') else dict(config)
        data['model_type'] = 'qwen3'
        data.pop('architectures', None)
        qcfg = Qwen3Config.from_dict(data)
        self.config = qcfg
        self.model = Qwen3Model(qcfg)

    def get_input_embeddings(self):
        return self.model.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.model.set_input_embeddings(value)

    def get_output_embeddings(self):
        return None

    def set_output_embeddings(self, value):
        raise NotImplementedError('Qwen3Model has no LM head in Ice Audio')

    def forward(self, input_ids=None, inputs_embeds=None, attention_mask=None, position_ids=None, return_dict=True, **kwargs):
        return self.model(input_ids=input_ids, inputs_embeds=inputs_embeds,
                          attention_mask=attention_mask, position_ids=position_ids,
                          return_dict=return_dict, **kwargs)

class iceAudio(PreTrainedModel):
    _supports_flex_attn = True
    _supports_flash_attn_2 = True
    _supports_sdpa = True
    config_class = iceAudioConfig

    def __init__(self, config: iceAudioConfig, language_model: Optional[PreTrainedModel] = None):
        super().__init__(config)

        if language_model is not None:

            self.language_model = language_model
        else:

            self.language_model = IceAudioLanguageModel(self.config.language_model_config)

        self.audio_embeddings = nn.Embedding(
            config.num_audio_codebook * config.audio_vocab_size,
            self.config.language_model_config.hidden_size,
        )
        self.register_buffer(
            "codebook_layer_offsets",
            torch.arange(config.num_audio_codebook) * config.audio_vocab_size,
        )

        self.audio_heads = nn.Linear(
            self.config.language_model_config.hidden_size,
            config.num_audio_codebook * config.audio_vocab_size,
            bias=False,
        )

        self.acoustic_adapter = None
        if getattr(config, "acoustic_adapter_enabled", True):
            self.acoustic_adapter = AcousticProsodyAdapter(
                hidden_size=self.config.language_model_config.hidden_size,
                ff_size=getattr(config, "acoustic_adapter_ff_size", 4096),
                num_layers=getattr(config, "acoustic_adapter_layers", 7),
            )

        self.normalized_audio_codebook_weights = [
            w / sum(config.audio_codebook_weights)
            for w in config.audio_codebook_weights
        ]

        self.post_init()

        self.text_tokenizer = None
        self.audio_tokenizer = None
        self.duration_estimator = None
        self.sampling_rate = None
        self._asr_pipe = None

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
        train_mode = kwargs.pop("train", False)
        load_asr = kwargs.pop("load_asr", False)
        asr_model_name = kwargs.pop("asr_model_name", "openai/whisper-large-v3-turbo")

        _prev_disable = logging.root.manager.disable
        logging.disable(logging.INFO)

        try:

            resolved_path = _resolve_model_path(pretrained_model_name_or_path)

            model = super().from_pretrained(resolved_path, *args, **kwargs)

            if not train_mode:
                model.text_tokenizer = AutoTokenizer.from_pretrained(resolved_path, trust_remote_code=True)

                audio_tokenizer_path = os.path.join(resolved_path, "audio_tokenizer")

                if not os.path.isdir(audio_tokenizer_path):
                    audio_tokenizer_path = _resolve_model_path(
                        "eustlb/higgs-audio-v2-tokenizer"
                    )

                tokenizer_device = (
                    "cpu" if str(model.device).startswith("mps") else model.device
                )
                model.audio_tokenizer = HiggsAudioV2TokenizerModel.from_pretrained(
                    audio_tokenizer_path, device_map=tokenizer_device
                )
                model.feature_extractor = AutoFeatureExtractor.from_pretrained(
                    audio_tokenizer_path
                )

                model.sampling_rate = model.feature_extractor.sampling_rate

                model.duration_estimator = RuleDurationEstimator()

                if load_asr:
                    model.load_asr_model(model_name=asr_model_name)
        finally:
            logging.disable(_prev_disable)

        return model

    def load_asr_model(self, model_name: str = "openai/whisper-large-v3-turbo"):

        from transformers import pipeline as hf_pipeline

        logger.info("Loading ASR model %s ...", model_name)
        asr_dtype = (
            torch.float16 if str(self.device).startswith("cuda") else torch.float32
        )

        model_name = _resolve_model_path(model_name)

        self._asr_pipe = hf_pipeline(
            "automatic-speech-recognition",
            model=model_name,
            dtype=asr_dtype,
            device_map=self.device,
        )
        logger.info("ASR model loaded on %s.", self.device)

    @torch.inference_mode()
    def transcribe(
        self,
        audio: Union[str, tuple],
    ) -> str:

        if self._asr_pipe is None:
            raise RuntimeError(
                "ASR model is not loaded. Call model.load_asr_model() first."
            )

        if isinstance(audio, str):
            return self._asr_pipe(audio)["text"].strip()
        else:
            waveform, sr = audio
            if isinstance(waveform, torch.Tensor):
                waveform = waveform.cpu().numpy()
            waveform = np.squeeze(waveform)
            audio_input = {
                "array": waveform,
                "sampling_rate": sr,
            }
            return self._asr_pipe(audio_input)["text"].strip()

    def get_input_embeddings(self):
        return self.language_model.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.language_model.set_input_embeddings(value)

    def _prepare_embed_inputs(
        self, input_ids: torch.Tensor, audio_mask: torch.Tensor
    ) -> torch.Tensor:

        text_embeds = self.get_input_embeddings()(input_ids[:, 0, :])

        shifted_ids = (
            input_ids * audio_mask.unsqueeze(1)
        ) + self.codebook_layer_offsets.view(1, -1, 1)

        audio_embeds = self.audio_embeddings(shifted_ids).sum(dim=1)
        if self.acoustic_adapter is not None:
            audio_embeds = self.acoustic_adapter.condition_audio_embeddings(
                audio_embeds, audio_mask
            )

        return torch.where(audio_mask.unsqueeze(-1), audio_embeds, text_embeds)

    def forward(
        self,
        input_ids: torch.LongTensor,
        audio_mask: torch.Tensor,
        labels: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        document_ids: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
    ):

        inputs_embeds = self._prepare_embed_inputs(input_ids, audio_mask)

        if attention_mask is None and document_ids is not None:
            if not _flex_attention_available:
                raise RuntimeError(
                    "flex_attention is not available in the current environment. "
                    "If you do not need flex_attention, set "
                    '"attn_implementation": "sdpa" in your training config.'
                )
            attention_mask = create_block_mask(
                _get_packed_mask(
                    document_ids[0].to(inputs_embeds.device),
                ),
                B=None,
                H=None,
                Q_LEN=input_ids.size(-1),
                KV_LEN=input_ids.size(-1),
                _compile=True,
                device=inputs_embeds.device,
            )

        llm_outputs = self.language_model(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            return_dict=True,
            position_ids=position_ids,
        )
        hidden_states = llm_outputs[0]
        if self.acoustic_adapter is not None:
            hidden_states = self.acoustic_adapter.refine_hidden_states(
                hidden_states, audio_mask
            )

        loss = None

        batch_size, seq_len, _ = hidden_states.shape
        logits_flat = self.audio_heads(hidden_states)
        audio_logits = logits_flat.view(
            batch_size,
            seq_len,
            self.config.num_audio_codebook,
            self.config.audio_vocab_size,
        ).permute(0, 2, 1, 3)

        if labels is not None:

            per_token_loss = torch.nn.functional.cross_entropy(
                audio_logits.permute(0, 3, 1, 2),
                labels,
                reduction="none",
                ignore_index=-100,
            )

            valid_mask = (labels != -100).float()

            layer_means = (per_token_loss * valid_mask).sum(
                dim=(0, 2)
            ) / valid_mask.sum(dim=(0, 2)).clamp(min=1.0)

            weights = torch.tensor(
                self.normalized_audio_codebook_weights, device=audio_logits.device
            )
            loss = (layer_means * weights).sum()

        return iceAudioModelOutput(
            loss=loss,
            logits=audio_logits,
        )

    def supported_language_ids(self) -> set[str]:
        return LANG_IDS

    def supported_language_names(self) -> set[str]:
        return LANG_NAMES

    @torch.inference_mode()
    def generate(
        self,
        text: Union[str, list[str]],
        language: Union[str, list[str], None] = None,
        ref_text: Union[str, list[str], None] = None,
        ref_audio: Union[
            str,
            list[str],
            tuple[torch.Tensor, int],
            list[tuple[torch.Tensor, int]],
            None,
        ] = None,
        voice_clone_prompt: Union[
            VoiceClonePrompt, list[VoiceClonePrompt], None
        ] = None,
        instruct: Union[str, list[str], None] = None,
        duration: Union[float, list[Optional[float]], None] = None,
        speed: Union[float, list[Optional[float]], None] = None,
        generation_config: Optional[iceAudioGenerationConfig] = None,
        **kwargs,
    ) -> list[np.ndarray]:

        if self.audio_tokenizer is None or self.text_tokenizer is None:
            raise RuntimeError(
                "Model is not loaded with audio/text tokenizers. Make sure you "
                "loaded the model with iceAudio.from_pretrained()."
            )
        gen_config = (
            generation_config
            if generation_config is not None
            else iceAudioGenerationConfig.from_dict(kwargs)
        )

        self.eval()

        full_task = self._preprocess_all(
            text=text,
            language=language,
            ref_text=ref_text,
            ref_audio=ref_audio,
            voice_clone_prompt=voice_clone_prompt,
            instruct=instruct,
            preprocess_prompt=gen_config.preprocess_prompt,
            speed=speed,
            duration=duration,
        )

        short_idx, long_idx = full_task.get_indices(
            gen_config, self.audio_tokenizer.config.frame_rate
        )

        results = [None] * full_task.batch_size

        if short_idx:
            short_task = full_task.slice_task(short_idx)
            short_results = self._generate_iterative(short_task, gen_config)
            for idx, res in zip(short_idx, short_results):
                results[idx] = res

        if long_idx:
            long_task = full_task.slice_task(long_idx)
            long_results = self._generate_chunked(long_task, gen_config)
            for idx, res in zip(long_idx, long_results):
                results[idx] = res

        generated_audios = []
        for i in range(full_task.batch_size):
            assert results[i] is not None, f"Result {i} was not generated"
            generated_audios.append(
                self._decode_and_post_process(
                    results[i], full_task.ref_rms[i], gen_config
                )
            )

        return generated_audios

    def create_voice_clone_prompt(
        self,
        ref_audio: Union[str, tuple[torch.Tensor, int]],
        ref_text: Optional[str] = None,
        preprocess_prompt: bool = True,
    ) -> VoiceClonePrompt:

        if self.audio_tokenizer is None:
            raise RuntimeError(
                "Audio tokenizer is not loaded. Make sure you loaded the model "
                "with iceAudio.from_pretrained()."
            )

        if isinstance(ref_audio, str):
            ref_wav = load_audio(ref_audio, self.sampling_rate)
        else:
            waveform, sr = ref_audio
            if isinstance(waveform, torch.Tensor):
                waveform = waveform.cpu().numpy()
            if waveform.ndim == 1:
                waveform = waveform[np.newaxis, :]
            if waveform.shape[0] > 1:
                waveform = np.mean(waveform, axis=0, keepdims=True)
            if sr != self.sampling_rate:
                waveform = torchaudio.functional.resample(
                    torch.from_numpy(waveform),
                    orig_freq=sr,
                    new_freq=self.sampling_rate,
                ).numpy()
            ref_wav = waveform

        ref_rms = float(np.sqrt(np.mean(ref_wav**2)))
        if 0 < ref_rms < 0.1:
            ref_wav = ref_wav * 0.1 / ref_rms

        if preprocess_prompt:

            if ref_text is None:
                ref_wav = trim_long_audio(
                    ref_wav, self.sampling_rate, trim_threshold=20.0
                )
            ref_wav = remove_silence(
                ref_wav,
                self.sampling_rate,
                mid_sil=200,
                lead_sil=100,
                trail_sil=200,
            )
            if ref_wav.shape[-1] == 0:
                raise ValueError(
                    "Reference audio is empty after silence removal. "
                    "Try setting preprocess_prompt=False."
                )

        ref_duration = ref_wav.shape[-1] / self.sampling_rate
        if ref_duration > 20.0:
            logger.warning(
                "Reference audio is %.1fs long (>20s). This may cause slower "
                "generation, higher memory usage, and degraded voice cloning "
                "quality. We recommend trimming it to 3-10s.",
                ref_duration,
            )

        if ref_text is None:
            if self._asr_pipe is None:
                logger.info("ASR model not loaded yet, loading on-the-fly ...")
                self.load_asr_model()
            ref_text = self.transcribe((ref_wav, self.sampling_rate))
            logger.debug("Auto-transcribed ref_text: %s", ref_text)

        chunk_size = self.audio_tokenizer.config.hop_length
        clip_size = int(ref_wav.shape[-1] % chunk_size)
        ref_wav = ref_wav[:, :-clip_size] if clip_size > 0 else ref_wav

        ref_wav_tensor = torch.from_numpy(ref_wav).to(self.audio_tokenizer.device)
        ref_audio_tokens = self.audio_tokenizer.encode(
            ref_wav_tensor.unsqueeze(0),
        ).audio_codes.squeeze(
            0
        )

        if preprocess_prompt:
            ref_text = add_punctuation(ref_text)

        return VoiceClonePrompt(
            ref_audio_tokens=ref_audio_tokens,
            ref_text=ref_text,
            ref_rms=ref_rms,
        )

    def _decode_and_post_process(
        self,
        tokens: Union[torch.Tensor, List[torch.Tensor]],
        rms: Union[float, None],
        gen_config: iceAudioGenerationConfig,
    ) -> np.ndarray:

        tokenizer_device = self.audio_tokenizer.device
        if isinstance(tokens, list):
            chunk_audios = [
                self.audio_tokenizer.decode(t.to(tokenizer_device).unsqueeze(0))
                .audio_values[0]
                .cpu()
                .numpy()
                for t in tokens
            ]
            audio_waveform = cross_fade_chunks(chunk_audios, self.sampling_rate)
        else:
            audio_waveform = (
                self.audio_tokenizer.decode(tokens.to(tokenizer_device).unsqueeze(0))
                .audio_values[0]
                .cpu()
                .numpy()
            )

        audio_waveform = self._post_process_audio(
            audio_waveform,
            postprocess_output=gen_config.postprocess_output,
            ref_rms=rms,
        )
        return audio_waveform.squeeze(0)

    def _post_process_audio(
        self,
        generated_audio: np.ndarray,
        postprocess_output: bool,
        ref_rms: Union[float, None],
    ) -> np.ndarray:

        if postprocess_output:
            generated_audio = remove_silence(
                generated_audio,
                self.sampling_rate,
                mid_sil=500,
                lead_sil=100,
                trail_sil=100,
            )

        if ref_rms is not None and ref_rms < 0.1:
            generated_audio = generated_audio * ref_rms / 0.1
        elif ref_rms is None:
            peak = np.abs(generated_audio).max()
            if peak > 1e-6:
                generated_audio = generated_audio / peak * 0.5

        generated_audio = fade_and_pad_audio(
            generated_audio,
            sample_rate=self.sampling_rate,
        )
        return generated_audio

    def _generate_chunked(
        self, task: GenerationTask, gen_config: iceAudioGenerationConfig
    ) -> List[List[torch.Tensor]]:

        all_chunks = []
        for i in range(task.batch_size):
            avg_tokens_per_char = task.target_lens[i] / len(task.texts[i])
            text_chunk_len = int(
                gen_config.audio_chunk_duration
                * self.audio_tokenizer.config.frame_rate
                / avg_tokens_per_char
            )
            chunks = chunk_text_punctuation(
                text=task.texts[i],
                chunk_len=text_chunk_len,
                min_chunk_len=3,
            )
            logger.debug(f"Item {i} chunked into {len(chunks)} pieces: {chunks}")
            all_chunks.append(chunks)

        has_ref = [t is not None for t in task.ref_audio_tokens]
        assert all(has_ref) or not any(has_ref), (
            "Chunked inference requires all items to either have or not have "
            "ref_audio. Mixed ref/non-ref is not supported."
        )

        max_num_chunks = max(len(c) for c in all_chunks)

        chunk_results = [[] for _ in range(task.batch_size)]

        def _run_batch(indices, texts, ref_audios, ref_texts):
            speed_list = task.speed
            target_lens = [
                self._estimate_target_tokens(
                    texts[j],
                    ref_texts[j],
                    ref_audios[j].size(-1) if ref_audios[j] is not None else None,
                    speed=speed_list[i] if speed_list else 1.0,
                )
                for j, i in enumerate(indices)
            ]
            sub_task = GenerationTask(
                batch_size=len(indices),
                texts=texts,
                target_lens=target_lens,
                langs=[task.langs[i] for i in indices],
                instructs=[task.instructs[i] for i in indices],
                ref_texts=ref_texts,
                ref_audio_tokens=ref_audios,
                ref_rms=[task.ref_rms[i] for i in indices],
                speed=[task.speed[i] for i in indices] if task.speed else None,
            )
            gen_tokens = self._generate_iterative(sub_task, gen_config)
            for j, idx in enumerate(indices):
                chunk_results[idx].append(gen_tokens[j])

        if all(has_ref):

            for ci in range(max_num_chunks):
                indices = [i for i in range(task.batch_size) if ci < len(all_chunks[i])]
                if not indices:
                    continue
                _run_batch(
                    indices,
                    texts=[all_chunks[i][ci] for i in indices],
                    ref_audios=[task.ref_audio_tokens[i] for i in indices],
                    ref_texts=[task.ref_texts[i] for i in indices],
                )
        else:

            indices_0 = [i for i in range(task.batch_size) if len(all_chunks[i]) > 0]
            _run_batch(
                indices_0,
                texts=[all_chunks[i][0] for i in indices_0],
                ref_audios=[None] * len(indices_0),
                ref_texts=[None] * len(indices_0),
            )
            first_chunk_map = {idx: chunk_results[idx][0] for idx in indices_0}

            for ci in range(1, max_num_chunks):
                indices = [i for i in range(task.batch_size) if ci < len(all_chunks[i])]
                if not indices:
                    continue
                _run_batch(
                    indices,
                    texts=[all_chunks[i][ci] for i in indices],
                    ref_audios=[first_chunk_map[i] for i in indices],
                    ref_texts=[all_chunks[i][0] for i in indices],
                )

        return chunk_results

    def _preprocess_all(
        self,
        text: Union[str, list[str]],
        language: Union[str, list[str], None] = None,
        ref_text: Union[str, list[str], None] = None,
        ref_audio: Union[
            str,
            list[str],
            tuple[torch.Tensor, int],
            list[tuple[torch.Tensor, int]],
            None,
        ] = None,
        voice_clone_prompt: Union[
            VoiceClonePrompt, list[VoiceClonePrompt], None
        ] = None,
        instruct: Union[str, list[str], None] = None,
        preprocess_prompt: bool = True,
        speed: Union[float, list[Optional[float]], None] = None,
        duration: Union[float, list[Optional[float]], None] = None,
    ) -> GenerationTask:

        if isinstance(text, str):
            text_list = [text]
        else:
            assert isinstance(
                text, list
            ), "text should be a string or a list of strings"
            text_list = text
        batch_size = len(text_list)

        language_list = self._ensure_list(language, batch_size)
        language_list = [_resolve_language(lang) for lang in language_list]
        instruct_list = self._ensure_list(instruct, batch_size)
        for i, s in enumerate(instruct_list):
            if s is None:
                continue
            use_zh = bool(text_list[i] and _ZH_RE.search(text_list[i]))
            instruct_list[i] = _resolve_instruct(s, use_zh=use_zh)

        if voice_clone_prompt is not None and (
            ref_text is not None or ref_audio is not None
        ):
            logger.warning(
                "Both voice_clone_prompt and ref_text/ref_audio are provided. "
                "ref_text/ref_audio will be ignored."
            )
        if voice_clone_prompt is None and ref_audio is not None:

            ref_text_list = self._ensure_list(ref_text, batch_size, auto_repeat=False)
            ref_audio_list = self._ensure_list(ref_audio, batch_size, auto_repeat=False)

            voice_clone_prompt = []
            for i in range(len(ref_text_list)):
                voice_clone_prompt.append(
                    self.create_voice_clone_prompt(
                        ref_audio=ref_audio_list[i],
                        ref_text=ref_text_list[i],
                        preprocess_prompt=preprocess_prompt,
                    )
                )

        voice_clone_prompt_list = self._ensure_list(voice_clone_prompt, batch_size)
        if voice_clone_prompt_list[0] is not None:
            ref_text_list = [vc.ref_text for vc in voice_clone_prompt_list]
            ref_audio_tokens_list = [
                vc.ref_audio_tokens for vc in voice_clone_prompt_list
            ]
            ref_rms_list = [vc.ref_rms for vc in voice_clone_prompt_list]
        else:
            ref_text_list = [None] * batch_size
            ref_audio_tokens_list = [None] * batch_size
            ref_rms_list = [None] * batch_size

        if speed is not None:
            if isinstance(speed, (int, float)):
                user_speed = [float(speed)] * batch_size
            else:
                user_speed = list(speed)
        else:
            user_speed = None

        if duration is not None:
            if isinstance(duration, (int, float)):
                durations = [float(duration)] * batch_size
            else:
                durations = list(duration)
        else:
            durations = None

        num_target_tokens_list = []
        for i in range(batch_size):
            has_dur = durations is not None and durations[i] is not None
            item_speed = 1.0 if has_dur else (user_speed[i] if user_speed else 1.0)
            est = self._estimate_target_tokens(
                text_list[i],
                ref_text_list[i],
                ref_audio_tokens_list[i].size(-1)
                if ref_audio_tokens_list[i] is not None
                else None,
                speed=item_speed,
            )
            num_target_tokens_list.append(est)

        speed_list: Optional[List[float]] = None
        if durations is not None:
            frame_rate = self.audio_tokenizer.config.frame_rate
            speed_list = []
            for i in range(batch_size):
                if durations[i] is not None:
                    target_tokens = max(1, int(durations[i] * frame_rate))
                    est = num_target_tokens_list[i]
                    speed_list.append(est / target_tokens if target_tokens > 0 else 1.0)
                    num_target_tokens_list[i] = target_tokens
                else:
                    s = user_speed[i] if user_speed else None
                    speed_list.append(s if s is not None else 1.0)
        elif user_speed is not None:
            speed_list = [s if s is not None else 1.0 for s in user_speed]

        return GenerationTask(
            batch_size=batch_size,
            texts=text_list,
            target_lens=num_target_tokens_list,
            langs=language_list,
            instructs=instruct_list,
            ref_texts=ref_text_list,
            ref_audio_tokens=ref_audio_tokens_list,
            ref_rms=ref_rms_list,
            speed=speed_list,
        )

    def _estimate_target_tokens(self, text, ref_text, num_ref_audio_tokens, speed=1.0):
        if num_ref_audio_tokens is None or ref_text is None or len(ref_text) == 0:
            ref_text = "Nice to meet you."
            num_ref_audio_tokens = 25

        est = self.duration_estimator.estimate_duration(
            text, ref_text, num_ref_audio_tokens
        )
        if speed > 0 and speed != 1.0:
            est = est / speed
        return max(1, int(est))

    def _ensure_list(
        self, x: Union[Any, List[Any]], batch_size: int, auto_repeat: bool = True
    ) -> List[Any]:
        x_list = x if isinstance(x, list) else [x]
        if len(x_list) not in (
            1,
            batch_size,
        ):
            raise ValueError(
                f"should be either the number of the text or 1, but got {len(x_list)}"
            )
        if auto_repeat and len(x_list) == 1 and batch_size is not None:
            x_list = x_list * batch_size
        return x_list

    def _prepare_inference_inputs(
        self,
        text: str,
        num_target_tokens: int,
        ref_text: Optional[str] = None,
        ref_audio_tokens: Optional[torch.Tensor] = None,
        lang: Optional[str] = None,
        instruct: Optional[str] = None,
        denoise: bool = True,
    ):
        style_text = ""
        legacy_runtime = getattr(self,'runtime_mode','independent_training') == 'legacy_audio_compatible'
        if denoise and ref_audio_tokens is not None:
            style_text += "<|denoise|>" if legacy_runtime else "<|ice_r_instruction|>denoise<|ice_r_text_end|>"
        lang_str = lang if lang else "None"
        instruct_str = instruct if instruct else "None"
        if legacy_runtime:
            style_text += f"<|lang_start|>{lang_str}<|lang_end|>"
            style_text += f"<|instruct_start|>{instruct_str}<|instruct_end|>"
        else:
            style_text += f"<|ice_r_lang|>{lang_str}<|ice_r_text_end|>"
            style_text += f"<|ice_r_instruction|>{instruct_str}<|ice_r_text_end|>"

        style_tokens = (
            self.text_tokenizer(style_text, return_tensors="pt")
            .input_ids.repeat(self.config.num_audio_codebook, 1)
            .unsqueeze(0)
        ).to(
            self.device
        )

        full_text = _combine_text(ref_text=ref_text, text=text)
        wrapped_text = f"<|text_start|>{full_text}<|text_end|>" if legacy_runtime else f"<|ice_r_text|>{full_text}<|ice_r_text_end|>"
        text_tokens = (
            _tokenize_with_nonverbal_tags(wrapped_text, self.text_tokenizer)
            .repeat(self.config.num_audio_codebook, 1)
            .unsqueeze(0)
        ).to(
            self.device
        )

        target_audio_tokens = torch.full(
            (1, self.config.num_audio_codebook, num_target_tokens),
            self.config.audio_mask_id,
            dtype=torch.long,
            device=self.device,
        )

        parts = [style_tokens, text_tokens]
        if ref_audio_tokens is not None:
            parts.append(ref_audio_tokens.unsqueeze(0).to(self.device))
        parts.append(target_audio_tokens)
        cond_input_ids = torch.cat(parts, dim=2)

        cond_total_length = cond_input_ids.shape[2]
        cond_audio_start_idx = cond_total_length - num_target_tokens
        if ref_audio_tokens is not None:
            cond_audio_start_idx -= ref_audio_tokens.size(-1)

        cond_audio_mask = torch.zeros(
            1, cond_total_length, dtype=torch.bool, device=self.device
        )
        cond_audio_mask[0, cond_audio_start_idx:] = True

        return {
            "input_ids": cond_input_ids,
            "audio_mask": cond_audio_mask,
        }

    def _generate_iterative(
        self, task: GenerationTask, gen_config: iceAudioGenerationConfig
    ) -> List[torch.Tensor]:

        B = task.batch_size

        for i in range(B):
            logger.debug(
                "Item %d — text: %s | ref_text: %s | instruct: %s | lang: %s | target_tokens: %d",
                i,
                task.texts[i],
                task.ref_texts[i],
                task.instructs[i],
                task.langs[i],
                task.target_lens[i],
            )

        inputs_list = [
            self._prepare_inference_inputs(
                task.texts[i],
                task.target_lens[i],
                task.ref_texts[i],
                task.ref_audio_tokens[i],
                task.langs[i],
                task.instructs[i],
                gen_config.denoise,
            )
            for i in range(B)
        ]

        c_lens = [inp["input_ids"].size(2) for inp in inputs_list]
        max_c_len = max(c_lens)
        pad_id = self.config.audio_mask_id

        batch_input_ids = torch.full(
            (2 * B, self.config.num_audio_codebook, max_c_len),
            pad_id,
            dtype=torch.long,
            device=self.device,
        )
        batch_audio_mask = torch.zeros(
            (2 * B, max_c_len), dtype=torch.bool, device=self.device
        )
        batch_attention_mask = torch.zeros(
            (2 * B, 1, max_c_len, max_c_len), dtype=torch.bool, device=self.device
        )

        for i, inp in enumerate(inputs_list):
            c_len, u_len = c_lens[i], task.target_lens[i]

            batch_input_ids[i, :, :c_len] = inp["input_ids"]
            batch_audio_mask[i, :c_len] = inp["audio_mask"]
            batch_attention_mask[i, :, :c_len, :c_len] = True

            batch_input_ids[B + i, :, :u_len] = inp["input_ids"][..., -u_len:]
            batch_audio_mask[B + i, :u_len] = inp["audio_mask"][..., -u_len:]
            batch_attention_mask[B + i, :, :u_len, :u_len] = True
            if max_c_len > u_len:
                pad_diag = torch.arange(u_len, max_c_len, device=self.device)
                batch_attention_mask[B + i, :, pad_diag, pad_diag] = True

        tokens = torch.full(
            (B, self.config.num_audio_codebook, max(task.target_lens)),
            self.config.audio_mask_id,
            dtype=torch.long,
            device=self.device,
        )

        timesteps = _get_time_steps(
            t_start=0.0,
            t_end=1.0,
            num_step=gen_config.num_step,
            t_shift=gen_config.t_shift,
        ).tolist()
        schedules = []
        for t_len in task.target_lens:
            total_mask = t_len * self.config.num_audio_codebook
            rem = total_mask
            sched = []
            for step in range(gen_config.num_step):
                num = (
                    rem
                    if step == gen_config.num_step - 1
                    else min(
                        math.ceil(total_mask * (timesteps[step + 1] - timesteps[step])),
                        rem,
                    )
                )
                sched.append(int(num))
                rem -= int(num)
            schedules.append(sched)

        layer_ids = torch.arange(
            self.config.num_audio_codebook, device=self.device
        ).view(1, -1, 1)

        for step in range(gen_config.num_step):
            batch_logits = self(
                input_ids=batch_input_ids,
                audio_mask=batch_audio_mask,
                attention_mask=batch_attention_mask,
            ).logits.to(torch.float32)

            for i in range(B):
                k = schedules[i][step]
                if k <= 0:
                    continue

                c_len, t_len = c_lens[i], task.target_lens[i]

                c_logits = batch_logits[i : i + 1, :, c_len - t_len : c_len, :]
                u_logits = batch_logits[B + i : B + i + 1, :, :t_len, :]

                pred_tokens, scores = self._predict_tokens_with_scoring(
                    c_logits, u_logits, gen_config
                )

                scores = scores - (layer_ids * gen_config.layer_penalty_factor)

                if gen_config.position_temperature > 0.0:
                    scores = _gumbel_sample(scores, gen_config.position_temperature)

                sample_tokens = tokens[i : i + 1, :, :t_len]
                scores.masked_fill_(
                    sample_tokens != self.config.audio_mask_id, -float("inf")
                )

                _, topk_idx = torch.topk(scores.flatten(), k)
                flat_tokens = sample_tokens.flatten()
                flat_tokens[topk_idx] = pred_tokens.flatten()[topk_idx]
                sample_tokens.copy_(flat_tokens.view_as(sample_tokens))

                tokens[i : i + 1, :, :t_len] = sample_tokens
                batch_input_ids[i : i + 1, :, c_len - t_len : c_len] = sample_tokens
                batch_input_ids[B + i : B + i + 1, :, :t_len] = sample_tokens

        return [tokens[i, :, : task.target_lens[i]] for i in range(B)]

    def _predict_tokens_with_scoring(self, c_logits, u_logits, gen_config):
        if gen_config.guidance_scale != 0:
            c_log_probs = F.log_softmax(c_logits, dim=-1)
            u_log_probs = F.log_softmax(u_logits, dim=-1)
            log_probs = torch.log_softmax(
                c_log_probs + gen_config.guidance_scale * (c_log_probs - u_log_probs),
                dim=-1,
            )
        else:
            log_probs = F.log_softmax(c_logits, dim=-1)

        log_probs[..., self.config.audio_mask_id] = -float("inf")

        if gen_config.class_temperature > 0.0:
            filtered_probs = _filter_top_k(log_probs, ratio=0.1)
            pred_tokens = _gumbel_sample(
                filtered_probs, gen_config.class_temperature
            ).argmax(dim=-1)
        else:
            pred_tokens = log_probs.argmax(dim=-1)

        confidence_scores = log_probs.max(dim=-1)[0]

        return pred_tokens, confidence_scores



class iceAudioForCausalLM(iceAudio):
    config_class = iceAudioConfig

def _get_packed_mask(document_ids):
    return partial(_mask_mod_packed, document_ids)


def _mask_mod_packed(document_ids, b, h, q_idx, kv_idx):

    same_doc = document_ids[q_idx] == document_ids[kv_idx]
    return same_doc


def _resolve_language(language: Optional[str]) -> Union[str, None]:
    

    if language is None or language.lower() == "none":
        return None
    if language in LANG_IDS:
        return language
    key = language.lower()
    if key in LANG_NAME_TO_ID:
        return LANG_NAME_TO_ID[key]
    raise ValueError(
        f"Unsupported language '{language}'. ICE-R supports exactly the 590 entries "
        "in languages.json; unsupported values are rejected explicitly."
    )


def _resolve_instruct(
    instruct: Optional[str], use_zh: bool = False
) -> Union[str, None]:

    if instruct is None:
        return None

    instruct_str = instruct.strip()
    if not instruct_str:
        return None

    raw_items = re.split(r"\s*[,，]\s*", instruct_str)
    raw_items = [x for x in raw_items if x]

    unknown = []
    normalised = []
    for raw in raw_items:
        n = raw.strip().lower()
        if n in _INSTRUCT_ALL_VALID:
            normalised.append(n)
        else:
            sug = difflib.get_close_matches(n, _INSTRUCT_ALL_VALID, n=1, cutoff=0.6)
            unknown.append((raw, n, sug[0] if sug else None))

    if unknown:
        lines = []
        for raw, n, sug in unknown:
            if sug:
                lines.append(f"  '{raw}' -> '{n}' (unsupported; did you mean '{sug}'?)")
            else:
                lines.append(f"  '{raw}' -> '{n}' (unsupported)")
        err = (
            f"Unsupported instruct items found in {instruct_str}:\n"
            + "\n".join(lines)
            + "\n\nValid English items: "
            + ", ".join(sorted(_INSTRUCT_VALID_EN))
            + "\nValid Chinese items: "
            + "，".join(sorted(_INSTRUCT_VALID_ZH))
            + "\n\nTip: Use only English or only Chinese instructs. "
            "English instructs should use comma + space (e.g. "
            "'male, indian accent'),\nChinese instructs should use full-width "
            "comma (e.g. '男，河南话')."
        )
        raise ValueError(err)

    has_dialect = any(n.endswith("话") for n in normalised)
    has_accent = any(" accent" in n for n in normalised)

    if has_dialect and has_accent:
        raise ValueError(
            "Cannot mix Chinese dialect and English accent in a single instruct. "
            "Dialects are for Chinese speech, accents for English speech."
        )

    if has_dialect:
        use_zh = True
    elif has_accent:
        use_zh = False

    if use_zh:
        normalised = [_INSTRUCT_EN_TO_ZH.get(n, n) for n in normalised]
    else:
        normalised = [_INSTRUCT_ZH_TO_EN.get(n, n) for n in normalised]

    conflicts = []
    for cat in _INSTRUCT_MUTUALLY_EXCLUSIVE:
        hits = [n for n in normalised if n in cat]
        if len(hits) > 1:
            conflicts.append(hits)
    if conflicts:
        parts = []
        for group in conflicts:
            parts.append(" vs ".join(f"'{x}'" for x in group))
        raise ValueError(
            "Conflicting instruct items within the same category: "
            + "; ".join(parts)
            + ". Each category (gender, age, pitch, style, accent, dialect) "
            "allows at most one item."
        )

    has_zh = any(any("\u4e00" <= c <= "\u9fff" for c in n) for n in normalised)
    separator = "，" if has_zh else ", "

    return separator.join(normalised)


def _filter_top_k(logits: torch.Tensor, ratio: float = 0.1) -> torch.Tensor:
    k = math.ceil(ratio * logits.shape[-1])
    val, ind = logits.topk(k, dim=-1)
    probs = torch.full_like(logits, float("-inf"))
    probs.scatter_(-1, ind, val)
    return probs


def _gumbel_sample(logits: torch.Tensor, temperature: float) -> torch.Tensor:
    scaled_logits = logits / temperature
    u = torch.rand_like(scaled_logits)
    gumbel_noise = -torch.log(-torch.log(u + 1e-10) + 1e-10)
    return scaled_logits + gumbel_noise


def _get_time_steps(
    t_start: float = 0.0,
    t_end: float = 1.0,
    num_step: int = 10,
    t_shift: float = 1.0,
    device: torch.device = torch.device("cpu"),
) -> torch.Tensor:
    timesteps = torch.linspace(t_start, t_end, num_step + 1).to(device)
    timesteps = t_shift * timesteps / (1 + (t_shift - 1) * timesteps)
    return timesteps


_NONVERBAL_PATTERN = re.compile(
    r"\[(laughter|sigh|confirmation-en|question-en|question-ah|question-oh|"
    r"question-ei|question-yi|surprise-ah|surprise-oh|surprise-wa|"
    r"surprise-yo|dissatisfaction-hnn)\]"
)


def _tokenize_with_nonverbal_tags(text: str, tokenizer) -> torch.Tensor:
    parts = []
    last_end = 0
    for m in _NONVERBAL_PATTERN.finditer(text):
        if m.start() > last_end:
            segment = text[last_end : m.start()]
            ids = tokenizer(segment, add_special_tokens=False).input_ids
            if ids:
                parts.append(ids)
        tag_ids = tokenizer(m.group(), add_special_tokens=False).input_ids
        if tag_ids:
            parts.append(tag_ids)
        last_end = m.end()
    if last_end < len(text):
        segment = text[last_end:]
        ids = tokenizer(segment, add_special_tokens=False).input_ids
        if ids:
            parts.append(ids)

    if not parts:
        result = tokenizer(text, return_tensors="pt").input_ids
    else:
        combined = []
        for p in parts:
            combined.extend(p)
        result = torch.tensor([combined], dtype=torch.long)
    return result


def _combine_text(text, ref_text: Optional[str] = None) -> str:

    if ref_text:
        full_text = ref_text.strip() + " " + text.strip()
    else:
        full_text = text.strip()

    full_text = re.sub(r"[\r\n]+", "", full_text)

    full_text = full_text.replace("\uff08", "(").replace("\uff09", ")")

    full_text = re.sub(r"[ \t]+", " ", full_text)

    chinese_range = r"[\u4e00-\u9fff]"
    pattern = rf"(?<={chinese_range})\s+|\s+(?={chinese_range})"
    full_text = re.sub(pattern, "", full_text)

    return full_text

AutoConfig.register("ice-012-audio", iceAudioConfig, exist_ok=True)
AutoModel.register(iceAudioConfig, iceAudio)
AutoModelForCausalLM.register(iceAudioConfig, iceAudioForCausalLM)


class FunctionPreservingCapacity(nn.Module):
    def __init__(self, hidden_size=1024, expert_width=7560, experts=6):
        super().__init__()
        self.experts = nn.ModuleList([
            nn.Sequential(nn.Linear(hidden_size, expert_width, bias=False), nn.SiLU(), nn.Linear(expert_width, hidden_size, bias=False))
            for _ in range(experts)
        ])
        self.routing_memory = nn.Parameter(torch.zeros(8102719))
        self.output_gate = nn.Parameter(torch.zeros(1))
        for e in self.experts:
            nn.init.normal_(e[0].weight, mean=0.0, std=0.02)
            nn.init.zeros_(e[2].weight)
    def forward(self, x):
        if self.output_gate.detach().abs().item() == 0.0:
            return x
        y=sum(e(x) for e in self.experts) / len(self.experts)
        route=self.routing_memory.mean().to(dtype=x.dtype, device=x.device)
        return x + self.output_gate.to(dtype=x.dtype) * (y + route)

class IceRLanguageModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config=config
        self.model=IndependentBackbone(config)
    def get_input_embeddings(self): return self.model.get_input_embeddings()
    def set_input_embeddings(self,value): self.model.embed=value
    def forward(self,input_ids=None,inputs_embeds=None,attention_mask=None,return_dict=True,position_ids=None,**kwargs):
        return self.model(input_ids=input_ids,inputs_embeds=inputs_embeds,attention_mask=attention_mask,position_ids=position_ids,return_dict=True)

class iceReengineeredAudioForCausalLM(iceAudioForCausalLM):
    config_class = iceReengineeredAudioConfig
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
        cfg = kwargs.get('config')
        if cfg is None:
            from transformers import AutoConfig
            cfg = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True, local_files_only=kwargs.get('local_files_only', False))
        nested = getattr(cfg, 'language_model_config', None)
        if isinstance(nested, dict):
            cfg.language_model_config = iceAudioLanguageConfig(**nested)
        kwargs['config'] = cfg
        return super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
    def __init__(self, config, *args, **kwargs):
        super().__init__(config, *args, **kwargs)
        self.runtime_mode=getattr(config,'runtime_mode','independent_training')
        if self.runtime_mode != 'legacy_audio_compatible':
            self.language_model=IceRLanguageModel(config.language_model_config)
        self.architecture_version='ice-r/3.1.0'
    def _prepare_embed_inputs(self, input_ids, audio_mask):
        return super()._prepare_embed_inputs(input_ids, audio_mask)
    def supported_language_ids(self):
        return set(LANG_IDS)
    def supported_language_names(self):
        return set(LANG_NAMES)