| 1 | from typing import Protocol |
| 2 | |
| 3 | import torch |
| 4 | |
| 5 | from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier |
| 6 | from ltx_core.components.protocols import DiffusionStepProtocol |
| 7 | from ltx_core.types import LatentState |
| 8 | from ltx_pipelines.utils.constants import VIDEO_LATENT_CHANNELS, VIDEO_SCALE_FACTORS |
| 9 | |
| 10 | |
| 11 | class PipelineComponents: |
| 12 | """ |
| 13 | Container class for pipeline components used throughout the LTX pipelines. |
| 14 | Attributes: |
| 15 | dtype (torch.dtype): Default torch dtype for tensors in the pipeline. |
| 16 | device (torch.device): Target device to place tensors and modules on. |
| 17 | video_scale_factors (SpatioTemporalScaleFactors): Scale factors (T, H, W) for VAE latent space. |
| 18 | video_latent_channels (int): Number of channels in the video latent representation. |
| 19 | video_patchifier (VideoLatentPatchifier): Patchifier instance for video latents. |
| 20 | audio_patchifier (AudioPatchifier): Patchifier instance for audio latents. |
| 21 | """ |
| 22 | |
| 23 | def __init__( |
| 24 | self, |
| 25 | dtype: torch.dtype, |
| 26 | device: torch.device, |
| 27 | ): |
| 28 | self.dtype = dtype |
| 29 | self.device = device |
| 30 | |
| 31 | self.video_scale_factors = VIDEO_SCALE_FACTORS |
| 32 | self.video_latent_channels = VIDEO_LATENT_CHANNELS |
| 33 | |
| 34 | self.video_patchifier = VideoLatentPatchifier(patch_size=1) |
| 35 | self.audio_patchifier = AudioPatchifier(patch_size=1) |
| 36 | |
| 37 | |
| 38 | class DenoisingFunc(Protocol): |
| 39 | """ |
| 40 | Protocol for a denoising function used in the LTX pipeline. |
| 41 | Args: |
| 42 | video_state (LatentState): The current latent state for video. |
| 43 | audio_state (LatentState): The current latent state for audio. |
| 44 | sigmas (torch.Tensor): A 1D tensor of sigma values for each diffusion step. |
| 45 | step_index (int): Index of the current denoising step. |
| 46 | Returns: |
| 47 | tuple[torch.Tensor, torch.Tensor]: The denoised video and audio tensors. |
| 48 | """ |
| 49 | |
| 50 | def __call__( |
| 51 | self, video_state: LatentState, audio_state: LatentState, sigmas: torch.Tensor, step_index: int |
| 52 | ) -> tuple[torch.Tensor, torch.Tensor]: ... |
| 53 | |
| 54 | |
| 55 | class DenoisingLoopFunc(Protocol): |
| 56 | """ |
| 57 | Protocol for a denoising loop function used in the LTX pipeline. |
| 58 | Args: |
| 59 | sigmas (torch.Tensor): A 1D tensor of sigma values for each diffusion step. |
| 60 | video_state (LatentState): The current latent state for video. |
| 61 | audio_state (LatentState): The current latent state for audio. |
| 62 | stepper (DiffusionStepProtocol): The diffusion step protocol to use. |
| 63 | Returns: |
| 64 | tuple[LatentState, LatentState]: The denoised video and audio latent states. |
| 65 | """ |
| 66 | |
| 67 | def __call__( |
| 68 | self, |
| 69 | sigmas: torch.Tensor, |
| 70 | video_state: LatentState, |
| 71 | audio_state: LatentState, |
| 72 | stepper: DiffusionStepProtocol, |
| 73 | ) -> tuple[torch.Tensor, torch.Tensor]: ... |
| 74 |