BeatSeat
044240d8eb
Signed-off-by: BeatSeat <wendavid552@gmail.com> Signed-off-by: David <david@podcast.ai> Signed-off-by: wendavid552 <wendavid552@gmail.com> Signed-off-by: WenDavid <davidzhangcf@foxmail.com>
1063 lines
45 KiB
Python
1063 lines
45 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""IndexTTS2 Stage 0: GPT-2 AR Talker with vLLM-native PagedAttention.
|
|
|
|
Predicts mel codes autoregressively and collects hidden_states as latent
|
|
for Stage 1 (S2Mel + BigVGAN).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import math
|
|
import re
|
|
from collections.abc import Iterable
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from vllm.compilation.decorators import support_torch_compile
|
|
from vllm.config import VllmConfig
|
|
from vllm.distributed import get_pp_group
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
|
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
|
|
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
|
from vllm.model_executor.models.gpt2 import GPT2Block
|
|
from vllm.model_executor.models.utils import (
|
|
is_pp_missing_parameter,
|
|
make_empty_intermediate_tensors_factory,
|
|
make_layers,
|
|
maybe_prefix,
|
|
)
|
|
from vllm.sequence import IntermediateTensors
|
|
|
|
from vllm_omni.data_entry_keys import OmniPayload
|
|
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
|
from vllm_omni.utils.speaker_cache import get_speaker_cache
|
|
|
|
from .configuration_indextts2 import IndexTTS2Config
|
|
from .gpt.conformer_encoder import ConformerEncoder
|
|
from .gpt.embeddings import LearnedPositionEmbeddings
|
|
from .gpt.perceiver import PerceiverResampler
|
|
from .preprocess_utils import (
|
|
compute_fbank,
|
|
load_campplus,
|
|
load_qwen_emotion,
|
|
load_reference_audio,
|
|
load_semantic_codec,
|
|
load_wav2vec2,
|
|
resolve_model_file,
|
|
wav2vec_extract,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def _find_most_similar_cosine(query: torch.Tensor, matrix: torch.Tensor) -> int:
|
|
sims = F.cosine_similarity(query.float(), matrix.float(), dim=1)
|
|
return int(torch.argmax(sims).item())
|
|
|
|
|
|
@support_torch_compile(
|
|
dynamic_arg_dims={
|
|
"input_ids": 0,
|
|
"positions": 0,
|
|
"intermediate_tensors": 0,
|
|
"inputs_embeds": 0,
|
|
}
|
|
)
|
|
class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
|
"""vLLM-native GPT-2 AR talker for IndexTTS2.
|
|
|
|
Stage 0 of the two-stage pipeline. Predicts mel codes (8194 vocab)
|
|
and accumulates hidden_states as latent for Stage 1 S2Mel decoder.
|
|
"""
|
|
|
|
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
|
super().__init__()
|
|
self.vllm_config = vllm_config
|
|
self.model_path = vllm_config.model_config.model
|
|
self.config: IndexTTS2Config = vllm_config.model_config.hf_config # type: ignore[assignment]
|
|
gpt_cfg = self.config.gpt
|
|
|
|
self.model_dim = gpt_cfg["model_dim"]
|
|
self.num_layers = gpt_cfg["layers"]
|
|
self.num_heads = gpt_cfg["heads"]
|
|
self.max_mel_tokens = gpt_cfg["max_mel_tokens"]
|
|
self.max_text_tokens = gpt_cfg["max_text_tokens"]
|
|
self.number_mel_codes = gpt_cfg["number_mel_codes"]
|
|
self.start_mel_token = gpt_cfg["start_mel_token"]
|
|
self.stop_mel_token = gpt_cfg["stop_mel_token"]
|
|
self.number_text_tokens = gpt_cfg.get("number_text_tokens", 12000)
|
|
self.condition_num_latent = gpt_cfg.get("condition_num_latent", 32)
|
|
|
|
# --- Flags for vLLM-Omni framework ---
|
|
self.have_multimodal_outputs = True
|
|
self.has_preprocess = True
|
|
self.has_postprocess = True
|
|
self.requires_raw_input_tokens = True
|
|
self.enable_update_additional_information = True
|
|
# S2Mel consumes the complete code sequence only at request end, but
|
|
# the connector accumulator needs to observe each appendable AR delta.
|
|
# Keep runner request-end gating disabled and emit per-step rows from
|
|
# make_omni_output; the full-payload builder reconstructs the sequence.
|
|
self.omni_payload_at_request_end = False
|
|
self.omni_request_end_token_ids = ()
|
|
self._speaker_cache = get_speaker_cache()
|
|
self.gpu_resident_buffer_keys: set[tuple[str, str]] = {
|
|
("codes", "mel"),
|
|
("hidden_states", "latent"),
|
|
("meta", "mel_start_offset"),
|
|
("meta", "latent_acc"),
|
|
("meta", "mel_code_count"),
|
|
}
|
|
|
|
# --- GPT-2 transformer (vLLM-native with PagedAttention) ---
|
|
# Build a GPT2Config-like object for GPT2Block
|
|
from transformers import GPT2Config as HFGpt2Config
|
|
|
|
max_seq_len = self.max_mel_tokens + self.max_text_tokens + self.condition_num_latent + 4
|
|
gpt2_config = HFGpt2Config(
|
|
vocab_size=self.number_mel_codes,
|
|
n_positions=max_seq_len,
|
|
n_ctx=max_seq_len,
|
|
n_embd=self.model_dim,
|
|
n_layer=self.num_layers,
|
|
n_head=self.num_heads,
|
|
activation_function="gelu_new",
|
|
layer_norm_epsilon=1e-5,
|
|
)
|
|
# Store for GPT2Block construction
|
|
cache_config = vllm_config.cache_config
|
|
quant_config = vllm_config.quant_config
|
|
self.start_layer, self.end_layer, self.h = make_layers(
|
|
gpt2_config.n_layer,
|
|
lambda prefix: GPT2Block(gpt2_config, cache_config, quant_config, prefix=prefix),
|
|
prefix=f"{prefix}.h",
|
|
)
|
|
self.ln_f = nn.LayerNorm(self.model_dim, eps=gpt2_config.layer_norm_epsilon)
|
|
self.final_norm = nn.LayerNorm(self.model_dim, eps=gpt2_config.layer_norm_epsilon)
|
|
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
|
|
["hidden_states"], self.model_dim
|
|
)
|
|
|
|
# --- Embeddings ---
|
|
self.text_embedding = nn.Embedding(self.number_text_tokens + 1, self.model_dim)
|
|
self.mel_embedding = nn.Embedding(self.number_mel_codes, self.model_dim)
|
|
self.text_pos_embedding = LearnedPositionEmbeddings(self.max_text_tokens + 2, self.model_dim)
|
|
self.mel_pos_embedding = LearnedPositionEmbeddings(self.max_mel_tokens + 2 + 1, self.model_dim)
|
|
self.speed_emb = nn.Embedding(2, self.model_dim)
|
|
self.emo_layer = nn.Linear(self.model_dim, self.model_dim)
|
|
self.emovec_layer = nn.Linear(1024, self.model_dim)
|
|
|
|
# --- Mel head (logits) ---
|
|
if get_pp_group().is_last_rank:
|
|
self.mel_head = ParallelLMHead(
|
|
self.number_mel_codes,
|
|
self.model_dim,
|
|
quant_config=quant_config,
|
|
prefix=maybe_prefix(prefix, "mel_head"),
|
|
)
|
|
else:
|
|
self.mel_head = None # type: ignore[assignment]
|
|
self.logits_processor = LogitsProcessor(self.number_mel_codes)
|
|
|
|
# --- Conditioning modules (speaker / emotion) ---
|
|
cond_cfg = gpt_cfg.get("condition_module", {})
|
|
emo_cond_cfg = gpt_cfg.get("emo_condition_module", {})
|
|
|
|
self.conditioning_encoder = ConformerEncoder(
|
|
input_size=1024,
|
|
output_size=cond_cfg.get("output_size", 512),
|
|
linear_units=cond_cfg.get("linear_units", 2048),
|
|
attention_heads=cond_cfg.get("attention_heads", 8),
|
|
num_blocks=cond_cfg.get("num_blocks", 6),
|
|
input_layer=cond_cfg.get("input_layer", "conv2d2"),
|
|
)
|
|
self.perceiver_encoder = PerceiverResampler(
|
|
self.model_dim,
|
|
dim_context=cond_cfg.get("output_size", 512),
|
|
ff_mult=cond_cfg.get("perceiver_mult", 2),
|
|
heads=cond_cfg.get("attention_heads", 8),
|
|
num_latents=self.condition_num_latent,
|
|
)
|
|
|
|
self.emo_conditioning_encoder = ConformerEncoder(
|
|
input_size=1024,
|
|
output_size=emo_cond_cfg.get("output_size", 512),
|
|
linear_units=emo_cond_cfg.get("linear_units", 1024),
|
|
attention_heads=emo_cond_cfg.get("attention_heads", 4),
|
|
num_blocks=emo_cond_cfg.get("num_blocks", 4),
|
|
input_layer=emo_cond_cfg.get("input_layer", "conv2d2"),
|
|
)
|
|
self.emo_perceiver_encoder = PerceiverResampler(
|
|
1024,
|
|
dim_context=emo_cond_cfg.get("output_size", 512),
|
|
ff_mult=emo_cond_cfg.get("perceiver_mult", 2),
|
|
heads=emo_cond_cfg.get("attention_heads", 4),
|
|
num_latents=1,
|
|
)
|
|
|
|
# Padding masks for perceiver cross-attention
|
|
self.cond_mask_pad = nn.ConstantPad1d((self.condition_num_latent, 0), True)
|
|
self.emo_cond_mask_pad = nn.ConstantPad1d((1, 0), True)
|
|
|
|
# --- Lazy-loaded external models ---
|
|
self._w2v_stat: torch.Tensor | None = None
|
|
self._emo_matrix: torch.Tensor | None = None
|
|
self._spk_matrix: torch.Tensor | None = None
|
|
self._text_tokenizer: Any = None
|
|
|
|
# Initialize embeddings per GPT-2 convention
|
|
for emb in [self.text_embedding, self.mel_embedding]:
|
|
emb.weight.data.normal_(mean=0.0, std=0.02)
|
|
self.speed_emb.weight.data.normal_(mean=0.0, std=0.0)
|
|
|
|
# ------------------------------------------------------------------
|
|
# vLLM required hooks
|
|
# ------------------------------------------------------------------
|
|
|
|
def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
|
|
return self.mel_embedding(input_ids)
|
|
|
|
def _extract_mel_offsets(self, positions: torch.Tensor, kwargs: dict[str, Any]) -> torch.Tensor:
|
|
"""Build per-token mel_start_offset from model_intermediate_buffer."""
|
|
info_dicts = kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or []
|
|
if "runtime_additional_information" in kwargs and "model_intermediate_buffer" not in kwargs:
|
|
logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer")
|
|
if not info_dicts:
|
|
logger.warning_once("[OFFSET] no model_intermediate_buffer available; defaulting mel offsets to 0")
|
|
return torch.zeros_like(positions)
|
|
offsets = []
|
|
for info in info_dicts:
|
|
if isinstance(info, dict):
|
|
meta = info.get("meta", {})
|
|
val = meta.get("mel_start_offset", 0)
|
|
offsets.append(int(val) if not isinstance(val, int) else val)
|
|
else:
|
|
offsets.append(0)
|
|
return torch.tensor(offsets, dtype=positions.dtype, device=positions.device)
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
intermediate_tensors: IntermediateTensors | None = None,
|
|
inputs_embeds: torch.Tensor | None = None,
|
|
**kwargs: Any,
|
|
) -> torch.Tensor | IntermediateTensors:
|
|
"""AR transformer forward.
|
|
|
|
During prefill: inputs_embeds is set by preprocess(), input_ids ignored.
|
|
During decode: inputs_embeds is mel_embedding(input_ids) + mel_pos.
|
|
"""
|
|
if get_pp_group().is_first_rank:
|
|
if inputs_embeds is None:
|
|
logger.warning("[FORWARD] inputs_embeds=None — fallback decode path")
|
|
mel_start_offset = self._extract_mel_offsets(positions, kwargs)
|
|
mel_pos = positions - mel_start_offset
|
|
mel_pos = torch.clamp(mel_pos, min=0)
|
|
inputs_embeds = self.mel_embedding(input_ids) + self.mel_pos_embedding.emb(mel_pos)
|
|
hidden_states = inputs_embeds
|
|
else:
|
|
assert intermediate_tensors is not None
|
|
hidden_states = intermediate_tensors["hidden_states"]
|
|
|
|
for layer in self.h[self.start_layer : self.end_layer]:
|
|
hidden_states = layer(hidden_states)
|
|
|
|
if not get_pp_group().is_last_rank:
|
|
return IntermediateTensors({"hidden_states": hidden_states})
|
|
|
|
return self.final_norm(self.ln_f(hidden_states))
|
|
|
|
def compute_logits(
|
|
self,
|
|
hidden_states: torch.Tensor | OmniOutput,
|
|
sampling_metadata: Any = None,
|
|
) -> torch.Tensor | None:
|
|
if isinstance(hidden_states, OmniOutput):
|
|
hidden_states = hidden_states.text_hidden_states
|
|
if hidden_states is None or self.mel_head is None:
|
|
return None
|
|
return self.logits_processor(self.mel_head, hidden_states)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Omni multimodal output plumbing
|
|
# ------------------------------------------------------------------
|
|
|
|
def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
|
|
"""Collect mel_codes and hidden_states from intermediate buffer."""
|
|
if isinstance(model_outputs, OmniOutput):
|
|
return model_outputs
|
|
|
|
hidden = model_outputs
|
|
info_dicts = kwargs.get("model_intermediate_buffer")
|
|
if info_dicts is None:
|
|
info_dicts = kwargs.get("runtime_additional_information") or []
|
|
if "runtime_additional_information" in kwargs and "model_intermediate_buffer" not in kwargs:
|
|
logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer")
|
|
|
|
device = hidden.device
|
|
latent_dim = hidden.shape[-1]
|
|
_zero_codes = torch.zeros(0, 1, dtype=torch.long, device=device)
|
|
_zero_latent = torch.zeros(0, latent_dim, dtype=hidden.dtype, device=device)
|
|
|
|
# Per-request lists: to_payload_element dispatches lists via
|
|
# element[idx], giving correct per-request mapping even during
|
|
# mixed prefill+decode steps. Zero-length placeholders for prefill
|
|
# requests are harmless — torch.cat ignores them in accumulation.
|
|
#
|
|
# Export semantics are appendable deltas for the existing full-payload
|
|
# accumulator: codes.mel is a 2-D row [1, 1], hidden_states.latent is
|
|
# the aligned one-row [1, D] slice from meta.latent_acc.
|
|
mel_codes_list: list[torch.Tensor] = []
|
|
latent_list: list[torch.Tensor] = []
|
|
meta_list: list[dict[str, Any]] = []
|
|
|
|
for info in info_dicts:
|
|
if not isinstance(info, dict):
|
|
mel_codes_list.append(_zero_codes)
|
|
latent_list.append(_zero_latent)
|
|
meta_list.append({})
|
|
continue
|
|
codes = info.get("codes", {})
|
|
mc = codes.get("mel")
|
|
info_meta = info.get("meta", {})
|
|
# mel_len now comes from the int counter in meta, not from
|
|
# code_seq.shape[0] (codes.mel is a single-token delta).
|
|
mel_len = int(info_meta.get("mel_code_count", 0))
|
|
if isinstance(mc, torch.Tensor) and mc.numel() > 0 and mel_len > 0:
|
|
code_seq = mc.reshape(-1)
|
|
mel_codes_list.append(code_seq[-1:].reshape(1, 1).contiguous())
|
|
else:
|
|
mel_codes_list.append(_zero_codes)
|
|
|
|
lat_acc = info_meta.get("latent_acc")
|
|
if isinstance(lat_acc, torch.Tensor) and lat_acc.numel() > 0 and mel_len > 0:
|
|
lat_seq = lat_acc.reshape(-1, lat_acc.shape[-1]) if lat_acc.ndim != 2 else lat_acc
|
|
idx = min(max(mel_len - 1, 0), int(lat_seq.shape[0]) - 1)
|
|
latent_list.append(lat_seq[idx : idx + 1].contiguous())
|
|
else:
|
|
hs = info.get("hidden_states", {})
|
|
lat = hs.get("latent")
|
|
if isinstance(lat, torch.Tensor) and lat.numel() > 0:
|
|
lat_seq = lat.reshape(-1, lat.shape[-1]) if lat.ndim != 2 else lat
|
|
latent_list.append(lat_seq[-1:].contiguous())
|
|
else:
|
|
latent_list.append(_zero_latent)
|
|
|
|
req_meta: dict[str, Any] = {}
|
|
if mel_len == 1:
|
|
s_ref = info_meta.get("S_ref")
|
|
ref_mel = info_meta.get("ref_mel")
|
|
style = info_meta.get("style")
|
|
if isinstance(s_ref, torch.Tensor):
|
|
req_meta["S_ref"] = s_ref
|
|
if isinstance(ref_mel, torch.Tensor):
|
|
req_meta["ref_mel"] = ref_mel
|
|
if isinstance(style, torch.Tensor):
|
|
req_meta["style"] = style
|
|
meta_list.append(req_meta)
|
|
|
|
if not any(t.numel() > 0 for t in mel_codes_list):
|
|
return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
|
|
|
|
mm: OmniPayload = {
|
|
"codes": {"mel": mel_codes_list},
|
|
"hidden_states": {"latent": latent_list},
|
|
}
|
|
if any(meta_list):
|
|
mm["meta"] = meta_list
|
|
|
|
return OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm)
|
|
|
|
# ------------------------------------------------------------------
|
|
# preprocess / postprocess
|
|
# ------------------------------------------------------------------
|
|
|
|
def preprocess(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
input_embeds: torch.Tensor | None,
|
|
**info_dict: Any,
|
|
) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
|
|
"""Build prompt embeddings for prefill; compute mel embeddings for decode.
|
|
|
|
Prefill layout::
|
|
|
|
[conds(32) + emo_vec(1) + duration(2)] [text_emb + text_pos] [start_mel + mel_pos(0)]
|
|
|
|
Decode: ``mel_embedding(token) + mel_pos_embedding(step)``.
|
|
"""
|
|
span_len = int(input_ids.shape[0])
|
|
meta = info_dict.get("meta", {})
|
|
|
|
# --- Decode path: span_len==1 and mel_start_offset already set ---
|
|
if span_len == 1 and isinstance(meta, dict) and meta.get("mel_start_offset") is not None:
|
|
return self._preprocess_decode(input_ids, info_dict, meta)
|
|
|
|
# --- Prefill path ---
|
|
return self._preprocess_prefill(input_ids, input_embeds, info_dict, span_len)
|
|
|
|
def _preprocess_decode(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
info_dict: dict[str, Any],
|
|
meta: dict[str, Any],
|
|
) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
|
|
"""Lightweight decode step: embed token + track code count."""
|
|
# Track code count via int counter instead of accumulating a full tensor
|
|
# (the connector reconstructs the full sequence from per-step deltas).
|
|
prev_count = meta.get("mel_code_count", 0)
|
|
if not isinstance(prev_count, int):
|
|
prev_count = int(prev_count)
|
|
decode_step = prev_count + 1
|
|
|
|
# Only keep the current token; no O(T) history accumulation.
|
|
new_codes = input_ids.detach()
|
|
|
|
mel_pos = input_ids.new_full((1,), decode_step)
|
|
embeds = self.mel_embedding(input_ids) + self.mel_pos_embedding.emb(mel_pos)
|
|
|
|
return (
|
|
input_ids,
|
|
embeds,
|
|
{
|
|
"codes": {"mel": new_codes},
|
|
"meta": {"mel_code_count": decode_step},
|
|
},
|
|
)
|
|
|
|
def _preprocess_prefill(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
input_embeds: torch.Tensor | None,
|
|
info_dict: dict[str, Any],
|
|
span_len: int,
|
|
) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
|
|
"""Full conditioning pipeline — runs once during prefill."""
|
|
additional_information = info_dict.get("additional_information")
|
|
if isinstance(additional_information, dict):
|
|
merged: dict[str, Any] = {k: v for k, v in info_dict.items() if k != "additional_information"}
|
|
for k, v in additional_information.items():
|
|
merged.setdefault(k, v)
|
|
info_dict = merged
|
|
|
|
device = input_ids.device
|
|
|
|
def _first(key, default=None):
|
|
v = info_dict.get(key)
|
|
return v[0] if isinstance(v, list) and v else default
|
|
|
|
text = _first("text")
|
|
if not text:
|
|
raise ValueError("Missing additional_information.text for IndexTTS2 talker.")
|
|
|
|
voice_path = _first("voice")
|
|
if voice_path is None:
|
|
raise ValueError("IndexTTS2 requires voice (reference audio path) in additional_information.")
|
|
|
|
emo_audio_path = _first("emo_audio")
|
|
emo_text = _first("emo_text")
|
|
use_emo_text = bool(_first("use_emo_text", False))
|
|
emo_vector = _first("emo_vector")
|
|
emo_alpha = float(_first("emo_alpha", 1.0))
|
|
use_random = bool(_first("use_random", False))
|
|
_raw_emo_voice = _first("emo_voice_name")
|
|
emo_voice_name = str(_raw_emo_voice).strip().lower() if _raw_emo_voice else None
|
|
|
|
# --- Speaker cache lookup ---
|
|
_voice_name = _first("voice_name")
|
|
_voice_created_at = int(_first("voice_created_at", 0))
|
|
_speaker_cache_key = None
|
|
_cached = None
|
|
if _voice_name:
|
|
_speaker_cache_key = self._speaker_cache.make_cache_key(_voice_name, "indextts2", _voice_created_at)
|
|
_cached = self._speaker_cache.get(_speaker_cache_key)
|
|
|
|
# --- Load audio and extract features ---
|
|
# Keep the speaker-cache hit path truly lazy: cached speaker artifacts are
|
|
# sufficient for conditioning and same-speaker emotion reuse, so avoid
|
|
# loading/resampling reference audio and touching the Wav2Vec2 singleton
|
|
# unless a miss (or an uncached separate emotion audio) actually needs it.
|
|
wav_16k: torch.Tensor | None = None
|
|
w2v_model = None
|
|
w2v_proc = None
|
|
|
|
if _cached is not None:
|
|
s_ref = _cached["S_ref"].to(device)
|
|
style = _cached["style"].to(device)
|
|
ref_mel = _cached["ref_mel"].to(device)
|
|
spk_cond_emb = _cached["spk_cond_emb"].to(device)
|
|
else:
|
|
wav_16k, wav_22k = load_reference_audio(voice_path, device, mode="speaker")
|
|
w2v_model, w2v_proc = load_wav2vec2(self.model_path, device)
|
|
self._ensure_w2v_stat_loaded(device)
|
|
spk_cond_emb = wav2vec_extract(wav_16k, w2v_model, w2v_proc, device, self._w2v_stat)
|
|
|
|
semantic_codec = load_semantic_codec(self.model_path, self.config.semantic_codec, device)
|
|
with torch.no_grad():
|
|
_, s_ref = semantic_codec.quantize(spk_cond_emb) # [B, T, 1024] quantized embeddings
|
|
|
|
campplus = load_campplus(self.model_path, device)
|
|
fbank = compute_fbank(wav_16k, device)
|
|
with torch.no_grad():
|
|
style = campplus(fbank) # [1, 192]
|
|
|
|
ref_mel = self._compute_mel_22k(wav_22k, device) # [1, 80, T_ref]
|
|
|
|
if _speaker_cache_key is not None:
|
|
self._speaker_cache.put(
|
|
_speaker_cache_key,
|
|
{
|
|
"S_ref": s_ref.detach().cpu().contiguous(),
|
|
"style": style.detach().cpu().contiguous(),
|
|
"ref_mel": ref_mel.detach().cpu().contiguous(),
|
|
"spk_cond_emb": spk_cond_emb.detach().cpu().contiguous(),
|
|
},
|
|
)
|
|
|
|
model_dtype = next(self.conditioning_encoder.parameters()).dtype
|
|
spk_cond_emb = spk_cond_emb.to(device=device, dtype=model_dtype)
|
|
spk_lens = torch.tensor([spk_cond_emb.shape[1]], device=device, dtype=torch.long)
|
|
speech_cond, mask = self.conditioning_encoder(spk_cond_emb, spk_lens)
|
|
conds_mask = self.cond_mask_pad(mask.squeeze(1))
|
|
conds = self.perceiver_encoder(speech_cond, conds_mask) # [1, 32, D]
|
|
|
|
emo_vec = self._compute_emotion_vector(
|
|
wav_16k=wav_16k,
|
|
speaker_audio_path=voice_path,
|
|
emo_audio_path=emo_audio_path,
|
|
emo_voice_name=emo_voice_name,
|
|
main_text=text,
|
|
use_emo_text=use_emo_text,
|
|
emo_text=emo_text,
|
|
emo_vector=emo_vector,
|
|
emo_alpha=emo_alpha,
|
|
use_random=use_random,
|
|
style=style,
|
|
spk_cond_emb=spk_cond_emb,
|
|
w2v_model=w2v_model,
|
|
w2v_proc=w2v_proc,
|
|
device=device,
|
|
)
|
|
|
|
speed_zero = torch.zeros(1, device=device, dtype=torch.long)
|
|
speed_one = torch.ones(1, device=device, dtype=torch.long)
|
|
duration_emb = self.speed_emb(speed_zero)
|
|
duration_emb_half = self.speed_emb(speed_one)
|
|
|
|
conds_with_emo = conds + emo_vec.unsqueeze(1) # [1, 32, D]
|
|
conds_prefix = torch.cat(
|
|
[conds_with_emo, duration_emb_half.unsqueeze(1), duration_emb.unsqueeze(1)],
|
|
dim=1,
|
|
) # [1, 34, D]
|
|
|
|
text_tokens = self._tokenize_text(text, device) # [1, L]
|
|
text_emb = self.text_embedding(text_tokens) + self.text_pos_embedding(text_tokens)
|
|
|
|
start_mel = self.mel_embedding(torch.tensor([[self.start_mel_token]], device=device, dtype=torch.long))
|
|
start_mel_pos = self.mel_pos_embedding.get_fixed_embedding(0, device)
|
|
start_mel = start_mel + start_mel_pos # [1, 1, D]
|
|
|
|
inputs_embeds = torch.cat([conds_prefix, text_emb, start_mel], dim=1).squeeze(0)
|
|
mel_start_offset = int(conds_prefix.shape[1]) + int(text_emb.shape[1])
|
|
prefill_mel_index = int(inputs_embeds.shape[0]) - 1
|
|
|
|
if int(inputs_embeds.shape[0]) < span_len:
|
|
pad_n = span_len - int(inputs_embeds.shape[0])
|
|
pad_emb = torch.zeros(pad_n, self.model_dim, device=device, dtype=inputs_embeds.dtype)
|
|
inputs_embeds = torch.cat([inputs_embeds, pad_emb], dim=0)
|
|
elif int(inputs_embeds.shape[0]) > span_len:
|
|
inputs_embeds = inputs_embeds[:span_len]
|
|
|
|
input_ids_out = torch.full_like(input_ids, self.start_mel_token)
|
|
|
|
info_update: dict[str, Any] = {
|
|
"meta": {
|
|
"mel_start_offset": mel_start_offset,
|
|
"prefill_mel_index": prefill_mel_index,
|
|
"mel_code_count": 0,
|
|
"S_ref": s_ref.cpu().contiguous(),
|
|
"ref_mel": ref_mel.cpu().contiguous(),
|
|
"style": style.cpu().contiguous(),
|
|
"latent_acc": torch.zeros(0, self.model_dim, dtype=inputs_embeds.dtype, device=device),
|
|
},
|
|
"codes": {"mel": torch.zeros(0, dtype=torch.long, device=device)},
|
|
"hidden_states": {"latent": torch.zeros(0, self.model_dim, dtype=inputs_embeds.dtype, device=device)},
|
|
}
|
|
|
|
return input_ids_out, inputs_embeds, info_update
|
|
|
|
def postprocess(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
multimodal_outputs: Any = None,
|
|
**kwargs: Any,
|
|
) -> dict[str, Any]:
|
|
"""Store current decode-step hidden state for Stage 1.
|
|
|
|
Only the current row is kept; the connector's full-payload accumulator
|
|
reconstructs the complete latent sequence from per-step deltas emitted
|
|
by make_omni_output.
|
|
"""
|
|
if hidden_states.shape[0] != 1:
|
|
meta = kwargs.get("meta", {})
|
|
prefill_mel_index = int(meta.get("prefill_mel_index", hidden_states.shape[0] - 1))
|
|
prefill_mel_index = min(max(prefill_mel_index, 0), hidden_states.shape[0] - 1)
|
|
return {"meta": {"latent_acc": hidden_states[prefill_mel_index : prefill_mel_index + 1].detach()}}
|
|
|
|
return {"meta": {"latent_acc": hidden_states.detach()}}
|
|
|
|
def _compute_mel_22k(self, wav_22k: torch.Tensor, device: torch.device) -> torch.Tensor:
|
|
"""Compute 80-band mel spectrogram at 22.05kHz for Stage 1."""
|
|
s2mel_cfg = self.config.s2mel.get("preprocess_params", {})
|
|
sr = s2mel_cfg.get("sr", 22050)
|
|
spect = s2mel_cfg.get("spect_params", {})
|
|
n_fft = spect.get("n_fft", 1024)
|
|
hop_length = spect.get("hop_length", 256)
|
|
win_length = spect.get("win_length", 1024)
|
|
n_mels = spect.get("n_mels", 80)
|
|
|
|
# mel_spectrogram is pure torch with per-device caches — run on the
|
|
# wav's device (GPU after load_reference_audio) instead of CPU.
|
|
wav = wav_22k.float()
|
|
if wav.ndim == 1:
|
|
wav = wav.unsqueeze(0)
|
|
|
|
from .s2mel.modules.audio import mel_spectrogram as mel_fn
|
|
|
|
mel = mel_fn(wav, n_fft, n_mels, sr, hop_length, win_length, 0, None)
|
|
return mel.to(device=device) # [1, 80, T]
|
|
|
|
# ------------------------------------------------------------------
|
|
# Emotion conditioning
|
|
# ------------------------------------------------------------------
|
|
|
|
def _compute_emotion_vector(
|
|
self,
|
|
*,
|
|
wav_16k: torch.Tensor | None,
|
|
speaker_audio_path: Any,
|
|
emo_audio_path: str | None,
|
|
emo_voice_name: str | None,
|
|
main_text: str,
|
|
use_emo_text: bool,
|
|
emo_text: str | None,
|
|
emo_vector: Any,
|
|
emo_alpha: float,
|
|
use_random: bool,
|
|
style: torch.Tensor,
|
|
spk_cond_emb: torch.Tensor,
|
|
w2v_model: Any,
|
|
w2v_proc: Any,
|
|
device: torch.device,
|
|
) -> torch.Tensor:
|
|
"""Full emotion control logic aligned with official infer_v2.py. Returns [1, D]."""
|
|
# --- Mutual exclusion: text/vector guidance clears emotion reference audio ---
|
|
# Mirrors official infer_v2.py: use_emo_text is evaluated before explicit
|
|
# vectors, so text guidance wins if both are supplied by an internal caller.
|
|
if use_emo_text:
|
|
if emo_text is None:
|
|
emo_text = main_text
|
|
emo_vector = self._predict_emotion_from_text(emo_text, device)
|
|
if use_emo_text or emo_vector is not None:
|
|
emo_audio_path = None
|
|
|
|
# --- Validate emo_vector (user-provided or text-predicted) ---
|
|
if emo_vector is not None:
|
|
if isinstance(emo_vector, torch.Tensor):
|
|
emo_vector = emo_vector.tolist()
|
|
elif isinstance(emo_vector, np.ndarray):
|
|
emo_vector = emo_vector.tolist()
|
|
if not isinstance(emo_vector, list) or len(emo_vector) != len(self._DESIRED_ORDER):
|
|
raise ValueError(
|
|
"IndexTTS2 emo_vector must contain 8 values in order: "
|
|
"happy, angry, sad, afraid, disgusted, melancholic, surprised, calm"
|
|
)
|
|
if any(isinstance(v, bool) for v in emo_vector):
|
|
raise ValueError("IndexTTS2 emo_vector values must be numeric")
|
|
try:
|
|
emo_vector = [float(v) for v in emo_vector]
|
|
except (TypeError, ValueError) as exc:
|
|
raise ValueError("IndexTTS2 emo_vector values must be numeric") from exc
|
|
if not all(math.isfinite(v) for v in emo_vector):
|
|
raise ValueError("IndexTTS2 emo_vector values must be finite")
|
|
if not all(0.0 <= v <= 1.2 for v in emo_vector):
|
|
raise ValueError("IndexTTS2 emo_vector values must be between 0.0 and 1.2")
|
|
|
|
# Alpha pre-scaling (official infer_v2.py:414-418):
|
|
# when emo_alpha != 1.0, scale emo_vector down before overlay.
|
|
emo_vector_scale = max(0.0, min(1.0, emo_alpha))
|
|
if emo_vector_scale != 1.0:
|
|
emo_vector = [int(x * emo_vector_scale * 10000) / 10000 for x in emo_vector]
|
|
|
|
# --- Path B: emo_vector (8-dim) → spk_matrix cosine lookup ---
|
|
emovec_mat = None
|
|
weight_sum = 0.0
|
|
if emo_vector is not None:
|
|
weight_sum = sum(emo_vector)
|
|
vec = torch.tensor(emo_vector, device=device, dtype=torch.float32)
|
|
emovec_mat = self._compute_emo_from_distribution(vec, style, device, use_random=use_random)
|
|
|
|
# --- Audio-based emotion with alpha blending ---
|
|
if emo_audio_path is None:
|
|
# No separate emotion audio: reuse speaker w2v-bert features instead of
|
|
# reloading + re-extracting the same audio (saves ~0.5s per prefill).
|
|
# Official code resamples orig→16k directly for emotion vs orig→22k→16k
|
|
# for speaker; the resulting w2v features are near-identical.
|
|
effective_alpha = 1.0
|
|
emo_cond_emb = spk_cond_emb
|
|
else:
|
|
effective_alpha = emo_alpha
|
|
# emotion audio cache (enabled when emo_voice_name is provided)
|
|
_emo_cache_key = None
|
|
_emo_cached = None
|
|
if emo_voice_name:
|
|
_emo_cache_key = self._speaker_cache.make_cache_key(emo_voice_name, "indextts2_emo", 0)
|
|
_emo_cached = self._speaker_cache.get(_emo_cache_key)
|
|
|
|
if _emo_cached is not None:
|
|
emo_cond_emb = _emo_cached["spk_cond_emb"].to(device)
|
|
logger.debug("[EMO] cache HIT for %s", emo_voice_name)
|
|
else:
|
|
emo_wav_16k, _ = load_reference_audio(emo_audio_path, device, mode="emotion")
|
|
if w2v_model is None or w2v_proc is None:
|
|
w2v_model, w2v_proc = load_wav2vec2(self.model_path, device)
|
|
self._ensure_w2v_stat_loaded(device)
|
|
emo_cond_emb = wav2vec_extract(emo_wav_16k, w2v_model, w2v_proc, device, self._w2v_stat)
|
|
if _emo_cache_key is not None:
|
|
self._speaker_cache.put(
|
|
_emo_cache_key,
|
|
{
|
|
"spk_cond_emb": emo_cond_emb.detach().cpu().contiguous(),
|
|
},
|
|
)
|
|
logger.debug("[EMO] cache MISS, cached for %s", emo_voice_name)
|
|
|
|
emovec = self._merge_emovec(spk_cond_emb, emo_cond_emb, effective_alpha)
|
|
|
|
# --- emo_vector overlay (official infer_v2.py:560-561) ---
|
|
# emovec_mat is raw feat2.pt weighted sum [1, model_dim], NOT projected.
|
|
# Mix: emovec_mat contributes the emotion prototype; (1-sum)*emovec
|
|
# preserves the audio-based emotion as a "base color".
|
|
if emovec_mat is not None:
|
|
emovec = emovec_mat.to(dtype=emovec.dtype) + (1.0 - weight_sum) * emovec
|
|
|
|
return emovec
|
|
|
|
_MELANCHOLIC_WORDS = frozenset(
|
|
{
|
|
"低落",
|
|
"melancholy",
|
|
"melancholic",
|
|
"depression",
|
|
"depressed",
|
|
"gloomy",
|
|
}
|
|
)
|
|
|
|
_DESIRED_ORDER = ["高兴", "愤怒", "悲伤", "恐惧", "反感", "低落", "惊讶", "自然"]
|
|
|
|
def _predict_emotion_from_text(self, text: str, device: torch.device) -> list[float] | None:
|
|
"""Use QwenEmotion CausalLM to predict 8-dim emotion vector from text (aligned with official)."""
|
|
model, tokenizer = load_qwen_emotion(
|
|
self.model_path,
|
|
device,
|
|
trust_remote_code=self.vllm_config.model_config.trust_remote_code,
|
|
)
|
|
if model is None or tokenizer is None:
|
|
return None
|
|
|
|
messages = [
|
|
{"role": "system", "content": "文本情感分类"},
|
|
{"role": "user", "content": text},
|
|
]
|
|
prompt_text = tokenizer.apply_chat_template(
|
|
messages,
|
|
tokenize=False,
|
|
add_generation_prompt=True,
|
|
enable_thinking=False,
|
|
)
|
|
model_inputs = tokenizer([prompt_text], return_tensors="pt").to(model.device)
|
|
|
|
with torch.no_grad():
|
|
generated_ids = model.generate(
|
|
**model_inputs,
|
|
max_new_tokens=512,
|
|
pad_token_id=tokenizer.eos_token_id,
|
|
)
|
|
output_ids = generated_ids[0][len(model_inputs.input_ids[0]) :].tolist()
|
|
|
|
try:
|
|
index = len(output_ids) - output_ids[::-1].index(151668)
|
|
except ValueError:
|
|
index = 0
|
|
content_str = tokenizer.decode(output_ids[index:], skip_special_tokens=True)
|
|
|
|
try:
|
|
content = json.loads(content_str)
|
|
except json.JSONDecodeError:
|
|
content = {
|
|
m.group(1): float(m.group(2)) for m in re.finditer(r'([^\s":.,]+?)"?\s*:\s*([\d.]+)', content_str)
|
|
}
|
|
|
|
text_lower = text.lower()
|
|
if any(w in text_lower for w in self._MELANCHOLIC_WORDS):
|
|
content["悲伤"], content["低落"] = content.get("低落", 0.0), content.get("悲伤", 0.0)
|
|
|
|
max_score, min_score = 1.2, 0.0
|
|
emo_vector = []
|
|
for cn_key in self._DESIRED_ORDER:
|
|
val = content.get(cn_key, 0.0)
|
|
if isinstance(val, str):
|
|
try:
|
|
val = float(val)
|
|
except ValueError:
|
|
val = 0.0
|
|
emo_vector.append(max(min_score, min(max_score, val)))
|
|
|
|
if all(v <= 0.0 for v in emo_vector):
|
|
emo_vector[7] = 1.0 # default calm
|
|
|
|
return emo_vector
|
|
|
|
def _compute_emo_from_distribution(
|
|
self,
|
|
emo_dist: torch.Tensor,
|
|
style: torch.Tensor,
|
|
device: torch.device,
|
|
use_random: bool = False,
|
|
) -> torch.Tensor:
|
|
"""From 8-dim emotion distribution + CAMPPlus style, build emovec_mat via spk_matrix cosine lookup.
|
|
|
|
Returns [1, model_dim] (feat2.pt rows are already model_dim).
|
|
"""
|
|
self._ensure_matrix_loaded("_emo_matrix", "emo_matrix", "Emotion matrix")
|
|
self._ensure_matrix_loaded("_spk_matrix", "spk_matrix", "Speaker matrix")
|
|
if self._emo_matrix is None or self._spk_matrix is None:
|
|
return torch.zeros(1, self.model_dim, device=device)
|
|
|
|
weight_vector = emo_dist.to(device=device, dtype=torch.float32)
|
|
if weight_vector.ndim == 2:
|
|
weight_vector = weight_vector.squeeze(0) # [8]
|
|
|
|
if use_random:
|
|
indices = [int(torch.randint(0, sub.shape[0], (1,), device=device).item()) for sub in self._spk_matrix]
|
|
else:
|
|
indices = [_find_most_similar_cosine(style, sub.to(device)) for sub in self._spk_matrix]
|
|
emo_rows = torch.cat(
|
|
[sub.to(device)[idx].unsqueeze(0) for idx, sub in zip(indices, self._emo_matrix)], dim=0
|
|
) # [8, model_dim]
|
|
|
|
return (weight_vector.unsqueeze(1) * emo_rows).sum(0).unsqueeze(0) # [1, model_dim]
|
|
|
|
def _compute_audio_emo_vec(self, cond_emb: torch.Tensor) -> torch.Tensor:
|
|
"""Conformer → Perceiver → emovec_layer → emo_layer. Input: [1, T, 1024], output: [1, D]."""
|
|
emo_dtype = next(self.emo_conditioning_encoder.parameters()).dtype
|
|
cond_emb = cond_emb.to(device=cond_emb.device, dtype=emo_dtype)
|
|
cond_lens = torch.tensor([cond_emb.shape[1]], device=cond_emb.device, dtype=torch.long)
|
|
emo_cond, emo_mask = self.emo_conditioning_encoder(cond_emb, cond_lens)
|
|
emo_conds_mask = self.emo_cond_mask_pad(emo_mask.squeeze(1))
|
|
emo_percept = self.emo_perceiver_encoder(emo_cond, emo_conds_mask) # [1, 1, 1024]
|
|
emo_vec_syn = self.emovec_layer(emo_percept.squeeze(1)) # [1, D]
|
|
return self.emo_layer(emo_vec_syn) # [1, D]
|
|
|
|
def _merge_emovec(
|
|
self,
|
|
spk_cond_emb: torch.Tensor,
|
|
emo_cond_emb: torch.Tensor,
|
|
alpha: float,
|
|
) -> torch.Tensor:
|
|
"""Alpha-blend speaker and emotion audio vectors: base + alpha * (emo - base)."""
|
|
if emo_cond_emb is spk_cond_emb or alpha == 1.0:
|
|
return self._compute_audio_emo_vec(emo_cond_emb)
|
|
if alpha == 0.0:
|
|
return self._compute_audio_emo_vec(spk_cond_emb)
|
|
base_vec = self._compute_audio_emo_vec(spk_cond_emb)
|
|
emo_vec = self._compute_audio_emo_vec(emo_cond_emb)
|
|
return base_vec + alpha * (emo_vec - base_vec) # [1, D]
|
|
|
|
# ------------------------------------------------------------------
|
|
# Text tokenizer
|
|
# ------------------------------------------------------------------
|
|
|
|
def _tokenize_text(self, text: str, device: torch.device) -> torch.Tensor:
|
|
"""Tokenize text using the BPE tokenizer, add start/stop tokens."""
|
|
if self._text_tokenizer is None:
|
|
from .tokenizer import IndexTTS2Tokenizer
|
|
|
|
bpe_path = resolve_model_file(self.model_path, "bpe.model")
|
|
if bpe_path is None:
|
|
raise FileNotFoundError(f"BPE model not found in {self.model_path}")
|
|
self._text_tokenizer = IndexTTS2Tokenizer(bpe_path, model_dir=self.model_path)
|
|
|
|
token_ids = self._text_tokenizer.encode(text, add_special_tokens=False)
|
|
|
|
# Wrap with start/stop text tokens
|
|
start_text = 0
|
|
stop_text = 1
|
|
token_ids = [start_text] + token_ids + [stop_text]
|
|
return torch.tensor([token_ids], device=device, dtype=torch.long)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Lazy loading helpers
|
|
# ------------------------------------------------------------------
|
|
|
|
def _ensure_w2v_stat_loaded(self, device: torch.device) -> None:
|
|
if self._w2v_stat is not None:
|
|
return
|
|
stat_path = resolve_model_file(self.model_path, self.config.w2v_stat)
|
|
if stat_path is not None:
|
|
raw = torch.load(stat_path, map_location="cpu", weights_only=True)
|
|
if isinstance(raw, dict):
|
|
self._w2v_stat = torch.stack([raw["mean"], raw["var"]]) # [2, 1024]
|
|
else:
|
|
self._w2v_stat = raw
|
|
else:
|
|
logger.warning("Wav2Vec2-BERT stat file not found in %s", self.model_path)
|
|
|
|
def _ensure_matrix_loaded(self, attr: str, cfg_key: str, label: str) -> None:
|
|
if getattr(self, attr) is not None:
|
|
return
|
|
path = resolve_model_file(self.model_path, getattr(self.config, cfg_key))
|
|
if path is not None:
|
|
raw = torch.load(path, map_location="cpu", weights_only=True)
|
|
setattr(self, attr, torch.split(raw, list(self.config.emo_num)))
|
|
else:
|
|
logger.warning("%s not found in %s", label, self.model_path)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Weight loading
|
|
# ------------------------------------------------------------------
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
"""Load weights from IndexTTS2 checkpoint (gpt.pth).
|
|
|
|
Bypasses vllm's default .pt iterator because the model directory
|
|
contains raw-tensor files (feat1.pt, feat2.pt) that are not state
|
|
dicts and would crash ``pt_weights_iterator``.
|
|
|
|
Weight mapping from gpt.pth checkpoint → vLLM model params:
|
|
gpt.h.{i}.* → h.{i}.* (strip ``gpt.`` prefix)
|
|
gpt.ln_f.* → ln_f.*
|
|
final_norm.* → final_norm.*
|
|
text_head.* → (skipped)
|
|
<everything else> → identity
|
|
"""
|
|
_ = weights # don't iterate — raw .pt files crash the default iterator
|
|
|
|
ckpt_path = resolve_model_file(self.model_path, self.config.gpt_checkpoint)
|
|
if ckpt_path is None:
|
|
raise FileNotFoundError(f"IndexTTS2 GPT checkpoint {self.config.gpt_checkpoint!r} not found")
|
|
state = torch.load(ckpt_path, map_location="cpu", weights_only=True)
|
|
|
|
prefix_map = {
|
|
"gpt.h.": "h.",
|
|
"gpt.ln_f.": "ln_f.",
|
|
"final_norm.": "final_norm.",
|
|
"text_head.": "_skip_text_head.",
|
|
}
|
|
|
|
params_dict = dict(self.named_parameters(remove_duplicate=False))
|
|
loaded_params: set[str] = set()
|
|
|
|
for name, loaded_weight in state.items():
|
|
# Skip attention masks
|
|
if ".attn.bias" in name or ".attn.masked_bias" in name:
|
|
continue
|
|
# Skip wpe (null position embeddings)
|
|
if ".wpe." in name or ".wte." in name:
|
|
continue
|
|
|
|
# Remap checkpoint name to model name
|
|
mapped_name = name
|
|
for old_prefix, new_prefix in prefix_map.items():
|
|
if name.startswith(old_prefix):
|
|
mapped_name = new_prefix + name[len(old_prefix) :]
|
|
break
|
|
|
|
# Skip unmapped / intentionally skipped weights
|
|
if mapped_name.startswith("_skip_"):
|
|
continue
|
|
|
|
if is_pp_missing_parameter(mapped_name, self):
|
|
continue
|
|
|
|
if mapped_name not in params_dict:
|
|
logger.debug("Skipping unrecognized weight: %s → %s", name, mapped_name)
|
|
continue
|
|
|
|
param = params_dict[mapped_name]
|
|
|
|
# GPT-2 Conv1D → Linear transpose
|
|
for conv1d_name in ["c_attn", "c_proj", "c_fc"]:
|
|
if conv1d_name not in mapped_name:
|
|
continue
|
|
if not mapped_name.endswith(".weight"):
|
|
continue
|
|
loaded_weight = loaded_weight.t()
|
|
|
|
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
|
weight_loader(param, loaded_weight)
|
|
loaded_params.add(mapped_name)
|
|
|
|
logger.info(
|
|
"Loaded %d weights for IndexTTS2TalkerForConditionalGeneration",
|
|
len(loaded_params),
|
|
)
|
|
required_prefixes = [
|
|
"text_embedding.",
|
|
"mel_embedding.",
|
|
"text_pos_embedding.",
|
|
"mel_pos_embedding.",
|
|
"speed_emb.",
|
|
"emo_layer.",
|
|
"emovec_layer.",
|
|
"conditioning_encoder.",
|
|
"perceiver_encoder.",
|
|
"emo_conditioning_encoder.",
|
|
"emo_perceiver_encoder.",
|
|
"h.",
|
|
"ln_f.",
|
|
"final_norm.",
|
|
]
|
|
if self.mel_head is not None:
|
|
required_prefixes.append("mel_head.")
|
|
missing_prefixes = [
|
|
prefix for prefix in required_prefixes if not any(name.startswith(prefix) for name in loaded_params)
|
|
]
|
|
if missing_prefixes:
|
|
raise RuntimeError(
|
|
"IndexTTS2 GPT checkpoint did not load required parameter groups: " + ", ".join(missing_prefixes)
|
|
)
|
|
|
|
# Ensure all sub-modules on correct device (some nn.Parameter created
|
|
# with torch.Tensor() end up on CPU even when model is on CUDA).
|
|
target_device = next((p.device for p in self.h.parameters()), torch.device("cpu"))
|
|
for mod in [
|
|
self.conditioning_encoder,
|
|
self.emo_conditioning_encoder,
|
|
self.perceiver_encoder,
|
|
self.emo_perceiver_encoder,
|
|
]:
|
|
mod.to(target_device)
|
|
|
|
self._warmup_external_models(target_device)
|
|
return loaded_params
|
|
|
|
def _warmup_external_models(self, device: torch.device) -> None:
|
|
"""Preload external models to eliminate first-request latency."""
|
|
try:
|
|
load_wav2vec2(self.model_path, device)
|
|
load_semantic_codec(self.model_path, self.config.semantic_codec, device)
|
|
load_campplus(self.model_path, device)
|
|
self._ensure_w2v_stat_loaded(device)
|
|
logger.info("External models preloaded: Wav2Vec2, RepCodec, CAMPPlus")
|
|
except Exception as e:
|
|
logger.warning("Failed to preload external models: %s", e)
|