diff --git a/extensions_built_in/audio_models/yue2/src/tokenizer.py b/extensions_built_in/audio_models/yue2/src/tokenizer.py index 8d7209a..064d44f 100644 --- a/extensions_built_in/audio_models/yue2/src/tokenizer.py +++ b/extensions_built_in/audio_models/yue2/src/tokenizer.py @@ -78,8 +78,9 @@ def _rebuild_rotary(model: nn.Module) -> int: class SemanticTokenizer(nn.Module): """Waveform -> per-frame codec ids (0..32767) at 25 Hz.""" - def __init__(self, head_path: str, mert_repo: str = MERT_REPO): + def __init__(self, head_path: str, mert_repo: str = MERT_REPO, debug: bool = False): super().__init__() + self.debug = debug from transformers import AutoFeatureExtractor, AutoModel self.processor = AutoFeatureExtractor.from_pretrained(mert_repo, trust_remote_code=True) @@ -149,6 +150,6 @@ class SemanticTokenizer(nn.Module): mono = torchaudio.functional.resample(mono, sample_rate, MERT_SAMPLE_RATE) feats = self.mert_features(mono.to(self.device)) out = self.tokens_from_features(feats) - if out.unique().numel() < 8: + if self.debug and out.unique().numel() < 8: print(f"YuE2 tokenizer: degenerate token stream (unique {out.unique().numel()}); MERT features nan {torch.isnan(feats).sum().item()}") return out diff --git a/extensions_built_in/audio_models/yue2/yue2_model.py b/extensions_built_in/audio_models/yue2/yue2_model.py index 07280e7..1423f55 100644 --- a/extensions_built_in/audio_models/yue2/yue2_model.py +++ b/extensions_built_in/audio_models/yue2/yue2_model.py @@ -33,6 +33,7 @@ from toolkit.basic import flush from toolkit.config_modules import GenerateImageConfig from toolkit.dto import DTO from toolkit.models.v2.resolver import resolve_named_file +from toolkit.paths import MODELS_PATH from toolkit.print import print_acc from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler @@ -43,6 +44,24 @@ from .src.tokenizer import HEAD_FILE, HEAD_REPO, NAR_LORA_FILE, SemanticTokenize from .src.vae import SAMPLE_RATE, YuE2VAE DEFAULT_CHECKPOINT = "Comfy-Org/YuE2/checkpoints/yue2_3b_bf16.safetensors" +# community tokenizer files have no ComfyUI folder; they live under /ai_toolkit/ +AI_TOOLKIT_DIR = "ai_toolkit" + + +def resolve_ai_toolkit_file(path: str, component: str) -> str: + """Local file, or 'org/repo/file' downloaded into /ai_toolkit/.""" + if os.path.exists(path): + return path + splits = path.split("/") + if len(splits) < 3: + raise ValueError(f"Invalid {component} path: {path}. Must be a local file or 'org/repo/filename'.") + local_dir = os.path.join(MODELS_PATH, AI_TOOLKIT_DIR) + candidate = os.path.join(local_dir, splits[-1]) + if os.path.exists(candidate): + return candidate + import huggingface_hub + + return huggingface_hub.hf_hub_download(repo_id="/".join(splits[:2]), filename="/".join(splits[2:]), local_dir=local_dir) scheduler_config = { "num_train_timesteps": 1000, @@ -141,15 +160,24 @@ class YuE2AudioModel(BaseAudioModel): self.ar_max_tokens = int(kw.get("ar_max_tokens", 0)) # optional separate learning rate for the AR expert's LoRA (Adam ignores loss scale) self.ar_lr_multiplier = float(kw.get("ar_lr_multiplier", 1.0)) + # trust region: KL(base || lora) on the AR next-token distributions, base = LoRA switched off + self.ar_kl_weight = float(kw.get("ar_kl_weight", 0.0)) self._loss_log_every = int(kw.get("loss_log_every", 25)) + # extra per-readout AR diagnostics (targets, unique tokens, p(END), p(target)) + self.debug = bool(kw.get("debug", False)) self._loss_log_step = 0 self.merge_nar_lora = bool(kw.get("merge_nar_lora", False)) self.nar_lora_path = kw.get("nar_lora_path", f"{HEAD_REPO}/{NAR_LORA_FILE}") self.semantic_head_path = kw.get("semantic_head_path", f"{HEAD_REPO}/{HEAD_FILE}") self.sample_max_seconds = float(kw.get("sample_max_seconds", 120.0)) self.sample_ar_temperature = float(kw.get("sample_ar_temperature", 1.0)) + # reference sampler value; 1.0 lets a memorized song replay (the penalty knocks it off path) + self.sample_ar_repetition_penalty = float(kw.get("sample_ar_repetition_penalty", 1.2)) self.semantic_tokenizer: Optional[SemanticTokenizer] = None self._pending_aux_loss = None + self._pending_ar_ce = None + self._pending_ar_kl = None + self.additional_loss_logs = {} @staticmethod def get_train_scheduler(): @@ -175,7 +203,7 @@ class YuE2AudioModel(BaseAudioModel): self.model = YuE2Model.load_from_state_dict(sd, dtype=dtype) if self.merge_nar_lora: - lora_path = resolve_named_file(self.nar_lora_path, component="yue2 nar lora") + lora_path = resolve_ai_toolkit_file(self.nar_lora_path, component="yue2 nar lora") self.print_and_status_update(f"Merging NAR adapter {os.path.basename(lora_path)}") merge_nar_lora(self.model, lora_path) vae_sd = {k[len("vae.") :]: v for k, v in sd.items() if k.startswith("vae.")} @@ -196,10 +224,10 @@ class YuE2AudioModel(BaseAudioModel): def _get_semantic_tokenizer(self) -> SemanticTokenizer: if self.semantic_tokenizer is None: - head_path = resolve_named_file(self.semantic_head_path, component="yue2 semantic head") + head_path = resolve_ai_toolkit_file(self.semantic_head_path, component="yue2 semantic head") # log only: this can run mid-training and a status update would replace "Training" in the UI print_acc("Loading YuE2 semantic tokenizer (MERT-v2-FullSong + head)") - self.semantic_tokenizer = SemanticTokenizer(head_path) + self.semantic_tokenizer = SemanticTokenizer(head_path, debug=self.debug) return self.semantic_tokenizer def get_transformer_block_names(self) -> Optional[List[str]]: @@ -262,17 +290,25 @@ class YuE2AudioModel(BaseAudioModel): start, end = getattr(batch, "yue2_window", (0, noise.shape[1])) return (noise - batch.latents)[:, start:end].detach() - def scale_loss(self, loss): + def get_additional_loss(self, pred: torch.Tensor, target: torch.Tensor): + """AR next-token CE from the last forward, added to the flow loss by the trainer and + logged as its own term (`loss/ar_ce`, unweighted).""" aux = self._pending_aux_loss self._pending_aux_loss = None if aux is None: - return loss + self.additional_loss_logs = {} + return None + ce = self._pending_ar_ce.detach().float().item() + self.additional_loss_logs = {"loss/ar_ce": ce} + kl = self._pending_ar_kl + if kl is not None: + self.additional_loss_logs["loss/ar_kl"] = kl.detach().float().item() self._loss_log_step += 1 - if self._loss_log_every > 0 and self._loss_log_step % self._loss_log_every == 0: - # the trainer's bar shows the sum; the two terms move on different scales + if self.debug and self._loss_log_every > 0 and self._loss_log_step % self._loss_log_every == 0: info = getattr(self, "_pending_aux_info", "") - tqdm.write(f"yue2 step {self._loss_log_step}: flow {loss.detach().float().mean().item():.4f} ar_ce {aux.detach().float().item() / max(self.ar_loss_weight, 1e-8):.4f} {info}") - return loss + aux + kl_txt = f" ar_kl {kl.detach().float().item():.4f}" if kl is not None else "" + tqdm.write(f"yue2 step {self._loss_log_step}: ar_ce {ce:.4f}{kl_txt}" + (f" {info}" if info else "")) + return aux @staticmethod def _train_prefix(prefix_embeds: torch.Tensor, prefix_mask: torch.Tensor) -> torch.Tensor: @@ -283,7 +319,7 @@ class YuE2AudioModel(BaseAudioModel): def _ar_info(self, logits: torch.Tensor, ids: torch.Tensor, prefix: torch.Tensor, total: int, whole_song: bool) -> str: """Readout diagnostics: what the AR loss was computed over.""" - if self._loss_log_every <= 0 or (self._loss_log_step + 1) % self._loss_log_every != 0: + if not self.debug or self._loss_log_every <= 0 or (self._loss_log_step + 1) % self._loss_log_every != 0: return "" with torch.no_grad(): p = logits.softmax(-1) @@ -291,6 +327,27 @@ class YuE2AudioModel(BaseAudioModel): return (f"[ar targets {ids.shape[0]} unique {ids.unique().numel()} (song {total}, whole {whole_song}) " f"prefix {prefix.shape[0]} p(END|prefix) {p[0, MUSIC_END].item():.3f} p(target) median {p_target.median().item():.3f}]") + def _ar_kl(self, ar_embeds: torch.Tensor, logits: torch.Tensor): + """KL(base || lora) per position, mean. Base pass = same prefill with the LoRA inactive, no grad.""" + net = getattr(self, "_network", None) + n = logits.shape[0] + was_active = getattr(net, "is_active", None) + if net is not None: + net.is_active = False + try: + with torch.no_grad(): + _, base_hidden = self.model.ar.prefill(ar_embeds, return_hidden=True) + base_hidden = base_hidden[0, -n - 1 : -1] + finally: + if net is not None: + net.is_active = was_active + kl = logits.new_zeros(()) + for s0 in range(0, n, 1024): # chunked: two fp32 [n, vocab] log-softmaxes at once would be ~10 GB + with torch.no_grad(): + base_logp = torch.log_softmax(self.model.ar.model.lm_head(base_hidden[s0 : s0 + 1024]).float(), -1) + kl = kl + torch.nn.functional.kl_div(torch.log_softmax(logits[s0 : s0 + 1024], -1), base_logp, log_target=True, reduction="sum") + return kl / n + def _ar_inputs(self, prefix: torch.Tensor, tokens: torch.Tensor, end_token: bool = True): """One item: prefix (unpadded [L, H]) + codec tokens [+ MUSIC_END] -> embeds [1, L', H], ids [n].""" ids = tokens.long() + CODEC_OFFSET @@ -356,7 +413,7 @@ class YuE2AudioModel(BaseAudioModel): prefix_embeds = text_embeddings.text_embeds.to(device, self.torch_dtype) prefix_mask = text_embeddings.attention_mask.to(device) train_ar = self.ar_loss_weight > 0 and torch.is_grad_enabled() - preds, lm_losses = [], [] + preds, lm_losses, kl_losses = [], [], [] for i in range(latent_model_input.shape[0]): song = tokens_all[i].to(device) total = song.shape[0] @@ -372,6 +429,8 @@ class YuE2AudioModel(BaseAudioModel): _, hidden = model.ar.prefill(ar_embeds, return_hidden=True) logits = model.ar.model.lm_head(hidden[0, -ar_ids.shape[0] - 1 : -1]).float() lm_losses.append(torch.nn.functional.cross_entropy(logits, ar_ids)) + if self.ar_kl_weight > 0: + kl_losses.append(self._ar_kl(ar_embeds, logits)) self._pending_aux_info = self._ar_info(logits, ar_ids, prefix, total, whole_song) with torch.no_grad(): cache, _ = model.ar.prefill(embeds) @@ -380,13 +439,19 @@ class YuE2AudioModel(BaseAudioModel): if train_ar: logits = model.ar.model.lm_head(hidden[0, -ids.shape[0] - 1 : -1]).float() lm_losses.append(torch.nn.functional.cross_entropy(logits, ids)) + if self.ar_kl_weight > 0: + kl_losses.append(self._ar_kl(embeds, logits)) self._pending_aux_info = self._ar_info(logits, ids, prefix, total, whole_song) # the flow loss must not train the AR through its KV cache: only next-token CE shapes the AR cache = [(k.detach(), v.detach()) for k, v in cache] pred = model.nar(latent_model_input[i : i + 1].to(device, self.torch_dtype), t[i : i + 1], cache, embeds.shape[1]) preds.append(pred) if lm_losses: - self._pending_aux_loss = torch.stack(lm_losses).mean() * self.ar_loss_weight + self._pending_ar_ce = torch.stack(lm_losses).mean() + self._pending_ar_kl = torch.stack(kl_losses).mean() if kl_losses else None + self._pending_aux_loss = self._pending_ar_ce * self.ar_loss_weight + if self._pending_ar_kl is not None: + self._pending_aux_loss = self._pending_aux_loss + self._pending_ar_kl * self.ar_kl_weight return torch.cat(preds, 0) # ------------------------------------------------------------------ @@ -395,6 +460,7 @@ class YuE2AudioModel(BaseAudioModel): def get_generation_pipeline(self): return self.pipeline + def generate_single_audio(self, pipeline: YuE2Pipeline, gen_config: GenerateImageConfig, conditional_embeds: PromptEmbeds, unconditional_embeds, generator, extra): if self.model.device == torch.device("cpu"): self.model.to(self.device_torch) @@ -417,7 +483,8 @@ class YuE2AudioModel(BaseAudioModel): try: codec = pipeline.generate_codec_tokens( - prefix, max_tokens=max_tokens, seed=gen_config.seed, temperature=self.sample_ar_temperature, progress=token_progress, + prefix, max_tokens=max_tokens, seed=gen_config.seed, temperature=self.sample_ar_temperature, + repetition_penalty=self.sample_ar_repetition_penalty, progress=token_progress, ) finally: token_bar.close() diff --git a/ui/src/app/jobs/new/options.tsx b/ui/src/app/jobs/new/options.tsx index 0d6a7e3..b45bf76 100644 --- a/ui/src/app/jobs/new/options.tsx +++ b/ui/src/app/jobs/new/options.tsx @@ -2,7 +2,12 @@ import React from 'react'; import Link from 'next/link'; import { GroupedSelectOption, SelectOption, JobConfig, ConfigDoc } from '@/types'; import { defaultSliderConfig } from './jobConfig'; -import { defaultAudioSampleConfig, defaultSampleConfig, defaultIdeogramSamplesConfig, defaultYue2SampleConfig } from '@/helpers/defaultSamples'; +import { + defaultAudioSampleConfig, + defaultSampleConfig, + defaultIdeogramSamplesConfig, + defaultYue2SampleConfig, +} from '@/helpers/defaultSamples'; type Control = 'depth' | 'line' | 'pose' | 'inpaint'; @@ -1394,9 +1399,61 @@ export const modelArchs: ModelArch[] = [ 'config.process[0].datasets[x].cache_latents_to_disk': [true, true], // blank captions break lyric following; the AR must always see the prefix 'config.process[0].datasets[x].caption_dropout_rate': [0, 0.05], + 'config.process[0].model.model_kwargs': [{ sample_ar_repetition_penalty: 1.2, ar_kl_weight: 0.2 }, {}], }, // native YuE2 prompt: style text, a [Lyrics] line, the lyrics hasMultiLinePrompts: true, + modelNotes: ( +
+

+ Experimental. The AR (composition) model memorizes quickly and does not work well on small datasets. A handful + of songs is enough for it to learn the exact token sequence of each song; after that it stops generalizing and + free-running samples drift away from the training material. Expect to need a large, varied dataset for the AR + side to learn a style rather than the songs themselves. +

+

+ YuE2 is two experts on one backbone. The AR expert reads the style line and lyrics and writes the song + as a sequence of semantic codec tokens (25 per second). The NAR expert then renders those tokens into + audio latents with flow matching. Training a LoRA here trains both: next-token loss on the AR over the whole + song from its start, flow loss on the NAR over a random window. +

+

+ The AR loss is the one to watch (loss/ar_ce). It starts near 5 and, on a small dataset, falls + toward 0 within a few hundred steps, which is memorization. ar_kl_weight in model kwargs anchors + the AR to the base model so it cannot collapse onto the training songs; ar_lr_multiplier and a + smaller AR rank slow it further. Style comes mostly from the NAR, lyric following from the AR. Do not use + caption dropout: a blank prompt breaks lyric following. +

+

+ The official audio-to-token encoder is unreleased. Training uses the community tokenizer by Kytra ( + + @sin_ceriously + + ), a MERT-v2-FullSong head that maps real audio to YuE2 codec tokens:{' '} + + Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4 + + . It is downloaded on first use. Its tokens are an approximation of the model's own, so rendered samples will + not be bit-faithful to the training audio even when the AR replays a song exactly. Though it is extremely + close. +

+

+ Prompt format: a style line, then [Lyrics], then the lyrics with bracketed section headers such + as [Verse 1] and [Chorus]. The Qwen3-Omni captioner has a YuE2 preset that writes + captions in this layout. Sample length is set by the Duration field in the sample section. +

+
+ ), disableSections: ['network.conv'], additionalSections: ['model.low_vram', 'sample.duration'], }, @@ -1832,7 +1889,6 @@ export const jobTypeOptions: JobTypeOption[] = [ }, ]; - const MODEL_PREFIX = 'config.process[0].model.'; const SAMPLE_PREFIX = 'config.process[0].sample'; // training-only model settings: the training adapter (e.g. Z-Image Turbo's @@ -1869,7 +1925,8 @@ export const getGenerateDefaults = (arch: ModelArch): GenerateDefaults => { if (value === undefined || value === null || value === '') continue; if (field === 'sample_steps') sample.num_inference_steps = value; else if (field === 'neg') sample.negative_prompt = value; - else if (['width', 'height', 'guidance_scale', 'num_frames', 'fps', 'duration'].includes(field)) sample[field] = value; + else if (['width', 'height', 'guidance_scale', 'num_frames', 'fps', 'duration'].includes(field)) + sample[field] = value; } } // the engine defaults to convrot8; the training default qtype is the @@ -1878,7 +1935,8 @@ export const getGenerateDefaults = (arch: ModelArch): GenerateDefaults => { if (model.quantize_te && (!model.qtype_te || model.qtype_te === 'qfloat8')) model.qtype_te = 'convrot8'; const sections = arch.additionalSections || []; const gen = arch.generate || {}; - const modality: GenerateModality = gen.modality || (arch.group === 'audio' ? 'audio' : arch.isVideoModel ? 'video' : 'image'); + const modality: GenerateModality = + gen.modality || (arch.group === 'audio' ? 'audio' : arch.isVideoModel ? 'video' : 'image'); if (modality !== 'video') { delete sample.num_frames; delete sample.fps; @@ -1890,7 +1948,8 @@ export const getGenerateDefaults = (arch: ModelArch): GenerateDefaults => { modality, model: { ...model, ...(gen.model || {}) }, sample: { ...sample, ...(gen.sample || {}) }, - needsControlImage: gen.needsControlImage ?? (sections.includes('sample.ctrl_img') || sections.includes('sample.multi_ctrl_imgs')), + needsControlImage: + gen.needsControlImage ?? (sections.includes('sample.ctrl_img') || sections.includes('sample.multi_ctrl_imgs')), sizeLocked: gen.sizeLocked ?? false, sampleTags: arch.sampleTags, };