# Adapted for diffusers from multimodal-art-projection/YuE at commit ef1936f2ee39fe8de486a0f47a481c95f8d4da87. # Licensed under Apache-2.0; see LICENSE. from __future__ import annotations import time import torch from .protocol import ABC_END, CODEC_OFFSET, CODEC_SIZE, CONTEXT, EOD, MUSIC_END class YuE2StaticKVCache: """Preallocated token cache for eager decoding; returns views of the filled prefix without copying history.""" def __init__(self, num_layers, batch_size, num_kv_heads, max_seq_len, head_dim, dtype, device): self.num_layers = num_layers self.max_seq_len = max_seq_len self.seen_tokens = 0 shape = (batch_size, num_kv_heads, max_seq_len, head_dim) self.keys = [torch.zeros(shape, dtype=dtype, device=device) for _ in range(num_layers)] self.values = [torch.zeros(shape, dtype=dtype, device=device) for _ in range(num_layers)] def get_seq_length(self): return self.seen_tokens def update(self, key, value, layer_idx): start, end = self.seen_tokens, self.seen_tokens + key.shape[1] if end > self.max_seq_len: raise ValueError(f"KV cache capacity {self.max_seq_len} exceeded by {end}; generation was not shortened") self.keys[layer_idx][:, :, start:end] = key.transpose(1, 2) self.values[layer_idx][:, :, start:end] = value.transpose(1, 2) if layer_idx == self.num_layers - 1: self.seen_tokens = end return self.keys[layer_idx][:, :, :end].transpose(1, 2), self.values[layer_idx][:, :, :end].transpose(1, 2) def synchronize(device): if device.type == "cuda": torch.cuda.synchronize(device) elif device.type == "mps": torch.mps.synchronize() def window_penalty(logits, recent_ids, penalty): if penalty == 1.0 or len(recent_ids) == 0: return logits recent = torch.as_tensor(recent_ids, dtype=torch.long, device=logits.device).reshape(1, -1) freq = torch.zeros_like(logits) freq.scatter_add_(-1, recent, torch.ones_like(recent, dtype=logits.dtype)) alpha = penalty**freq return torch.where(logits < 0, logits * alpha, logits / alpha) def distribution(logits, sampling, history, step, phase, legacy_off=False): # Planning-off requests keep the release's BF16 logits and top-3 floor; other modes sample from FP32 logits. scores = logits.clone() if legacy_off else logits.float().clone() end = ABC_END if phase == "abc" else MUSIC_END allowed = torch.full_like(scores, float("-inf")) if phase == "abc": allowed[..., :EOD] = 0 else: allowed[..., CODEC_OFFSET : CODEC_OFFSET + CODEC_SIZE] = 0 allowed[..., end] = 0 scores = scores + allowed if step < sampling.min_tokens: scores[..., end] = -torch.inf scores = window_penalty(scores, history[-sampling.penalty_window :], sampling.repetition_penalty) if sampling.temperature == 0: return scores if sampling.temperature != 1: scores = scores / sampling.temperature threshold = scores.topk(min(sampling.top_k, scores.shape[-1])).values[..., -1, None] scores = scores.masked_fill(scores < threshold, -torch.inf) if sampling.top_p < 1: values, indices = scores.sort(descending=True) probabilities = values.softmax(-1) removed = probabilities.cumsum(-1) - probabilities > sampling.top_p removed[..., : 3 if legacy_off else 1] = False values = values.masked_fill(removed, -torch.inf) scores = values.scatter(-1, indices, values) return scores @torch.inference_mode() def generate_tokens( transformer, prefix, sampling, seed, phase, device, negative=None, combine_logits=None, legacy_off=False, cancelled=None, on_token=None, graph_decoder=None, ): """Sample one stage's tokens. With `negative`, `combine_logits(conditional, unconditional)` applies guidance. `graph_decoder` (the `GraphAR` class) decodes with CUDA graphs on CUDA devices; elsewhere decoding stays eager. """ if len(prefix) + sampling.max_tokens > CONTEXT: raise ValueError("Prefix + requested generation budget exceeds 24576; no implicit truncation") if negative is not None and (combine_logits is None or len(negative) + sampling.max_tokens > CONTEXT): raise ValueError("Guidance needs `combine_logits` and a negative prefix that fits the context") if cancelled is not None and cancelled(): raise InterruptedError("Cancelled before prefill") # Both stages reset the request seed, as the release does. rng_device = device if device.type in {"cpu", "cuda"} else torch.device("cpu") generator = torch.Generator(device=rng_device).manual_seed(seed) config = transformer.config def prefill(ids): cache = YuE2StaticKVCache( num_layers=config.num_layers, batch_size=1, num_kv_heads=config.num_key_value_heads, max_seq_len=len(ids) + sampling.max_tokens, head_dim=config.attention_head_dim, dtype=transformer.dtype, device=device, ) logits = transformer(torch.tensor([ids], device=device), kv_cache=cache, logits_to_keep=1).logits return logits[:, -1, :], cache graph = None positive_cache = negative_cache = None synchronize(device) start = time.perf_counter() try: if graph_decoder is not None and device.type == "cuda": graph = graph_decoder( transformer, [prefix] if negative is None else [prefix, negative], sampling.max_tokens, device ) logits = graph.prefill() conditional = logits[:1] unconditional = logits[1:] if negative is not None else None else: conditional, positive_cache = prefill(prefix) unconditional = None if negative is not None: unconditional, negative_cache = prefill(negative) synchronize(device) prefill_seconds = time.perf_counter() - start history, first, eos = [], None, False end = ABC_END if phase == "abc" else MUSIC_END for step in range(sampling.max_tokens): if cancelled is not None and cancelled(): raise InterruptedError(f"Cancelled during {phase}") logits = conditional if negative is None else combine_logits(conditional, unconditional) scores = distribution(logits, sampling, history, step, phase, legacy_off) if sampling.temperature == 0: next_id = scores.argmax(-1, keepdim=True) else: probabilities = scores.softmax(-1) if device.type == "mps": next_id = torch.multinomial(probabilities.cpu(), 1, generator=generator).to(device) else: next_id = torch.multinomial(probabilities, 1, generator=generator) token = int(next_id.item()) if first is None: first = time.perf_counter() - start if on_token is not None: on_token(phase, token) if token == end: eos = True break history.append(token) if step + 1 < sampling.max_tokens: if graph is not None: branch_logits = graph.step(next_id) conditional = branch_logits[:1] unconditional = branch_logits[1:] if negative is not None else None else: conditional = transformer(next_id, kv_cache=positive_cache, logits_to_keep=1).logits[:, -1, :] if negative_cache is not None: unconditional = transformer(next_id, kv_cache=negative_cache, logits_to_keep=1).logits[ :, -1, : ] synchronize(device) seconds = time.perf_counter() - start count = len(history) + int(eos) timing = { "seconds": seconds, "prefill_seconds": prefill_seconds, "ttft_seconds": first, "output_tokens": count, "content_tokens": len(history), "output_tps": count / seconds, "prefix_tokens": len(prefix), "cfg_branches": 1 if negative is None else 2, "execution": "cuda_graph" if graph is not None else "eager", } return history, timing, not eos finally: if graph is not None: graph.close() positive_cache = negative_cache = None