Update Live-streaming code (auto-daily features)
This commit is contained in:
@@ -0,0 +1,956 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import yaml
|
||||
from accelerate import Accelerator
|
||||
from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration
|
||||
from einops import rearrange
|
||||
from torch.optim import AdamW
|
||||
from train_dots_tts import DotsTtsTrainingRun
|
||||
from transformers import get_cosine_schedule_with_warmup
|
||||
|
||||
from dots_tts.config import app as app_config
|
||||
from dots_tts.data import builders as data_module
|
||||
from dots_tts.models.dots_tts import model as dots_tts_model
|
||||
from dots_tts.models.dots_tts.config import MeanFlowConfig
|
||||
from dots_tts.models.dots_tts.core import DotsTtsForwardOutput
|
||||
from dots_tts.modules.backbone.dit import DiT
|
||||
from dots_tts.training import checkpoint as train_checkpoint
|
||||
from dots_tts.training import utils as train_utils
|
||||
from dots_tts.utils import util as util_module
|
||||
|
||||
_ALLOWED_TEACHER_SOLVERS = ("euler", "midpoint", "rk4")
|
||||
_ALLOWED_CFG_DISTILL_MODES = ("natural", "fused")
|
||||
_ALLOWED_ANCHOR_TARGETS = ("formula", "teacher")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MeanFlowSettings:
|
||||
teacher_model_path: str | None
|
||||
teacher_steps: int = 8
|
||||
teacher_solver: str = "euler"
|
||||
cfg_distill_mode: str = "fused"
|
||||
distill_cfg_scale: float = 1.2
|
||||
anchor_prob: float = 0.5
|
||||
anchor_target: str = "formula"
|
||||
time_sampling_mean: float = -0.4
|
||||
time_sampling_std: float = 1.0
|
||||
train_all_parameters: bool = False
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if int(self.teacher_steps) <= 0:
|
||||
raise ValueError("teacher_steps must be positive.")
|
||||
if self.teacher_solver not in _ALLOWED_TEACHER_SOLVERS:
|
||||
raise ValueError(
|
||||
f"teacher_solver must be one of {_ALLOWED_TEACHER_SOLVERS}, "
|
||||
f"got {self.teacher_solver!r}."
|
||||
)
|
||||
if self.cfg_distill_mode not in _ALLOWED_CFG_DISTILL_MODES:
|
||||
raise ValueError(
|
||||
"cfg_distill_mode must be one of "
|
||||
f"{_ALLOWED_CFG_DISTILL_MODES}, got {self.cfg_distill_mode!r}."
|
||||
)
|
||||
if self.anchor_target not in _ALLOWED_ANCHOR_TARGETS:
|
||||
raise ValueError(
|
||||
f"anchor_target must be one of {_ALLOWED_ANCHOR_TARGETS}, "
|
||||
f"got {self.anchor_target!r}."
|
||||
)
|
||||
if not 0.0 <= float(self.anchor_prob) <= 1.0:
|
||||
raise ValueError("anchor_prob must be in [0, 1].")
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"teacher_model_path": self.teacher_model_path,
|
||||
"teacher_steps": int(self.teacher_steps),
|
||||
"teacher_solver": self.teacher_solver,
|
||||
"cfg_distill_mode": self.cfg_distill_mode,
|
||||
"distill_cfg_scale": float(self.distill_cfg_scale),
|
||||
"anchor_prob": float(self.anchor_prob),
|
||||
"anchor_target": self.anchor_target,
|
||||
"time_sampling_mean": float(self.time_sampling_mean),
|
||||
"time_sampling_std": float(self.time_sampling_std),
|
||||
"train_all_parameters": bool(self.train_all_parameters),
|
||||
}
|
||||
|
||||
|
||||
def enable_meanflow_student(model: dots_tts_model.DotsTtsModel) -> None:
|
||||
meanflow_config = MeanFlowConfig(enabled=True, use_duration_embedding=True)
|
||||
model.config.meanflow = meanflow_config
|
||||
model.core.meanflow_config = meanflow_config
|
||||
model.core.mode = "meanflow"
|
||||
|
||||
old_dit = model.core.velocity_field_predictor
|
||||
if getattr(old_dit, "duration_embedder", None) is not None:
|
||||
return
|
||||
|
||||
new_dit = DiT(
|
||||
in_dim=model.core.fm_hidden_size,
|
||||
out_dim=model.core.latent_dim,
|
||||
transformer_config=model.core.config.DiT,
|
||||
mode="meanflow",
|
||||
)
|
||||
missing_keys, unexpected_keys = new_dit.load_state_dict(
|
||||
old_dit.state_dict(),
|
||||
strict=False,
|
||||
)
|
||||
missing_keys = [
|
||||
key for key in missing_keys if not key.startswith("duration_embedder.")
|
||||
]
|
||||
if missing_keys or unexpected_keys:
|
||||
raise RuntimeError(
|
||||
"Failed to initialize MeanFlow DiT from the pretrained flow-matching "
|
||||
f"DiT: missing={missing_keys[:5]} unexpected={unexpected_keys[:5]}"
|
||||
)
|
||||
duration_output = new_dit.duration_embedder.mlp[-1]
|
||||
nn.init.zeros_(duration_output.weight)
|
||||
nn.init.zeros_(duration_output.bias)
|
||||
model.core.velocity_field_predictor = new_dit
|
||||
|
||||
|
||||
class MeanFlowDotsTtsModel(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
student: dots_tts_model.DotsTtsModel,
|
||||
settings: MeanFlowSettings,
|
||||
):
|
||||
super().__init__()
|
||||
self.student = student
|
||||
self.settings = settings
|
||||
self._teacher_holder: dict[str, dots_tts_model.DotsTtsModel] = {}
|
||||
|
||||
@property
|
||||
def config(self):
|
||||
return self.student.config
|
||||
|
||||
@property
|
||||
def tokenizer(self):
|
||||
return self.student.tokenizer
|
||||
|
||||
@property
|
||||
def teacher(self) -> dots_tts_model.DotsTtsModel:
|
||||
teacher = self._teacher_holder.get("model")
|
||||
if teacher is None:
|
||||
raise RuntimeError("MeanFlow teacher model has not been initialized.")
|
||||
return teacher
|
||||
|
||||
def set_teacher(self, teacher: dots_tts_model.DotsTtsModel) -> None:
|
||||
for param in teacher.parameters():
|
||||
param.requires_grad_(False)
|
||||
teacher.eval()
|
||||
self._teacher_holder["model"] = teacher
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
super().to(*args, **kwargs)
|
||||
teacher = self._teacher_holder.get("model")
|
||||
if teacher is not None:
|
||||
self._teacher_holder["model"] = teacher.to(*args, **kwargs)
|
||||
self._teacher_holder["model"].eval()
|
||||
return self
|
||||
|
||||
def cuda(self, device=None):
|
||||
super().cuda(device)
|
||||
teacher = self._teacher_holder.get("model")
|
||||
if teacher is not None:
|
||||
self._teacher_holder["model"] = teacher.cuda(device).eval()
|
||||
return self
|
||||
|
||||
def train(self, mode: bool = True):
|
||||
super().train(mode)
|
||||
teacher = self._teacher_holder.get("model")
|
||||
if teacher is not None:
|
||||
teacher.eval()
|
||||
return self
|
||||
|
||||
def prepare_training_batch(self, data: dict[str, Any]) -> dict[str, Any]:
|
||||
return self.student.prepare_training_batch(data)
|
||||
|
||||
def save_pretrained(self, save_directory: str | Path) -> Path:
|
||||
return self.student.save_pretrained(save_directory)
|
||||
|
||||
def load_pretrained_weights(
|
||||
self, pretrained_model_name_or_path: str | Path
|
||||
) -> None:
|
||||
self.student.load_pretrained_weights(pretrained_model_name_or_path)
|
||||
|
||||
def set_cfg_droprate(
|
||||
self,
|
||||
cfg_droprate: float | None = None,
|
||||
xvec_drop_rate: float | None = None,
|
||||
) -> None:
|
||||
self.student.set_cfg_droprate(
|
||||
cfg_droprate=cfg_droprate,
|
||||
xvec_drop_rate=xvec_drop_rate,
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def compute_teacher_meanflow_target(
|
||||
self,
|
||||
*,
|
||||
xt: torch.Tensor,
|
||||
t: torch.Tensor,
|
||||
delta_t: torch.Tensor,
|
||||
prefix_data: dict[str, Any],
|
||||
g_cond: torch.Tensor | None,
|
||||
cfg_distill: bool,
|
||||
uncond_prefix_data: dict[str, Any] | None,
|
||||
uncond_g_cond: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
teacher_core = self.teacher.core
|
||||
teacher_dit = teacher_core.velocity_field_predictor
|
||||
io_helper = teacher_core.io_helper
|
||||
noisy_proj = teacher_core.coordinate_proj
|
||||
n_steps = int(self.settings.teacher_steps)
|
||||
solver = self.settings.teacher_solver
|
||||
cfg_scale = float(self.settings.distill_cfg_scale)
|
||||
|
||||
if solver not in _ALLOWED_TEACHER_SOLVERS:
|
||||
raise ValueError(f"Unsupported teacher solver: {solver!r}.")
|
||||
|
||||
device = xt.device
|
||||
batch_size = xt.size(0)
|
||||
latent_lens = prefix_data["latent_lens"]
|
||||
latent_patch_size = int(prefix_data["latent_patch_size"])
|
||||
anchor_mask = delta_t.float() == 0
|
||||
|
||||
autocast_device = "cuda" if device.type == "cuda" else "cpu"
|
||||
with torch.autocast(device_type=autocast_device, enabled=False):
|
||||
z = xt.float()
|
||||
cur_t = t.float()
|
||||
safe_dt = delta_t.float().clamp(min=1e-6)
|
||||
step_dt = safe_dt / n_steps
|
||||
|
||||
def evaluate(z_in: torch.Tensor, t_val: torch.Tensor) -> torch.Tensor:
|
||||
fm_seq = io_helper.replace_noise_latents_in_fm_seq(
|
||||
prefix_data,
|
||||
z_in.to(xt.dtype),
|
||||
noisy_proj,
|
||||
).float()
|
||||
vt = teacher_dit(
|
||||
x=fm_seq,
|
||||
timesteps=t_val,
|
||||
pos_ids=prefix_data["fm_pos_ids"],
|
||||
mask=prefix_data["fm_seq_mask"],
|
||||
attn_mask=prefix_data["fm_attn_mask"],
|
||||
g_cond=None if g_cond is None else g_cond.float(),
|
||||
)
|
||||
pred = io_helper.get_dit_outputs(
|
||||
pred_v=vt,
|
||||
fm_prefix_lengths=prefix_data["fm_prefix_lengths"],
|
||||
fm_gen_lengths=prefix_data["fm_gen_lengths"],
|
||||
fm_gen_patch_size=prefix_data["fm_gen_patch_size"],
|
||||
latent_patch_size=prefix_data["latent_patch_size"],
|
||||
)
|
||||
|
||||
if cfg_distill:
|
||||
if uncond_prefix_data is None:
|
||||
raise RuntimeError(
|
||||
"CFG distillation requires an uncond prefix."
|
||||
)
|
||||
fm_seq_u = io_helper.replace_noise_latents_in_fm_seq(
|
||||
uncond_prefix_data,
|
||||
z_in.to(xt.dtype),
|
||||
noisy_proj,
|
||||
).float()
|
||||
vt_u = teacher_dit(
|
||||
x=fm_seq_u,
|
||||
timesteps=t_val,
|
||||
pos_ids=uncond_prefix_data["fm_pos_ids"],
|
||||
mask=uncond_prefix_data["fm_seq_mask"],
|
||||
attn_mask=uncond_prefix_data["fm_attn_mask"],
|
||||
g_cond=None if uncond_g_cond is None else uncond_g_cond.float(),
|
||||
)
|
||||
pred_u = io_helper.get_dit_outputs(
|
||||
pred_v=vt_u,
|
||||
fm_prefix_lengths=uncond_prefix_data["fm_prefix_lengths"],
|
||||
fm_gen_lengths=uncond_prefix_data["fm_gen_lengths"],
|
||||
fm_gen_patch_size=uncond_prefix_data["fm_gen_patch_size"],
|
||||
latent_patch_size=uncond_prefix_data["latent_patch_size"],
|
||||
)
|
||||
pred = pred + cfg_scale * (pred - pred_u)
|
||||
return rearrange(pred, "n p d -> (n p) d")
|
||||
|
||||
v_init_flat = evaluate(z, cur_t)
|
||||
|
||||
def apply_velocity(
|
||||
z_cur: torch.Tensor,
|
||||
v_flat: torch.Tensor,
|
||||
*,
|
||||
dt_factor: float,
|
||||
) -> torch.Tensor:
|
||||
new_z = z_cur.clone()
|
||||
offset = 0
|
||||
for batch_idx in range(batch_size):
|
||||
length = int(latent_lens[batch_idx].item())
|
||||
if length <= 0:
|
||||
continue
|
||||
if not bool(anchor_mask[batch_idx].item()):
|
||||
new_z[batch_idx, :length, :] = z_cur[
|
||||
batch_idx, :length, :
|
||||
] + v_flat[offset : offset + length, :] * (
|
||||
step_dt[batch_idx] * float(dt_factor)
|
||||
)
|
||||
offset += length
|
||||
return new_z
|
||||
|
||||
if solver == "euler":
|
||||
v_flat = v_init_flat
|
||||
for step in range(n_steps):
|
||||
if step > 0:
|
||||
v_flat = evaluate(z, cur_t)
|
||||
z = apply_velocity(z, v_flat, dt_factor=1.0)
|
||||
cur_t = cur_t + step_dt
|
||||
elif solver == "midpoint":
|
||||
for step in range(n_steps):
|
||||
k1 = v_init_flat if step == 0 else evaluate(z, cur_t)
|
||||
z_mid = apply_velocity(z, k1, dt_factor=0.5)
|
||||
k2 = evaluate(z_mid, cur_t + 0.5 * step_dt)
|
||||
z = apply_velocity(z, k2, dt_factor=1.0)
|
||||
cur_t = cur_t + step_dt
|
||||
else:
|
||||
for step in range(n_steps):
|
||||
k1 = v_init_flat if step == 0 else evaluate(z, cur_t)
|
||||
z1 = apply_velocity(z, k1, dt_factor=0.5)
|
||||
k2 = evaluate(z1, cur_t + 0.5 * step_dt)
|
||||
z2 = apply_velocity(z, k2, dt_factor=0.5)
|
||||
k3 = evaluate(z2, cur_t + 0.5 * step_dt)
|
||||
z3 = apply_velocity(z, k3, dt_factor=1.0)
|
||||
k4 = evaluate(z3, cur_t + step_dt)
|
||||
z = apply_velocity(
|
||||
z,
|
||||
(k1 + 2.0 * k2 + 2.0 * k3 + k4) / 6.0,
|
||||
dt_factor=1.0,
|
||||
)
|
||||
cur_t = cur_t + step_dt
|
||||
|
||||
mean_velocity = (z - xt.float()) / safe_dt.view(-1, 1, 1)
|
||||
target_chunks = []
|
||||
offset = 0
|
||||
for batch_idx in range(batch_size):
|
||||
length = int(latent_lens[batch_idx].item())
|
||||
if length <= 0:
|
||||
continue
|
||||
if bool(anchor_mask[batch_idx].item()):
|
||||
target_b = v_init_flat[offset : offset + length, :]
|
||||
else:
|
||||
target_b = mean_velocity[batch_idx, :length, :]
|
||||
target_chunks.append(
|
||||
rearrange(target_b, "(n p) d -> n p d", p=latent_patch_size)
|
||||
)
|
||||
offset += length
|
||||
if not target_chunks:
|
||||
raise RuntimeError("Teacher rollout produced no MeanFlow target.")
|
||||
return torch.cat(target_chunks, dim=0).to(xt.dtype)
|
||||
|
||||
def forward(self, data: dict[str, Any]):
|
||||
loss_masks = data["loss_masks"]
|
||||
processed = self.student.prepare_training_inputs(data)
|
||||
processed["input_span_mask"] = data["input_span_mask"]
|
||||
processed["output_span_mask"] = data["output_span_mask"]
|
||||
outputs = self.meanflow_forward(processed)
|
||||
return self.student._compute_loss_terms(
|
||||
outputs,
|
||||
labels=processed["labels"],
|
||||
loss_masks=loss_masks,
|
||||
)
|
||||
|
||||
def meanflow_forward(self, data: dict[str, Any]) -> DotsTtsForwardOutput:
|
||||
core = self.student.core
|
||||
input_ids: torch.Tensor = data["input_ids"]
|
||||
input_ids_lengths: torch.Tensor = data["input_ids_lengths"]
|
||||
input_span_mask: torch.Tensor = data["input_span_mask"]
|
||||
output_span_mask: torch.Tensor = data["output_span_mask"]
|
||||
batch_size = input_ids.size(0)
|
||||
device = input_ids.device
|
||||
|
||||
latents: torch.Tensor | None = data.get("latents")
|
||||
latents_sampled: torch.Tensor | None = data.get("latents_sampled")
|
||||
latent_lengths: torch.Tensor | None = data.get("latent_lengths")
|
||||
has_latents = latents is not None or latents_sampled is not None
|
||||
|
||||
if has_latents:
|
||||
if latents_sampled is None:
|
||||
latents_sampled = core.io_helper.sample_from_latent(latents)
|
||||
patch_embeddings = core.patch_encoder(
|
||||
latents_sampled, x_lens=latent_lengths
|
||||
)
|
||||
valid_patch_counts = latent_lengths // core.latent_patch_size
|
||||
latents_sampled = core.io_helper.normalize(latents_sampled)
|
||||
else:
|
||||
latents_sampled = None
|
||||
patch_embeddings = None
|
||||
valid_patch_counts = torch.zeros(
|
||||
batch_size,
|
||||
dtype=torch.long,
|
||||
device=device,
|
||||
)
|
||||
|
||||
input_span_counts = input_span_mask.sum(dim=1)
|
||||
if input_span_counts.sum() > 0 and patch_embeddings is None:
|
||||
raise RuntimeError(
|
||||
"Found audio span tokens but no latents provided to compute patch embeddings."
|
||||
)
|
||||
|
||||
inputs_embeds = core.llm.get_input_embeddings()(input_ids)
|
||||
if patch_embeddings is not None:
|
||||
inputs_embeds = inputs_embeds.clone()
|
||||
patch_embeddings = patch_embeddings.to(inputs_embeds.dtype)
|
||||
for batch_idx in range(batch_size):
|
||||
span_num = int(input_span_counts[batch_idx].item())
|
||||
if span_num == 0:
|
||||
continue
|
||||
expected = int(valid_patch_counts[batch_idx].item())
|
||||
if expected != span_num:
|
||||
raise RuntimeError(
|
||||
f"Mismatch between span tokens ({span_num}) and latent patches "
|
||||
f"({expected}) for sample {batch_idx}."
|
||||
)
|
||||
indices = input_span_mask[batch_idx].nonzero(as_tuple=False).squeeze(-1)
|
||||
inputs_embeds[batch_idx, indices, :] = patch_embeddings[
|
||||
batch_idx,
|
||||
:span_num,
|
||||
:,
|
||||
]
|
||||
|
||||
_llm_attn_mask, llm_seq_mask, _ = core.causal_helper.create_causal_mask_and_pos(
|
||||
seq_lens=input_ids_lengths,
|
||||
max_len=input_ids.size(1),
|
||||
)
|
||||
llm_outputs = core.llm(
|
||||
inputs_embeds=inputs_embeds,
|
||||
attention_mask=llm_seq_mask.long(),
|
||||
use_cache=False,
|
||||
output_hidden_states=True,
|
||||
return_dict=True,
|
||||
)
|
||||
llm_logits = llm_outputs.logits
|
||||
llm_hidden = llm_outputs.hidden_states[-1]
|
||||
eos = core.eos_proj(llm_hidden.detach())
|
||||
|
||||
total_patches = int(output_span_mask.sum().item())
|
||||
if total_patches > 0 and latents_sampled is None:
|
||||
raise RuntimeError("MeanFlow training requested but latents are missing.")
|
||||
|
||||
if total_patches > 0:
|
||||
pred, target = self.meanflow_fm_segment(
|
||||
data,
|
||||
llm_hidden=llm_hidden,
|
||||
inputs_embeds=inputs_embeds,
|
||||
output_span_mask=output_span_mask,
|
||||
latents_sampled=latents_sampled,
|
||||
latent_lengths=latent_lengths,
|
||||
)
|
||||
else:
|
||||
pred, target = self.dummy_fm_forward(core, llm_hidden, device)
|
||||
|
||||
return DotsTtsForwardOutput(
|
||||
llm_logits=llm_logits,
|
||||
pred=pred,
|
||||
target=target,
|
||||
eos_out=eos,
|
||||
)
|
||||
|
||||
def meanflow_fm_segment(
|
||||
self,
|
||||
data: dict[str, Any],
|
||||
*,
|
||||
llm_hidden: torch.Tensor,
|
||||
inputs_embeds: torch.Tensor,
|
||||
output_span_mask: torch.Tensor,
|
||||
latents_sampled: torch.Tensor,
|
||||
latent_lengths: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
core = self.student.core
|
||||
teacher_core = self.teacher.core
|
||||
settings = self.settings
|
||||
batch_size = latents_sampled.size(0)
|
||||
device = latents_sampled.device
|
||||
latent_dtype = latents_sampled.dtype
|
||||
first_t = torch.randn(batch_size, device=device, dtype=latent_dtype)
|
||||
second_t = torch.randn(batch_size, device=device, dtype=latent_dtype)
|
||||
first_t = torch.sigmoid(
|
||||
first_t * float(settings.time_sampling_std)
|
||||
+ float(settings.time_sampling_mean)
|
||||
)
|
||||
second_t = torch.sigmoid(
|
||||
second_t * float(settings.time_sampling_std)
|
||||
+ float(settings.time_sampling_mean)
|
||||
)
|
||||
t_vec = torch.minimum(first_t, second_t)
|
||||
delta_t = (first_t - second_t).abs()
|
||||
anchor_mask = torch.rand(batch_size, device=device, dtype=latent_dtype) < float(
|
||||
settings.anchor_prob
|
||||
)
|
||||
delta_t = torch.where(anchor_mask, torch.zeros_like(delta_t), delta_t)
|
||||
z0 = torch.randn_like(latents_sampled)
|
||||
xt = core.fm_helper.sample_x_t(
|
||||
z0,
|
||||
latents_sampled,
|
||||
t_vec.view(-1, 1, 1).to(latent_dtype),
|
||||
)
|
||||
|
||||
fused_cfg = settings.cfg_distill_mode == "fused"
|
||||
if fused_cfg:
|
||||
cfg_mask = torch.zeros(batch_size, device=device, dtype=torch.bool)
|
||||
xvec_drop_mask = torch.zeros(batch_size, device=device, dtype=torch.bool)
|
||||
else:
|
||||
cfg_mask = torch.empty(
|
||||
batch_size, device=device, dtype=torch.float32
|
||||
).uniform_(0, 1) < float(core.cfg_droprate)
|
||||
xvec_drop_mask = torch.empty(
|
||||
batch_size, device=device, dtype=torch.float32
|
||||
).uniform_(0, 1) < float(core.xvec_drop_rate)
|
||||
|
||||
xvec_cond = core.xvec_proj(data["xvector"])
|
||||
vocal_mask = data.get("vocal_mask")
|
||||
if vocal_mask is None:
|
||||
vocal_mask = torch.ones(batch_size, device=device, dtype=torch.bool)
|
||||
xvec_cond = util_module.mask_data(xvec_cond, xvec_drop_mask & vocal_mask)
|
||||
|
||||
hiddens_for_fm = torch.where(
|
||||
output_span_mask.unsqueeze(-1),
|
||||
llm_hidden,
|
||||
inputs_embeds,
|
||||
)
|
||||
prefix_data = core.io_helper.prepare_meanflow_inputs_for_dit(
|
||||
hiddens=hiddens_for_fm,
|
||||
latents=latents_sampled,
|
||||
latent_lens=latent_lengths,
|
||||
hidden_proj=core.hidden_proj,
|
||||
latent_proj=core.latent_proj,
|
||||
noisy_proj=core.coordinate_proj,
|
||||
span_mask=output_span_mask,
|
||||
hidden_patch_size=core.hidden_patch_size,
|
||||
latent_patch_size=core.latent_patch_size,
|
||||
cfg_mask=cfg_mask,
|
||||
noise_latents=xt,
|
||||
)
|
||||
|
||||
uncond_prefix_data = None
|
||||
uncond_g_cond = None
|
||||
with torch.no_grad():
|
||||
teacher_xvec_cond = teacher_core.xvec_proj(data["xvector"])
|
||||
teacher_xvec_cond = util_module.mask_data(
|
||||
teacher_xvec_cond,
|
||||
xvec_drop_mask & vocal_mask,
|
||||
)
|
||||
teacher_prefix_data = (
|
||||
teacher_core.io_helper.prepare_meanflow_inputs_for_dit(
|
||||
hiddens=hiddens_for_fm,
|
||||
latents=latents_sampled,
|
||||
latent_lens=latent_lengths,
|
||||
hidden_proj=teacher_core.hidden_proj,
|
||||
latent_proj=teacher_core.latent_proj,
|
||||
noisy_proj=teacher_core.coordinate_proj,
|
||||
span_mask=output_span_mask,
|
||||
hidden_patch_size=teacher_core.hidden_patch_size,
|
||||
latent_patch_size=teacher_core.latent_patch_size,
|
||||
cfg_mask=cfg_mask,
|
||||
noise_latents=xt,
|
||||
)
|
||||
)
|
||||
if fused_cfg:
|
||||
uncond_prefix_data = (
|
||||
teacher_core.io_helper.prepare_meanflow_inputs_for_dit(
|
||||
hiddens=hiddens_for_fm,
|
||||
latents=latents_sampled,
|
||||
latent_lens=latent_lengths,
|
||||
hidden_proj=teacher_core.hidden_proj,
|
||||
latent_proj=teacher_core.latent_proj,
|
||||
noisy_proj=teacher_core.coordinate_proj,
|
||||
span_mask=output_span_mask,
|
||||
hidden_patch_size=teacher_core.hidden_patch_size,
|
||||
latent_patch_size=teacher_core.latent_patch_size,
|
||||
cfg_mask=torch.ones(
|
||||
batch_size, device=device, dtype=torch.bool
|
||||
),
|
||||
noise_latents=xt,
|
||||
)
|
||||
)
|
||||
uncond_g_cond = torch.zeros_like(teacher_xvec_cond)
|
||||
|
||||
teacher_target = self.compute_teacher_meanflow_target(
|
||||
xt=xt,
|
||||
t=t_vec,
|
||||
delta_t=delta_t,
|
||||
prefix_data=teacher_prefix_data,
|
||||
g_cond=teacher_xvec_cond,
|
||||
cfg_distill=fused_cfg,
|
||||
uncond_prefix_data=uncond_prefix_data,
|
||||
uncond_g_cond=uncond_g_cond,
|
||||
)
|
||||
if anchor_mask.any() and settings.anchor_target == "formula":
|
||||
target = self.replace_anchor_targets_with_formula(
|
||||
teacher_target,
|
||||
z0=z0,
|
||||
latents_sampled=latents_sampled,
|
||||
latent_lengths=latent_lengths,
|
||||
anchor_mask=anchor_mask,
|
||||
)
|
||||
else:
|
||||
target = teacher_target
|
||||
|
||||
student_vt = core.velocity_field_predictor(
|
||||
x=prefix_data["fm_seq"],
|
||||
timesteps=t_vec,
|
||||
duration=delta_t,
|
||||
pos_ids=prefix_data["fm_pos_ids"],
|
||||
mask=prefix_data["fm_seq_mask"],
|
||||
attn_mask=prefix_data["fm_attn_mask"],
|
||||
g_cond=xvec_cond,
|
||||
)
|
||||
pred = core.io_helper.get_dit_outputs(
|
||||
pred_v=student_vt,
|
||||
fm_prefix_lengths=prefix_data["fm_prefix_lengths"],
|
||||
fm_gen_lengths=prefix_data["fm_gen_lengths"],
|
||||
fm_gen_patch_size=prefix_data["fm_gen_patch_size"],
|
||||
latent_patch_size=prefix_data["latent_patch_size"],
|
||||
)
|
||||
return pred, target
|
||||
|
||||
def replace_anchor_targets_with_formula(
|
||||
self,
|
||||
teacher_target: torch.Tensor,
|
||||
*,
|
||||
z0: torch.Tensor,
|
||||
latents_sampled: torch.Tensor,
|
||||
latent_lengths: torch.Tensor,
|
||||
anchor_mask: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
core = self.student.core
|
||||
formula_target = core.fm_helper.compute_u_t(z0, latents_sampled)
|
||||
chunks = []
|
||||
offset = 0
|
||||
for batch_idx in range(latents_sampled.size(0)):
|
||||
length = int(latent_lengths[batch_idx].item())
|
||||
if length <= 0:
|
||||
continue
|
||||
patch_count = length // core.latent_patch_size
|
||||
if bool(anchor_mask[batch_idx].item()):
|
||||
chunks.append(
|
||||
rearrange(
|
||||
formula_target[batch_idx, :length, :],
|
||||
"(n p) d -> n p d",
|
||||
p=core.latent_patch_size,
|
||||
)
|
||||
)
|
||||
else:
|
||||
chunks.append(teacher_target[offset : offset + patch_count])
|
||||
offset += patch_count
|
||||
if not chunks:
|
||||
raise RuntimeError("Anchor target replacement produced no target.")
|
||||
return torch.cat(chunks, dim=0)
|
||||
|
||||
def dummy_fm_forward(
|
||||
self,
|
||||
core,
|
||||
llm_hidden: torch.Tensor,
|
||||
device: torch.device,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
dummy_length = core.latent_patch_size
|
||||
dummy_seq_h = llm_hidden.new_zeros((1, dummy_length, core.llm_hidden_size))
|
||||
dummy_seq_h = core.hidden_proj(dummy_seq_h) * 0.0
|
||||
dummy_seq_l = llm_hidden.new_zeros((1, dummy_length, core.latent_dim))
|
||||
dummy_seq_l = core.latent_proj(dummy_seq_l) * 0.0
|
||||
dummy_seq_c = llm_hidden.new_zeros((1, dummy_length, core.latent_dim))
|
||||
dummy_seq_c = core.coordinate_proj(dummy_seq_c) * 0.0
|
||||
dummy_seq = dummy_seq_h + dummy_seq_l + dummy_seq_c
|
||||
dummy_times = torch.zeros((1,), device=device, dtype=torch.float32)
|
||||
dummy_duration = torch.zeros((1,), device=device, dtype=torch.float32)
|
||||
dummy_attn_mask = torch.ones(
|
||||
(1, dummy_length, dummy_length),
|
||||
device=device,
|
||||
dtype=torch.bool,
|
||||
)
|
||||
dummy_out = core.velocity_field_predictor(
|
||||
x=dummy_seq,
|
||||
timesteps=dummy_times,
|
||||
duration=dummy_duration,
|
||||
attn_mask=dummy_attn_mask,
|
||||
)
|
||||
pred = dummy_out[:, -core.latent_patch_size :, :]
|
||||
return pred, pred.detach()
|
||||
|
||||
|
||||
class DotsTtsMeanFlowTrainingRun(DotsTtsTrainingRun):
|
||||
def __init__(
|
||||
self,
|
||||
cfg: app_config.AppConfig,
|
||||
*,
|
||||
meanflow_settings: MeanFlowSettings,
|
||||
debug_enabled: bool = False,
|
||||
):
|
||||
self.cfg = cfg
|
||||
self.meanflow_settings = meanflow_settings
|
||||
self.progress = train_utils.TrainProgress()
|
||||
self.max_train_steps = int(cfg.train.max_train_steps)
|
||||
self.grad_accumulation_steps = int(cfg.train.gradient_accumulation_steps)
|
||||
self.last_validation_step: int | None = None
|
||||
self.consecutive_empty_epochs = 0
|
||||
self.saved_latest_checkpoint = False
|
||||
self._last_log_step = 0
|
||||
self._last_log_time = 0.0
|
||||
self._debug_enabled = bool(debug_enabled)
|
||||
self._debug_batch_count = 0
|
||||
self._debug_audio_sample_rate = int(self.cfg.train_data.train_audio_sample_rate)
|
||||
|
||||
project_config = ProjectConfiguration(
|
||||
project_dir=self.cfg.train.output_dir,
|
||||
total_limit=self.cfg.train.max_checkpoints_to_keep,
|
||||
)
|
||||
ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)
|
||||
self.accelerator = Accelerator(
|
||||
kwargs_handlers=[ddp_kwargs],
|
||||
gradient_accumulation_steps=self.grad_accumulation_steps,
|
||||
log_with="tensorboard",
|
||||
project_config=project_config,
|
||||
step_scheduler_with_optimizer=False,
|
||||
)
|
||||
|
||||
util_module.seed_everything(self.cfg.train.seed)
|
||||
|
||||
student = dots_tts_model.DotsTtsModel.from_pretrained(
|
||||
self.cfg.train.pretrained_model_path
|
||||
)
|
||||
student.set_cfg_droprate(
|
||||
cfg_droprate=self.cfg.train.cfg_droprate,
|
||||
xvec_drop_rate=self.cfg.train.xvec_drop_rate,
|
||||
)
|
||||
enable_meanflow_student(student)
|
||||
if not bool(meanflow_settings.train_all_parameters):
|
||||
for param in student.parameters():
|
||||
param.requires_grad_(False)
|
||||
for param in student.core.velocity_field_predictor.parameters():
|
||||
param.requires_grad_(True)
|
||||
model = MeanFlowDotsTtsModel(student, meanflow_settings)
|
||||
|
||||
teacher_path = (
|
||||
meanflow_settings.teacher_model_path or self.cfg.train.pretrained_model_path
|
||||
)
|
||||
teacher = dots_tts_model.DotsTtsModel.from_pretrained(teacher_path)
|
||||
model.set_teacher(teacher)
|
||||
|
||||
optimizer = AdamW(
|
||||
(param for param in model.parameters() if param.requires_grad),
|
||||
lr=self.cfg.train.learning_rate,
|
||||
weight_decay=self.cfg.train.weight_decay,
|
||||
)
|
||||
scheduler = get_cosine_schedule_with_warmup(
|
||||
optimizer,
|
||||
num_warmup_steps=self.cfg.train.warmup_steps,
|
||||
num_training_steps=self.max_train_steps,
|
||||
)
|
||||
self.model, self.optimizer, self.scheduler = self.accelerator.prepare(
|
||||
model,
|
||||
optimizer,
|
||||
scheduler,
|
||||
)
|
||||
self.unwrapped_model = self.accelerator.unwrap_model(self.model)
|
||||
self.unwrapped_model.to(self.accelerator.device)
|
||||
|
||||
expected_sample_rate = int(self.unwrapped_model.config.vocoder.sample_rate)
|
||||
expected_audio_samples_per_llm_token = int(
|
||||
self.unwrapped_model.student.hop_size
|
||||
) * int(self.unwrapped_model.config.patch_size)
|
||||
if int(self.cfg.train_data.train_audio_sample_rate) != expected_sample_rate:
|
||||
raise ValueError(
|
||||
f"train_data.train_audio_sample_rate={int(self.cfg.train_data.train_audio_sample_rate)} "
|
||||
f"does not match the pretrained model sample rate {expected_sample_rate}."
|
||||
)
|
||||
if (
|
||||
int(self.cfg.train_data.audio_samples_per_llm_token)
|
||||
!= expected_audio_samples_per_llm_token
|
||||
):
|
||||
raise ValueError(
|
||||
"train_data.audio_samples_per_llm_token="
|
||||
f"{int(self.cfg.train_data.audio_samples_per_llm_token)} "
|
||||
"does not match the pretrained model audio token contract "
|
||||
f"{expected_audio_samples_per_llm_token}."
|
||||
)
|
||||
if self.cfg.val_data is not None:
|
||||
if int(self.cfg.val_data.train_audio_sample_rate) != expected_sample_rate:
|
||||
raise ValueError(
|
||||
f"val_data.train_audio_sample_rate={int(self.cfg.val_data.train_audio_sample_rate)} "
|
||||
f"does not match the pretrained model sample rate {expected_sample_rate}."
|
||||
)
|
||||
if (
|
||||
int(self.cfg.val_data.audio_samples_per_llm_token)
|
||||
!= expected_audio_samples_per_llm_token
|
||||
):
|
||||
raise ValueError(
|
||||
"val_data.audio_samples_per_llm_token="
|
||||
f"{int(self.cfg.val_data.audio_samples_per_llm_token)} "
|
||||
"does not match the pretrained model audio token contract "
|
||||
f"{expected_audio_samples_per_llm_token}."
|
||||
)
|
||||
|
||||
if self.accelerator.is_main_process:
|
||||
total_params = sum(
|
||||
param.numel() for param in self.unwrapped_model.parameters()
|
||||
)
|
||||
trainable_params = sum(
|
||||
param.numel()
|
||||
for param in self.unwrapped_model.parameters()
|
||||
if param.requires_grad
|
||||
)
|
||||
self.accelerator.print(f"Total parameters: {total_params:,}")
|
||||
self.accelerator.print(f"Trainable parameters: {trainable_params:,}")
|
||||
self.accelerator.print(
|
||||
f"MeanFlow teacher path: {Path(teacher_path).expanduser()}"
|
||||
)
|
||||
self.accelerator.print(
|
||||
f"Distributed type: {self.accelerator.distributed_type}"
|
||||
)
|
||||
|
||||
tokenizer = self.unwrapped_model.tokenizer
|
||||
self.tokenizer = tokenizer
|
||||
train_dataset = data_module.build_training_dataset(
|
||||
self.cfg.train_data,
|
||||
tokenizer=tokenizer,
|
||||
seed=int(self.cfg.train.seed),
|
||||
accelerator=self.accelerator,
|
||||
)
|
||||
self.train_loader = data_module.build_training_dataloader(
|
||||
train_dataset,
|
||||
self.cfg.train_data,
|
||||
tokenizer=tokenizer,
|
||||
)
|
||||
|
||||
self.val_loader = None
|
||||
if self.cfg.train.eval_interval is not None or self.cfg.train.run_eval_on_start:
|
||||
if self.cfg.val_data is None:
|
||||
raise ValueError(
|
||||
"Validation requires val_data when eval_interval or "
|
||||
"run_eval_on_start is enabled."
|
||||
)
|
||||
validation_data_cfg = self.cfg.val_data.model_copy(deep=True)
|
||||
validation_data_cfg.num_tokens_per_epoch = None
|
||||
val_dataset = data_module.build_validation_dataset(
|
||||
validation_data_cfg,
|
||||
tokenizer=tokenizer,
|
||||
seed=int(self.cfg.train.seed),
|
||||
accelerator=self.accelerator,
|
||||
)
|
||||
self.val_loader = data_module.build_validation_dataloader(
|
||||
val_dataset,
|
||||
validation_data_cfg,
|
||||
tokenizer=tokenizer,
|
||||
)
|
||||
|
||||
self._resume_if_available()
|
||||
self.train_loader.set_epoch(self.progress.epoch)
|
||||
|
||||
def _write_run_config(self) -> None:
|
||||
if not bool(getattr(self.accelerator, "is_main_process", True)):
|
||||
return
|
||||
output_dir = Path(self.cfg.train.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
config_path = output_dir / "config.yml"
|
||||
payload = self.cfg.to_dict()
|
||||
payload["meanflow_train"] = self.meanflow_settings.to_dict()
|
||||
with config_path.open("w", encoding="utf-8") as fout:
|
||||
yaml.safe_dump(
|
||||
payload,
|
||||
fout,
|
||||
sort_keys=False,
|
||||
allow_unicode=True,
|
||||
)
|
||||
|
||||
def _save_checkpoint(self, learning_rate: float) -> None:
|
||||
train_checkpoint.save_train_checkpoint(
|
||||
self.accelerator,
|
||||
self.model,
|
||||
self.optimizer,
|
||||
self.progress,
|
||||
self.cfg.train.output_dir,
|
||||
self.cfg.train.max_checkpoints_to_keep,
|
||||
self.train_loader.state_dict(),
|
||||
{
|
||||
"type": "transformers_cosine_with_warmup_meanflow",
|
||||
"global_step": int(self.progress.global_step),
|
||||
"base_lr": float(self.cfg.train.learning_rate),
|
||||
"current_lr": float(learning_rate),
|
||||
"warmup_steps": int(self.cfg.train.warmup_steps),
|
||||
"max_train_steps": int(self.max_train_steps),
|
||||
"meanflow": self.meanflow_settings.to_dict(),
|
||||
"state_dict": self.scheduler.state_dict(),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def parse_args(argv=None):
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Accelerate MeanFlow distillation entrypoint for dots.tts."
|
||||
)
|
||||
parser.add_argument("--config", default=app_config.DEFAULT_CONFIG_PATH)
|
||||
parser.add_argument(
|
||||
"--debug",
|
||||
action="store_true",
|
||||
help="Print training debug information.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--teacher-model-path",
|
||||
default=None,
|
||||
help=(
|
||||
"Frozen flow-matching teacher model path. Defaults to "
|
||||
"train.pretrained_model_path."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--teacher-steps", type=int, default=8)
|
||||
parser.add_argument(
|
||||
"--teacher-solver",
|
||||
choices=_ALLOWED_TEACHER_SOLVERS,
|
||||
default="euler",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cfg-distill-mode",
|
||||
choices=_ALLOWED_CFG_DISTILL_MODES,
|
||||
default="fused",
|
||||
)
|
||||
parser.add_argument("--distill-cfg-scale", type=float, default=1.2)
|
||||
parser.add_argument("--anchor-prob", type=float, default=0.5)
|
||||
parser.add_argument(
|
||||
"--anchor-target",
|
||||
choices=_ALLOWED_ANCHOR_TARGETS,
|
||||
default="formula",
|
||||
)
|
||||
parser.add_argument("--time-sampling-mean", type=float, default=-0.4)
|
||||
parser.add_argument("--time-sampling-std", type=float, default=1.0)
|
||||
parser.add_argument(
|
||||
"--train-all-parameters",
|
||||
action="store_true",
|
||||
help="Train all regular dots.tts parameters instead of only the DiT.",
|
||||
)
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def main(argv=None):
|
||||
args = parse_args(argv)
|
||||
settings = MeanFlowSettings(
|
||||
teacher_model_path=args.teacher_model_path,
|
||||
teacher_steps=args.teacher_steps,
|
||||
teacher_solver=args.teacher_solver,
|
||||
cfg_distill_mode=args.cfg_distill_mode,
|
||||
distill_cfg_scale=args.distill_cfg_scale,
|
||||
anchor_prob=args.anchor_prob,
|
||||
anchor_target=args.anchor_target,
|
||||
time_sampling_mean=args.time_sampling_mean,
|
||||
time_sampling_std=args.time_sampling_std,
|
||||
train_all_parameters=args.train_all_parameters,
|
||||
)
|
||||
return DotsTtsMeanFlowTrainingRun(
|
||||
app_config.load_config(args.config),
|
||||
meanflow_settings=settings,
|
||||
debug_enabled=args.debug,
|
||||
).run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user