83 lines
3.1 KiB
Python
83 lines
3.1 KiB
Python
# Student MeanFlow pour le DiT CosyVoice3 : v(x,t) -> u(x,t,r) (vitesse moyenne sur [t,r])
|
|
# Init: poids du teacher + gate r zero-init => student(x,t,r) == teacher v(x,t) a l'init.
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from cosyvoice.flow.DiT.dit import DiT, TimestepEmbedding
|
|
from cosyvoice.utils.mask import add_optional_chunk_mask
|
|
|
|
|
|
class MeanFlowDiT(DiT):
|
|
def __init__(self, **kw):
|
|
super().__init__(**kw)
|
|
self.r_embed = TimestepEmbedding(self.dim)
|
|
self.r_gate = nn.Linear(self.dim, self.dim)
|
|
nn.init.zeros_(self.r_gate.weight)
|
|
nn.init.zeros_(self.r_gate.bias)
|
|
|
|
@classmethod
|
|
def from_teacher(cls, teacher: DiT, **kw):
|
|
m = cls(**kw)
|
|
missing, unexpected = m.load_state_dict(teacher.state_dict(), strict=False)
|
|
assert not unexpected, f"poids teacher inattendus: {unexpected}"
|
|
assert all(k.startswith(("r_embed.", "r_gate.")) for k in missing), f"manquants hors student: {missing}"
|
|
return m
|
|
|
|
def forward(self, x, mask, mu, t, r=None, spks=None, cond=None, streaming=False):
|
|
# copie de DiT.forward avec injection de l'embedding d'intervalle r
|
|
x = x.transpose(1, 2)
|
|
mu = mu.transpose(1, 2)
|
|
cond = cond.transpose(1, 2)
|
|
spks = spks.unsqueeze(dim=1)
|
|
batch, seq_len = x.shape[0], x.shape[1]
|
|
if t.ndim == 0:
|
|
t = t.repeat(batch)
|
|
if r is None:
|
|
r = t
|
|
if r.ndim == 0:
|
|
r = r.repeat(batch)
|
|
|
|
t = self.time_embed(t) + self.r_gate(self.r_embed(r))
|
|
x = self.input_embed(x, cond, mu, spks.squeeze(1))
|
|
|
|
rope = self.rotary_embed.forward_from_seq_len(seq_len)
|
|
|
|
if self.long_skip_connection is not None:
|
|
residual = x
|
|
|
|
if streaming is True:
|
|
attn_mask = add_optional_chunk_mask(x, mask.bool(), False, False, 0, self.static_chunk_size, -1).unsqueeze(dim=1)
|
|
else:
|
|
attn_mask = add_optional_chunk_mask(x, mask.bool(), False, False, 0, 0, -1).repeat(1, x.size(1), 1).unsqueeze(dim=1)
|
|
|
|
for block in self.transformer_blocks:
|
|
x = block(x, t, mask=attn_mask.bool(), rope=rope)
|
|
|
|
if self.long_skip_connection is not None:
|
|
x = self.long_skip_connection(torch.cat((x, residual), dim=-1))
|
|
|
|
x = self.norm_out(x, t)
|
|
output = self.proj_out(x).transpose(1, 2)
|
|
return output
|
|
|
|
|
|
# kwargs exacts du DiT CosyVoice3-0.5B-2512 (cosyvoice3.yaml)
|
|
CV3_DIT_KWARGS = dict(
|
|
dim=1024, depth=22, heads=16, dim_head=64, ff_mult=2,
|
|
mel_dim=80, mu_dim=80, spk_dim=80, out_channels=80,
|
|
static_chunk_size=25 * 2, num_decoding_left_chunks=-1,
|
|
)
|
|
|
|
|
|
@torch.inference_mode()
|
|
def student_infer(student, z, mu, mask, spks, cond, t_span):
|
|
"""Inference few-step : x_{i+1} = x_i + (r-t) * u(x_i, t, r) sur la grille t_span (deja cosine-warpee)."""
|
|
x = z
|
|
for i in range(len(t_span) - 1):
|
|
t, r = t_span[i], t_span[i + 1]
|
|
bt = t.expand(x.size(0)).to(x.dtype)
|
|
br = r.expand(x.size(0)).to(x.dtype)
|
|
u = student(x, mask, mu, bt, r=br, spks=spks, cond=cond)
|
|
x = x + (r - t) * u
|
|
return x
|