BlueWaveSemi45/DramaboxCPU

1"""Audio reference conditioning item for IC-LoRA voice cloning."""2 3import torch4 5from ltx_core.components.patchifiers import AudioPatchifier6from ltx_core.conditioning.item import ConditioningItem7from ltx_core.tools import AudioLatentTools8from ltx_core.types import AudioLatentShape, LatentState9 10 11class AudioConditionByReferenceLatent(ConditioningItem):12    """Conditions audio generation on a reference audio latent for voice cloning.13 14    Mirrors VideoConditionByReferenceLatent but for audio:15    - Patchifies reference latent [B, C, T, F] -> [B, ref_T, 128]16    - Computes 1D temporal positions via AudioPatchifier17    - Sets denoise_mask = 1.0 - strength (strength=1.0 -> mask=0 -> frozen)18    - Builds ASYMMETRIC attention mask: target->ref=1 (attend), ref->target=0 (read-only)19    - APPENDS ref tokens to END of latent sequence (IC-LoRA pattern)20    - Uses OVERLAPPING positions (same coordinate space) so RoPE doesn't21      decay target->ref attention. The asymmetric mask provides the structural22      signal that ref tokens are conditioning, not reconstruction targets.23 24    Args:25        latent: Reference audio latent [B, C, T, F] (pre-VAE-encoded).26        strength: Conditioning strength. 1.0 = full (ref kept clean),27            0.0 = none (ref fully denoised). Default 1.0.28    """29 30    def __init__(self, latent: torch.Tensor, strength: float = 1.0):31        self.latent = latent32        self.strength = strength33 34    def apply_to(35        self,36        latent_state: LatentState,37        latent_tools: AudioLatentTools,38    ) -> LatentState:39        """Append reference audio tokens with positions and attention mask."""40        tokens = latent_tools.patchifier.patchify(self.latent)41 42        # Compute positions for the reference audio — small offset (0.5s) from43        # target start to avoid exact t=0 overlap (which causes ref content to44        # bleed into target start), while keeping RoPE decay minimal.45        # 0.5s / max_pos(20s) = 0.025 fractional — negligible RoPE decay.46        ref_shape = AudioLatentShape(47            batch=self.latent.shape[0],48            channels=self.latent.shape[1],49            frames=self.latent.shape[2],50            mel_bins=self.latent.shape[3],51        )52        positions = latent_tools.patchifier.get_patch_grid_bounds(53            output_shape=ref_shape,54            device=self.latent.device,55        )56        # Small offset to prevent t=0 position collision between target and ref57        positions = positions + 0.558 59        # Denoise mask: 0 for frozen (strength=1.0), 1 for fully denoised (strength=0.0)60        denoise_mask = torch.full(61            size=(*tokens.shape[:2], 1),62            fill_value=1.0 - self.strength,63            device=self.latent.device,64            dtype=torch.float32,65        )66 67        # Build ASYMMETRIC attention mask manually.68        # Structure:69        #              target (N)    ref (M)70        #            ┌────────────┬──────────┐71        #   target   │    1.0     │   1.0    │  target attends to everything72        #    (N)     │            │          │73        #            ├────────────┼──────────┤74        #    ref     │    0.0     │   1.0    │  ref only attends to itself75        #    (M)     │            │          │76        #            └────────────┴──────────┘77        #78        # This makes reference tokens "read-only conditioning":79        # - Target tokens freely attend to ref (voice cloning signal)80        # - Ref tokens don't attend to noisy target (stays clean/stable)81        batch_size = tokens.shape[0]82        num_target = latent_state.latent.shape[1]83        num_ref = tokens.shape[1]84        total = num_target + num_ref85 86        # Use float32 for the [0,1] mask — _prepare_self_attention_mask converts87        # to log-space bias in the model's compute dtype before it reaches attention.88        mask = torch.zeros(89            (batch_size, total, total),90            device=self.latent.device,91            dtype=torch.float32,92        )93 94        # Incorporate existing mask if present, otherwise full attention for target95        if latent_state.attention_mask is not None:96            mask[:, :num_target, :num_target] = latent_state.attention_mask97        else:98            mask[:, :num_target, :num_target] = 1.099 100        # Target -> ref: FULL attention (target can read reference voice)101        mask[:, :num_target, num_target:] = 1.0102 103        # Ref -> target: BLOCKED (ref is read-only, doesn't see noisy target)104        # mask[:, num_target:, :num_target] remains 0.0105 106        # Ref -> ref: full self-attention within reference107        mask[:, num_target:, num_target:] = 1.0108 109        return LatentState(110            latent=torch.cat([latent_state.latent, tokens], dim=1),111            denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),112            positions=torch.cat([latent_state.positions, positions], dim=2),113            clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),114            attention_mask=mask,115        )116