Coverage for wrapper/vibevoice/configuration_vibevoice.py: 81%
113 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1# mypy: ignore-errors
2# Copy from https://github.com/microsoft/VibeVoice/blob/main/vibevoice/modular/configuration_vibevoice.py
4""" VibeVoice_AcousticTokenizer model configuration"""
6from typing import Any
7from typing import List
8from typing import Optional
10from transformers.configuration_utils import PretrainedConfig
11from transformers.utils import logging
13from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
15logger = logging.get_logger(__name__)
18class VibeVoiceAcousticTokenizerConfig(PretrainedConfig):
19 model_type = "vibevoice_acoustic_tokenizer"
21 def __init__(
22 self,
23 channels: int = 1,
24 corpus_normalize: float = 0.0,
25 causal: bool = True,
26 vae_dim: int = 64,
27 fix_std: float = 0.5,
28 std_dist_type: str = 'gaussian',
29 # common
30 mixer_layer: str = 'depthwise_conv',
31 conv_norm: str = 'none',
32 pad_mode: str = 'constant',
33 disable_last_norm: bool = True,
34 layernorm: str = 'RMSNorm',
35 layernorm_eps: float = 1e-5,
36 layernorm_elementwise_affine: bool = True,
37 conv_bias: bool = True,
38 layer_scale_init_value: float = 1e-6,
39 weight_init_value: float = 1e-2,
40 # encoder specific
41 encoder_n_filters: int = 32,
42 encoder_ratios: Optional[List[int]] = [8, 5, 5, 4, 2, 2],
43 encoder_depths: str = "3-3-3-3-3-3-8",
44 # decoder specific
45 decoder_n_filters: int = 32,
46 decoder_ratios: Optional[List[int]] = None, # if None, same as encoder
47 decoder_depths: Optional[str] = None,
48 **kwargs: Any
49 ) -> None:
50 super().__init__(**kwargs)
51 self.channels = channels
52 self.corpus_normalize = corpus_normalize
53 self.causal = causal
54 self.vae_dim = vae_dim
55 self.fix_std = fix_std
56 self.std_dist_type = std_dist_type
58 # common parameters
59 self.conv_norm = conv_norm
60 self.pad_mode = pad_mode
61 self.layernorm_eps = layernorm_eps
62 self.disable_last_norm = disable_last_norm
63 self.layernorm = layernorm
64 self.layernorm_elementwise_affine = layernorm_elementwise_affine
65 self.conv_bias = conv_bias
66 self.layer_scale_init_value = layer_scale_init_value
67 self.weight_init_value = weight_init_value
68 self.mixer_layer = mixer_layer
70 # encoder specific parameters
71 self.encoder_n_filters = encoder_n_filters
72 self.encoder_ratios = encoder_ratios
73 self.encoder_depths = encoder_depths
75 # decoder specific parameters
76 self.decoder_ratios = decoder_ratios if decoder_ratios is not None else encoder_ratios
77 self.decoder_n_filters = decoder_n_filters
78 self.decoder_depths = decoder_depths
81class VibeVoiceSemanticTokenizerConfig(PretrainedConfig):
82 model_type = "vibevoice_semantic_tokenizer"
84 def __init__(
85 self,
86 channels: int = 1,
87 corpus_normalize: float = 0.0,
88 causal: bool = True,
89 vae_dim: int = 64,
90 fix_std: float = 0,
91 std_dist_type: str = 'none',
92 # common
93 mixer_layer: str = 'depthwise_conv',
94 conv_norm: str = 'none',
95 pad_mode: str = 'constant',
96 disable_last_norm: bool = True,
97 layernorm: str = 'RMSNorm',
98 layernorm_eps: float = 1e-5,
99 layernorm_elementwise_affine: bool = True,
100 conv_bias: bool = True,
101 layer_scale_init_value: float = 1e-6,
102 weight_init_value: float = 1e-2,
103 # encoder specific
104 encoder_n_filters: int = 32,
105 encoder_ratios: Optional[List[int]] = [8, 5, 5, 4, 2, 2],
106 encoder_depths: str = "3-3-3-3-3-3-8",
107 **kwargs: Any
108 ) -> None:
109 super().__init__(**kwargs)
110 self.channels = channels
111 self.corpus_normalize = corpus_normalize
112 self.causal = causal
113 self.vae_dim = vae_dim
114 self.fix_std = fix_std
115 self.std_dist_type = std_dist_type
117 # common parameters
118 self.conv_norm = conv_norm
119 self.pad_mode = pad_mode
120 self.layernorm_eps = layernorm_eps
121 self.disable_last_norm = disable_last_norm
122 self.layernorm = layernorm
123 self.layernorm_elementwise_affine = layernorm_elementwise_affine
124 self.conv_bias = conv_bias
125 self.layer_scale_init_value = layer_scale_init_value
126 self.weight_init_value = weight_init_value
127 self.mixer_layer = mixer_layer
129 # encoder specific parameters
130 self.encoder_n_filters = encoder_n_filters
131 self.encoder_ratios = encoder_ratios
132 self.encoder_depths = encoder_depths
135class VibeVoiceDiffusionHeadConfig(PretrainedConfig):
136 model_type = "vibevoice_diffusion_head"
138 def __init__(
139 self,
140 hidden_size: int = 768,
141 head_layers: int = 4,
142 head_ffn_ratio: float = 3.0,
143 rms_norm_eps: float = 1e-5,
144 latent_size: int = 64,
145 speech_vae_dim: Optional[int] = None,
146 prediction_type: str = "v_prediction",
147 diffusion_type: str = "ddpm",
148 ddpm_num_steps: int = 1000,
149 ddpm_num_inference_steps: int = 20,
150 ddpm_beta_schedule: str = "cosine",
151 ddpm_batch_mul: int = 4,
152 **kwargs: Any
153 ) -> None:
154 self.hidden_size = hidden_size
155 self.head_layers = head_layers
156 self.head_ffn_ratio = head_ffn_ratio
157 self.rms_norm_eps = rms_norm_eps
158 self.latent_size = latent_size
159 self.speech_vae_dim = speech_vae_dim
160 self.prediction_type = prediction_type
161 self.diffusion_type = diffusion_type
162 self.ddpm_num_steps = ddpm_num_steps
163 self.ddpm_num_inference_steps = ddpm_num_inference_steps
164 self.ddpm_beta_schedule = ddpm_beta_schedule
165 self.ddpm_batch_mul = ddpm_batch_mul
167 super().__init__(**kwargs)
170class VibeVoiceConfig(PretrainedConfig):
171 model_type = "vibevoice"
172 is_composition = True
173 sub_configs = {
174 "acoustic_tokenizer_config": VibeVoiceAcousticTokenizerConfig,
175 "semantic_tokenizer_config": VibeVoiceSemanticTokenizerConfig,
176 "decoder_config": Qwen2Config,
177 "diffusion_head_config": VibeVoiceDiffusionHeadConfig,
178 }
179 # keys_to_ignore_at_inference = ["past_key_values"]
180 # Default tensor parallel plan for base model `Qwen2`
181 base_model_tp_plan = {
182 "layers.*.self_attn.q_proj": "colwise",
183 "layers.*.self_attn.k_proj": "colwise",
184 "layers.*.self_attn.v_proj": "colwise",
185 "layers.*.self_attn.o_proj": "rowwise",
186 "layers.*.mlp.gate_proj": "colwise",
187 "layers.*.mlp.up_proj": "colwise",
188 "layers.*.mlp.down_proj": "rowwise",
189 }
191 def __init__(
192 self,
193 acoustic_tokenizer_config: Optional[Any] = None,
194 semantic_tokenizer_config: Optional[Any] = None,
195 decoder_config: Optional[Any] = None,
196 diffusion_head_config: Optional[Any] = None,
197 **kwargs: Any
198 ) -> None:
199 # kwargs["_attn_implementation"] = "flash_attention_2"
200 kwargs["_attn_implementation_autoset"] = False
202 if acoustic_tokenizer_config is None:
203 self.acoustic_tokenizer_config = self.sub_configs["acoustic_tokenizer_config"]()
204 elif isinstance(acoustic_tokenizer_config, dict):
205 acoustic_tokenizer_config["model_type"] = "vibevoice_acoustic_tokenizer"
206 self.acoustic_tokenizer_config = self.sub_configs["acoustic_tokenizer_config"](**acoustic_tokenizer_config)
207 elif isinstance(acoustic_tokenizer_config, VibeVoiceAcousticTokenizerConfig):
208 # If an instance of the config class is provided
209 self.acoustic_tokenizer_config = acoustic_tokenizer_config
211 if semantic_tokenizer_config is None:
212 self.semantic_tokenizer_config = self.sub_configs["semantic_tokenizer_config"]()
213 elif isinstance(semantic_tokenizer_config, dict):
214 semantic_tokenizer_config["model_type"] = "vibevoice_semantic_tokenizer"
215 self.semantic_tokenizer_config = self.sub_configs["semantic_tokenizer_config"](**semantic_tokenizer_config)
216 elif isinstance(semantic_tokenizer_config, VibeVoiceSemanticTokenizerConfig):
217 # If an instance of the config class is provided
218 self.semantic_tokenizer_config = semantic_tokenizer_config
220 if decoder_config is None:
221 self.decoder_config = self.sub_configs["decoder_config"]()
222 elif isinstance(decoder_config, dict):
223 # If a dictionary is provided, instantiate the config class with it
224 # self.decoder_config = self.sub_configs["decoder_config"](**decoder_config)
225 if decoder_config.get("model_type", '') == "qwen2":
226 self.decoder_config = Qwen2Config(**decoder_config)
227 else:
228 raise ValueError(f"Unsupported decoder model type: {decoder_config.get('model_type', '')}")
229 elif isinstance(decoder_config, (Qwen2Config,)):
230 # If an instance of the config class is provided
231 self.decoder_config = decoder_config
233 if diffusion_head_config is None:
234 self.diffusion_head_config = self.sub_configs["diffusion_head_config"]()
235 elif isinstance(diffusion_head_config, dict):
236 diffusion_head_config["model_type"] = "vibevoice_diffusion_head"
237 self.diffusion_head_config = self.sub_configs["diffusion_head_config"](**diffusion_head_config)
238 elif isinstance(diffusion_head_config, VibeVoiceDiffusionHeadConfig):
239 # If an instance of the config class is provided
240 self.diffusion_head_config = diffusion_head_config
242 # other parameters
243 self.acoustic_vae_dim = getattr(self.acoustic_tokenizer_config, 'vae_dim', 64)
244 self.semantic_vae_dim = getattr(self.semantic_tokenizer_config, 'vae_dim', 128)
245 self.num_hidden_layers = self.decoder_config.num_hidden_layers
247 super().__init__(**kwargs)
250__all__ = [
251 "VibeVoiceAcousticTokenizerConfig",
252 "VibeVoiceSemanticTokenizerConfig",
253 "VibeVoiceDiffusionHeadConfig",
254 "VibeVoiceConfig"
255]