| 1 | from dataclasses import dataclass, replace |
| 2 | from typing import Protocol |
| 3 | |
| 4 | import torch |
| 5 | from torch._prims_common import DeviceLikeType |
| 6 | |
| 7 | from ltx_core.components.patchifiers import ( |
| 8 | AudioLatentShape, |
| 9 | AudioPatchifier, |
| 10 | VideoLatentPatchifier, |
| 11 | VideoLatentShape, |
| 12 | get_pixel_coords, |
| 13 | ) |
| 14 | from ltx_core.components.protocols import Patchifier |
| 15 | from ltx_core.types import LatentState, SpatioTemporalScaleFactors |
| 16 | |
| 17 | DEFAULT_SCALE_FACTORS = SpatioTemporalScaleFactors.default() |
| 18 | |
| 19 | |
| 20 | class LatentTools(Protocol): |
| 21 | """ |
| 22 | Tools for building latent states. |
| 23 | """ |
| 24 | |
| 25 | patchifier: Patchifier |
| 26 | target_shape: VideoLatentShape | AudioLatentShape |
| 27 | |
| 28 | def create_initial_state( |
| 29 | self, |
| 30 | device: DeviceLikeType, |
| 31 | dtype: torch.dtype, |
| 32 | initial_latent: torch.Tensor | None = None, |
| 33 | ) -> LatentState: |
| 34 | """ |
| 35 | Create an initial latent state. If initial_latent is provided, it will be used to create the latent state. |
| 36 | """ |
| 37 | ... |
| 38 | |
| 39 | def patchify(self, latent_state: LatentState) -> LatentState: |
| 40 | """ |
| 41 | Patchify the latent state. |
| 42 | """ |
| 43 | if latent_state.latent.shape != self.target_shape.to_torch_shape(): |
| 44 | raise ValueError( |
| 45 | f"Latent state has shape {latent_state.latent.shape}, expected shape is " |
| 46 | f"{self.target_shape.to_torch_shape()}" |
| 47 | ) |
| 48 | latent_state = latent_state.clone() |
| 49 | latent = self.patchifier.patchify(latent_state.latent) |
| 50 | clean_latent = self.patchifier.patchify(latent_state.clean_latent) |
| 51 | denoise_mask = self.patchifier.patchify(latent_state.denoise_mask) |
| 52 | return replace(latent_state, latent=latent, denoise_mask=denoise_mask, clean_latent=clean_latent) |
| 53 | |
| 54 | def unpatchify(self, latent_state: LatentState) -> LatentState: |
| 55 | """ |
| 56 | Unpatchify the latent state. |
| 57 | """ |
| 58 | latent_state = latent_state.clone() |
| 59 | latent = self.patchifier.unpatchify(latent_state.latent, output_shape=self.target_shape) |
| 60 | clean_latent = self.patchifier.unpatchify(latent_state.clean_latent, output_shape=self.target_shape) |
| 61 | denoise_mask = self.patchifier.unpatchify( |
| 62 | latent_state.denoise_mask, output_shape=self.target_shape.mask_shape() |
| 63 | ) |
| 64 | return replace(latent_state, latent=latent, denoise_mask=denoise_mask, clean_latent=clean_latent) |
| 65 | |
| 66 | def clear_conditioning(self, latent_state: LatentState) -> LatentState: |
| 67 | """ |
| 68 | Clear the conditioning from the latent state. This method removes extra tokens from the end of the latent. |
| 69 | Therefore, conditioning items should add extra tokens ONLY to the end of the latent. |
| 70 | """ |
| 71 | latent_state = latent_state.clone() |
| 72 | |
| 73 | num_tokens = self.patchifier.get_token_count(self.target_shape) |
| 74 | latent = latent_state.latent[:, :num_tokens] |
| 75 | clean_latent = latent_state.clean_latent[:, :num_tokens] |
| 76 | denoise_mask = torch.ones_like(latent_state.denoise_mask)[:, :num_tokens] |
| 77 | positions = latent_state.positions[:, :, :num_tokens] |
| 78 | |
| 79 | return LatentState( |
| 80 | latent=latent, |
| 81 | denoise_mask=denoise_mask, |
| 82 | positions=positions, |
| 83 | clean_latent=clean_latent, |
| 84 | attention_mask=None, |
| 85 | ) |
| 86 | |
| 87 | |
| 88 | @dataclass(frozen=True) |
| 89 | class VideoLatentTools(LatentTools): |
| 90 | """ |
| 91 | Tools for building video latent states. |
| 92 | """ |
| 93 | |
| 94 | patchifier: VideoLatentPatchifier |
| 95 | target_shape: VideoLatentShape |
| 96 | fps: float |
| 97 | scale_factors: SpatioTemporalScaleFactors = DEFAULT_SCALE_FACTORS |
| 98 | causal_fix: bool = True |
| 99 | |
| 100 | def create_initial_state( |
| 101 | self, |
| 102 | device: DeviceLikeType, |
| 103 | dtype: torch.dtype, |
| 104 | initial_latent: torch.Tensor | None = None, |
| 105 | ) -> LatentState: |
| 106 | if initial_latent is not None: |
| 107 | assert initial_latent.shape == self.target_shape.to_torch_shape(), ( |
| 108 | f"Latent shape {initial_latent.shape} does not match target shape {self.target_shape.to_torch_shape()}" |
| 109 | ) |
| 110 | else: |
| 111 | initial_latent = torch.zeros( |
| 112 | *self.target_shape.to_torch_shape(), |
| 113 | device=device, |
| 114 | dtype=dtype, |
| 115 | ) |
| 116 | |
| 117 | clean_latent = initial_latent.clone() |
| 118 | |
| 119 | denoise_mask = torch.ones( |
| 120 | *self.target_shape.mask_shape().to_torch_shape(), |
| 121 | device=device, |
| 122 | dtype=torch.float32, |
| 123 | ) |
| 124 | |
| 125 | latent_coords = self.patchifier.get_patch_grid_bounds( |
| 126 | output_shape=self.target_shape, |
| 127 | device=device, |
| 128 | ) |
| 129 | |
| 130 | positions = get_pixel_coords( |
| 131 | latent_coords=latent_coords, |
| 132 | scale_factors=self.scale_factors, |
| 133 | causal_fix=self.causal_fix, |
| 134 | ).float() |
| 135 | positions[:, 0, ...] = positions[:, 0, ...] / self.fps |
| 136 | |
| 137 | return self.patchify( |
| 138 | LatentState( |
| 139 | latent=initial_latent, |
| 140 | denoise_mask=denoise_mask, |
| 141 | positions=positions.to(dtype), |
| 142 | clean_latent=clean_latent, |
| 143 | ) |
| 144 | ) |
| 145 | |
| 146 | |
| 147 | @dataclass(frozen=True) |
| 148 | class AudioLatentTools(LatentTools): |
| 149 | """ |
| 150 | Tools for building audio latent states. |
| 151 | """ |
| 152 | |
| 153 | patchifier: AudioPatchifier |
| 154 | target_shape: AudioLatentShape |
| 155 | |
| 156 | def create_initial_state( |
| 157 | self, |
| 158 | device: DeviceLikeType, |
| 159 | dtype: torch.dtype, |
| 160 | initial_latent: torch.Tensor | None = None, |
| 161 | ) -> LatentState: |
| 162 | if initial_latent is not None: |
| 163 | assert initial_latent.shape == self.target_shape.to_torch_shape(), ( |
| 164 | f"Latent shape {initial_latent.shape} does not match target shape {self.target_shape.to_torch_shape()}" |
| 165 | ) |
| 166 | else: |
| 167 | initial_latent = torch.zeros( |
| 168 | *self.target_shape.to_torch_shape(), |
| 169 | device=device, |
| 170 | dtype=dtype, |
| 171 | ) |
| 172 | |
| 173 | clean_latent = initial_latent.clone() |
| 174 | |
| 175 | denoise_mask = torch.ones( |
| 176 | *self.target_shape.mask_shape().to_torch_shape(), |
| 177 | device=device, |
| 178 | dtype=torch.float32, |
| 179 | ) |
| 180 | |
| 181 | latent_coords = self.patchifier.get_patch_grid_bounds( |
| 182 | output_shape=self.target_shape, |
| 183 | device=device, |
| 184 | ) |
| 185 | |
| 186 | return self.patchify( |
| 187 | LatentState( |
| 188 | latent=initial_latent, denoise_mask=denoise_mask, positions=latent_coords, clean_latent=clean_latent |
| 189 | ) |
| 190 | ) |
| 191 |