Spaces:
Running on Zero
Running on Zero
AndrianBalanescu
feat(zerogpu): add Breeze TTS 2 experimental voice + remove image gen by policy
4a8de45 | from __future__ import annotations | |
| import torch | |
| from transformers.generation.logits_process import LogitsProcessor | |
| def mask_invalid_codec_token_logits( | |
| scores: torch.Tensor, | |
| *, | |
| codebook_size: int, | |
| token_vocab_size: int, | |
| ) -> torch.Tensor: | |
| """Mask reserved model tokens that are not valid codec codebook entries.""" | |
| if codebook_size <= 0 or token_vocab_size < codebook_size: | |
| raise ValueError( | |
| f"Invalid codec token range: codebook_size={codebook_size}, " | |
| f"token_vocab_size={token_vocab_size}" | |
| ) | |
| if token_vocab_size > scores.shape[-1]: | |
| raise ValueError( | |
| f"token_vocab_size={token_vocab_size} exceeds logits size {scores.shape[-1]}" | |
| ) | |
| if codebook_size < token_vocab_size: | |
| scores[..., codebook_size:token_vocab_size] = float("-inf") | |
| return scores | |
| class GeneratedTokenRepetitionPenaltyLogitsProcessor(LogitsProcessor): | |
| """Apply HF-style repetition penalty only to generated, in-vocab tokens.""" | |
| def __init__(self, penalty: float) -> None: | |
| if penalty <= 0: | |
| raise ValueError("penalty must be > 0") | |
| self.penalty = float(penalty) | |
| def __call__( | |
| self, input_ids: torch.LongTensor, scores: torch.FloatTensor | |
| ) -> torch.FloatTensor: | |
| # Breeze starts generation with 2D text prompt IDs, then replaces them | |
| # with 3D generated codec frames after the first decoding step. | |
| if input_ids.ndim == 2: | |
| return scores | |
| if input_ids.ndim != 3: | |
| raise ValueError( | |
| f"Expected 2D prompt IDs or 3D codec frames, got {input_ids.ndim}D" | |
| ) | |
| generated_ids = input_ids[..., 0] | |
| if self.penalty == 1.0 or generated_ids.numel() == 0: | |
| return scores | |
| vocab_size = scores.shape[-1] | |
| valid = (generated_ids >= 0) & (generated_ids < vocab_size) | |
| if not torch.any(valid): | |
| return scores | |
| # Invalid IDs use token 0 as a temporary gather/scatter target. Restore | |
| # token 0 explicitly afterward so those placeholders cannot affect it. | |
| safe_ids = generated_ids.masked_fill(~valid, 0) | |
| selected_scores = torch.gather(scores, 1, safe_ids) | |
| penalized_scores = torch.where( | |
| selected_scores < 0, | |
| selected_scores * self.penalty, | |
| selected_scores / self.penalty, | |
| ) | |
| processed = scores.clone() | |
| processed.scatter_(1, safe_ids, penalized_scores) | |
| original_zero = scores[:, 0] | |
| penalized_zero = torch.where( | |
| original_zero < 0, | |
| original_zero * self.penalty, | |
| original_zero / self.penalty, | |
| ) | |
| generated_zero = (valid & (generated_ids == 0)).any(dim=1) | |
| processed[:, 0] = torch.where(generated_zero, penalized_zero, original_zero) | |
| return processed | |