OzzyGT/YuE2-Modular

1# Adapted for diffusers from multimodal-art-projection/YuE at commit ef1936f2ee39fe8de486a0f47a481c95f8d4da87.2# Licensed under Apache-2.0; see LICENSE.3from __future__ import annotations4 5import time6 7import torch8 9from .protocol import ABC_END, CODEC_OFFSET, CODEC_SIZE, CONTEXT, EOD, MUSIC_END10 11 12class YuE2StaticKVCache:13    """Preallocated token cache for eager decoding; returns views of the filled prefix without copying history."""14 15    def __init__(self, num_layers, batch_size, num_kv_heads, max_seq_len, head_dim, dtype, device):16        self.num_layers = num_layers17        self.max_seq_len = max_seq_len18        self.seen_tokens = 019        shape = (batch_size, num_kv_heads, max_seq_len, head_dim)20        self.keys = [torch.zeros(shape, dtype=dtype, device=device) for _ in range(num_layers)]21        self.values = [torch.zeros(shape, dtype=dtype, device=device) for _ in range(num_layers)]22 23    def get_seq_length(self):24        return self.seen_tokens25 26    def update(self, key, value, layer_idx):27        start, end = self.seen_tokens, self.seen_tokens + key.shape[1]28        if end > self.max_seq_len:29            raise ValueError(f"KV cache capacity {self.max_seq_len} exceeded by {end}; generation was not shortened")30        self.keys[layer_idx][:, :, start:end] = key.transpose(1, 2)31        self.values[layer_idx][:, :, start:end] = value.transpose(1, 2)32        if layer_idx == self.num_layers - 1:33            self.seen_tokens = end34        return self.keys[layer_idx][:, :, :end].transpose(1, 2), self.values[layer_idx][:, :, :end].transpose(1, 2)35 36 37def synchronize(device):38    if device.type == "cuda":39        torch.cuda.synchronize(device)40    elif device.type == "mps":41        torch.mps.synchronize()42 43 44def window_penalty(logits, recent_ids, penalty):45    if penalty == 1.0 or len(recent_ids) == 0:46        return logits47    recent = torch.as_tensor(recent_ids, dtype=torch.long, device=logits.device).reshape(1, -1)48    freq = torch.zeros_like(logits)49    freq.scatter_add_(-1, recent, torch.ones_like(recent, dtype=logits.dtype))50    alpha = penalty**freq51    return torch.where(logits < 0, logits * alpha, logits / alpha)52 53 54def distribution(logits, sampling, history, step, phase, legacy_off=False):55    # Planning-off requests keep the release's BF16 logits and top-3 floor; other modes sample from FP32 logits.56    scores = logits.clone() if legacy_off else logits.float().clone()57    end = ABC_END if phase == "abc" else MUSIC_END58    allowed = torch.full_like(scores, float("-inf"))59    if phase == "abc":60        allowed[..., :EOD] = 061    else:62        allowed[..., CODEC_OFFSET : CODEC_OFFSET + CODEC_SIZE] = 063    allowed[..., end] = 064    scores = scores + allowed65    if step < sampling.min_tokens:66        scores[..., end] = -torch.inf67    scores = window_penalty(scores, history[-sampling.penalty_window :], sampling.repetition_penalty)68    if sampling.temperature == 0:69        return scores70    if sampling.temperature != 1:71        scores = scores / sampling.temperature72    threshold = scores.topk(min(sampling.top_k, scores.shape[-1])).values[..., -1, None]73    scores = scores.masked_fill(scores < threshold, -torch.inf)74    if sampling.top_p < 1:75        values, indices = scores.sort(descending=True)76        probabilities = values.softmax(-1)77        removed = probabilities.cumsum(-1) - probabilities > sampling.top_p78        removed[..., : 3 if legacy_off else 1] = False79        values = values.masked_fill(removed, -torch.inf)80        scores = values.scatter(-1, indices, values)81    return scores82 83 84@torch.inference_mode()85def generate_tokens(86    transformer,87    prefix,88    sampling,89    seed,90    phase,91    device,92    negative=None,93    combine_logits=None,94    legacy_off=False,95    cancelled=None,96    on_token=None,97    graph_decoder=None,98):99    """Sample one stage's tokens.100 101    With `negative`, `combine_logits(conditional, unconditional)` applies guidance. `graph_decoder` (the `GraphAR`102    class) decodes with CUDA graphs on CUDA devices; elsewhere decoding stays eager.103    """104    if len(prefix) + sampling.max_tokens > CONTEXT:105        raise ValueError("Prefix + requested generation budget exceeds 24576; no implicit truncation")106    if negative is not None and (combine_logits is None or len(negative) + sampling.max_tokens > CONTEXT):107        raise ValueError("Guidance needs `combine_logits` and a negative prefix that fits the context")108    if cancelled is not None and cancelled():109        raise InterruptedError("Cancelled before prefill")110    # Both stages reset the request seed, as the release does.111    rng_device = device if device.type in {"cpu", "cuda"} else torch.device("cpu")112    generator = torch.Generator(device=rng_device).manual_seed(seed)113    config = transformer.config114 115    def prefill(ids):116        cache = YuE2StaticKVCache(117            num_layers=config.num_layers,118            batch_size=1,119            num_kv_heads=config.num_key_value_heads,120            max_seq_len=len(ids) + sampling.max_tokens,121            head_dim=config.attention_head_dim,122            dtype=transformer.dtype,123            device=device,124        )125        logits = transformer(torch.tensor([ids], device=device), kv_cache=cache, logits_to_keep=1).logits126        return logits[:, -1, :], cache127 128    graph = None129    positive_cache = negative_cache = None130    synchronize(device)131    start = time.perf_counter()132    try:133        if graph_decoder is not None and device.type == "cuda":134            graph = graph_decoder(135                transformer, [prefix] if negative is None else [prefix, negative], sampling.max_tokens, device136            )137            logits = graph.prefill()138            conditional = logits[:1]139            unconditional = logits[1:] if negative is not None else None140        else:141            conditional, positive_cache = prefill(prefix)142            unconditional = None143            if negative is not None:144                unconditional, negative_cache = prefill(negative)145        synchronize(device)146        prefill_seconds = time.perf_counter() - start147        history, first, eos = [], None, False148        end = ABC_END if phase == "abc" else MUSIC_END149        for step in range(sampling.max_tokens):150            if cancelled is not None and cancelled():151                raise InterruptedError(f"Cancelled during {phase}")152            logits = conditional if negative is None else combine_logits(conditional, unconditional)153            scores = distribution(logits, sampling, history, step, phase, legacy_off)154            if sampling.temperature == 0:155                next_id = scores.argmax(-1, keepdim=True)156            else:157                probabilities = scores.softmax(-1)158                if device.type == "mps":159                    next_id = torch.multinomial(probabilities.cpu(), 1, generator=generator).to(device)160                else:161                    next_id = torch.multinomial(probabilities, 1, generator=generator)162            token = int(next_id.item())163            if first is None:164                first = time.perf_counter() - start165            if on_token is not None:166                on_token(phase, token)167            if token == end:168                eos = True169                break170            history.append(token)171            if step + 1 < sampling.max_tokens:172                if graph is not None:173                    branch_logits = graph.step(next_id)174                    conditional = branch_logits[:1]175                    unconditional = branch_logits[1:] if negative is not None else None176                else:177                    conditional = transformer(next_id, kv_cache=positive_cache, logits_to_keep=1).logits[:, -1, :]178                    if negative_cache is not None:179                        unconditional = transformer(next_id, kv_cache=negative_cache, logits_to_keep=1).logits[180                            :, -1, :181                        ]182        synchronize(device)183        seconds = time.perf_counter() - start184        count = len(history) + int(eos)185        timing = {186            "seconds": seconds,187            "prefill_seconds": prefill_seconds,188            "ttft_seconds": first,189            "output_tokens": count,190            "content_tokens": len(history),191            "output_tps": count / seconds,192            "prefix_tokens": len(prefix),193            "cfg_branches": 1 if negative is None else 2,194            "execution": "cuda_graph" if graph is not None else "eager",195        }196        return history, timing, not eos197    finally:198        if graph is not None:199            graph.close()200        positive_cache = negative_cache = None201