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', 'dũ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)