Spaces:
Running on Zero
Running on Zero
| import os | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from transformers import Qwen2Config, Qwen2Model | |
| from .istft import ISTFTHead | |
| _AUDIO_VAE_ATTN_ENV = "MING_TINY_AUDIO_VAE_ATTN" | |
| _AUDIO_VAE_ATTN_BACKENDS = { | |
| "config", | |
| "eager", | |
| "sdpa", | |
| "flash_attention_2", | |
| } | |
| def _build_qwen2_config(config_dict): | |
| """Build the AudioVAE backbone config with a portable attention backend. | |
| The training config records ``flash_attention_2`` for both AudioVAE | |
| backbones. Requiring that optional binary extension just to construct the | |
| HF model makes the standalone checkpoint unusable on an otherwise valid | |
| Torch/SDPA environment. SDPA is the safe inference default; an explicitly | |
| prepared runtime can opt back into FlashAttention with the environment | |
| variable below. | |
| """ | |
| backend = os.getenv(_AUDIO_VAE_ATTN_ENV, "config").strip().lower() | |
| if backend not in _AUDIO_VAE_ATTN_BACKENDS: | |
| raise ValueError( | |
| f"{_AUDIO_VAE_ATTN_ENV}={backend!r} is invalid; choose one of " | |
| f"{sorted(_AUDIO_VAE_ATTN_BACKENDS)}" | |
| ) | |
| config = Qwen2Config.from_dict(config_dict=config_dict) | |
| if backend != "config": | |
| config._attn_implementation = backend | |
| return config | |
| class StreamingLinearUpsample(nn.Module): | |
| def __init__(self, scale_factor=4): | |
| super().__init__() | |
| self.scale_factor = scale_factor | |
| self.upsampler = nn.Upsample(scale_factor=scale_factor, mode='linear', align_corners=False) | |
| def forward(self, x, state=None, is_last=False): | |
| # 初始化状态 | |
| if state is None: | |
| state = { | |
| 'prev_chunk': None, | |
| 'history_last': None, | |
| 'is_first': True | |
| } | |
| if x is None and not is_last: | |
| return None, state | |
| if state['is_first'] and is_last: | |
| out = self.upsampler(x.transpose(1, 2)).transpose(1, 2) | |
| return out, None # 结束后清除状态 | |
| output_chunks = [] | |
| if state['is_first']: | |
| state['prev_chunk'] = x | |
| state['is_first'] = False | |
| if not is_last: | |
| return None, state | |
| if state['prev_chunk'] is not None: | |
| p = state['prev_chunk'].transpose(1, 2) | |
| if state['history_last'] is None: | |
| lookahead = x[:, :1, :].transpose(1, 2) | |
| inp = torch.cat([p, lookahead], dim=2) | |
| up = self.upsampler(inp) | |
| out_prev = up[:, :, :p.size(2) * self.scale_factor] | |
| else: | |
| lookahead = x[:, :1, :].transpose(1, 2) | |
| inp = torch.cat([state['history_last'], p, lookahead], dim=2) | |
| up = self.upsampler(inp) | |
| start = self.scale_factor | |
| end = start + p.size(2) * self.scale_factor | |
| out_prev = up[:, :, start:end] | |
| output_chunks.append(out_prev.transpose(1, 2)) | |
| state['history_last'] = p[:, :, -1:] | |
| state['prev_chunk'] = x | |
| if is_last: | |
| p = state['prev_chunk'].transpose(1, 2) | |
| inp = torch.cat([state['history_last'], p], dim=2) | |
| up = self.upsampler(inp) | |
| out_last = up[:, :, self.scale_factor:] | |
| output_chunks.append(out_last.transpose(1, 2)) | |
| state = None # 结束 | |
| final_out = torch.cat(output_chunks, dim=1) if output_chunks else None | |
| return final_out, state | |
| class Encoder(nn.Module): | |
| def __init__(self, encoder_args, input_dim=320, hop_size=320, latent_dim=64, patch_size=-1): | |
| super().__init__() | |
| config = _build_qwen2_config(encoder_args) | |
| self.encoder = Qwen2Model(config) | |
| self.input_dim = input_dim | |
| self.hop_size = hop_size | |
| self.latent_dim = latent_dim | |
| self.fc1 = nn.Linear(input_dim, config.hidden_size, bias=False) | |
| self.fc2 = nn.Linear(config.hidden_size, config.hidden_size) | |
| self.fc3 = nn.Linear(config.hidden_size, latent_dim*2) | |
| self.norm = nn.LayerNorm(config.hidden_size) | |
| self.patch_size = patch_size | |
| if patch_size != -1: | |
| config.num_hidden_layers = 4 | |
| self.aggregator = Qwen2Model(config) | |
| self.cls_embed = nn.Parameter(torch.rand(1, 1, config.hidden_size)) | |
| self.cls_embed.data.normal_(0, 0.02) | |
| def get_frames(self, x): | |
| num_frames_total = (x.size(-1) + self.hop_size - 1) // self.hop_size # 向上取整的帧数 | |
| expected_len = (num_frames_total - 1) * self.hop_size + self.input_dim | |
| padding_needed = expected_len - x.size(-1) | |
| waveform = F.pad(x, (0, padding_needed), value=0.0) | |
| frames = waveform.unfold(dimension=-1, size=self.input_dim, step=self.hop_size) # [B, T, d] | |
| return frames | |
| def pad_patch_insert_cls(self, x): | |
| bsz, _, dim = x.size() | |
| num_frame = x.size(1) | |
| r = num_frame % self.patch_size | |
| pad_num = self.patch_size-r if r else 0 | |
| x = F.pad(x, (0, 0, 0, pad_num), value=0.0) # 帧数对齐到patch_size倍数 | |
| x = x.reshape(-1, self.patch_size, dim) | |
| x = torch.cat((x, self.cls_embed.expand(x.size(0), -1, -1)), dim=1) # 每个patch后插入一个cls | |
| x = x.reshape(bsz, -1, dim) | |
| return x | |
| def forward(self, waveform): | |
| x = self.get_frames(waveform) | |
| x = self.fc1(x) | |
| x = self.fc2(x) | |
| x = self.encoder(inputs_embeds=x) | |
| x = x.last_hidden_state | |
| # downsample | |
| if self.patch_size != -1: | |
| x = self.pad_patch_insert_cls(x) | |
| x = self.aggregator(inputs_embeds=x) | |
| x = x.last_hidden_state | |
| bsz, _, dim = x.size() | |
| x = x.reshape(-1, self.patch_size+1, dim) | |
| x = x[:, -1:, :].reshape(bsz, -1, dim) | |
| x = self.fc3(x) | |
| return x, waveform.unsqueeze(1) | |
| class Decoder(nn.Module): | |
| def __init__(self, decoder_args, output_dim=320, latent_dim=64, semantic_model=None, patch_size=-1): | |
| super().__init__() | |
| config = _build_qwen2_config(decoder_args) | |
| self.decoder = Qwen2Model(config) | |
| self.output_dim = output_dim | |
| self.latent_dim = latent_dim | |
| self.fc1 = nn.Linear(latent_dim, config.hidden_size) | |
| if semantic_model is not None: | |
| self.gelu = nn.GELU() | |
| self.fc2 = nn.Linear(config.hidden_size, semantic_model.audio_emb_dim) | |
| self.semantic_model = semantic_model | |
| self.fc3 = nn.Linear(semantic_model.audio_emb_dim, config.hidden_size) | |
| else: | |
| self.semantic_model = None | |
| self.hop_length = output_dim | |
| self.head = ISTFTHead(dim=config.hidden_size, n_fft=self.hop_length * 4, hop_length=self.hop_length, padding="same") | |
| self.patch_size = patch_size | |
| if self.patch_size != -1: | |
| self.upsampling = StreamingLinearUpsample(scale_factor=patch_size) | |
| def forward(self, x, only_semantic_emb=False, past_key_values=None, use_cache=False): | |
| x = self.fc1(x) | |
| if self.semantic_model is not None: | |
| x = self.fc2(self.gelu(x)) | |
| x, past_key_values = self.semantic_model(whipser_feats=x, past_key_values=past_key_values, use_cache=use_cache) | |
| unified_emb = x | |
| if only_semantic_emb: | |
| return unified_emb, past_key_values | |
| x = self.fc3(x) | |
| else: | |
| unified_emb = None | |
| if self.patch_size != -1: | |
| x = self.upsampling(x.transpose(1, 2)).transpose(1, 2) | |
| x = self.decoder(inputs_embeds=x) | |
| x = x.last_hidden_state | |
| x, _ = self.head(x) | |
| return x, unified_emb | |
| def low_level_reconstruct(self, x, past_key_values=None, use_cache=False, stream_state=None, last_chunk=False): | |
| upsample_state, audio_buffer, window_buffer = stream_state | |
| bsz, device, dtype = x.size(0), x.device, x.dtype | |
| x = self.fc1(x) | |
| if self.patch_size != -1: | |
| if use_cache: | |
| # streaming | |
| x, upsample_state = self.upsampling(x, state=upsample_state, is_last=last_chunk) | |
| if x is None: | |
| stream_state = (upsample_state, audio_buffer, window_buffer) | |
| return torch.empty(bsz, 1, 0, device=device, dtype=dtype), stream_state, past_key_values | |
| else: | |
| x = self.upsampling.upsampler(x.transpose(1, 2)).transpose(1, 2) | |
| outputs = self.decoder(inputs_embeds=x, past_key_values=past_key_values, use_cache=use_cache) | |
| past_key_values = outputs.past_key_values | |
| x = outputs.last_hidden_state | |
| x, _, audio_buffer, window_buffer = self.head(x, streaming=use_cache, audio_buffer=audio_buffer, window_buffer=window_buffer, last_chunk=last_chunk) | |
| stream_state = (upsample_state, audio_buffer, window_buffer) | |
| return x, stream_state, past_key_values | |