Automated PR - 2026-01-05
This commit is contained in:
sync-bot
@@ -0,0 +1,280 @@
|
||||
# LTX-Core
|
||||
|
||||
The foundational library for the LTX-2 Audio-Video generation model. This package contains the raw model definitions, component implementations, and loading logic used by `ltx-pipelines` and `ltx-trainer`.
|
||||
|
||||
## 📦 What's Inside?
|
||||
|
||||
- **`components/`**: Modular diffusion components (Schedulers, Guiders, Noisers, Patchifiers) following standard protocols
|
||||
- **`conditioning/`**: Tools for preparing latent states and applying conditioning (image, video, keyframes)
|
||||
- **`guidance/`**: Perturbation system for fine-grained control over attention mechanisms
|
||||
- **`loader/`**: Utilities for loading weights from `.safetensors`, fusing LoRAs, and managing memory
|
||||
- **`model/`**: PyTorch implementations of the LTX-2 Transformer, Video VAE, Audio VAE, Vocoder and Upscaler
|
||||
- **`text_encoders/gemma`**: Gemma text encoder implementation with tokenizers, feature extractors, and separate encoders for audio-video and video-only generation
|
||||
|
||||
## 🚀 Quick Start
|
||||
|
||||
`ltx-core` provides the building blocks (models, components, and utilities) needed to construct inference flows. For ready-made inference pipelines use [`ltx-pipelines`](../ltx-pipelines/) or [`ltx-trainer`](../ltx-trainer/) for training.
|
||||
|
||||
## 🔧 Installation
|
||||
|
||||
```bash
|
||||
# From the repository root
|
||||
uv sync --frozen
|
||||
|
||||
# Or install as a package
|
||||
pip install -e packages/ltx-core
|
||||
```
|
||||
|
||||
## Building Blocks Overview
|
||||
|
||||
`ltx-core` provides modular components that can be combined to build custom inference flows:
|
||||
|
||||
### Core Models
|
||||
|
||||
- **Transformer** ([`model/transformer/`](src/ltx_core/model/transformer/)): The 48-layer LTX-2 transformer with cross-modal attention for joint audio-video processing. Expects inputs in [`Modality`](src/ltx_core/model/transformer/modality.py) format
|
||||
- **Video VAE** ([`model/video_vae/`](src/ltx_core/model/video_vae/)): Encodes/decodes video pixels to/from latent space with temporal and spatial compression
|
||||
- **Audio VAE** ([`model/audio_vae/`](src/ltx_core/model/audio_vae/)): Encodes/decodes audio spectrograms to/from latent space
|
||||
- **Vocoder** ([`model/audio_vae/`](src/ltx_core/model/audio_vae/)): Neural vocoder that converts mel spectrograms to audio waveforms
|
||||
- **Text Encoder** ([`text_encoders/`](src/ltx_core/text_encoders/)): Gemma-based encoder that produces separate embeddings for video and audio conditioning
|
||||
- **Spatial Upscaler** ([`model/upsampler/`](src/ltx_core/model/upsampler/)): Upsamples latent representations for higher-resolution generation
|
||||
|
||||
### Diffusion Components
|
||||
|
||||
- **Schedulers** ([`components/schedulers.py`](src/ltx_core/components/schedulers.py)): Noise schedules (LTX2Scheduler, LinearQuadratic, Beta) that control the denoising process
|
||||
- **Guiders** ([`components/guiders.py`](src/ltx_core/components/guiders.py)): Guidance strategies (CFG, STG, APG) for controlling generation quality and adherence to prompts
|
||||
- **Noisers** ([`components/noisers.py`](src/ltx_core/components/noisers.py)): Add noise to latents according to the diffusion schedule
|
||||
- **Patchifiers** ([`components/patchifiers.py`](src/ltx_core/components/patchifiers.py)): Convert between spatial latents `[B, C, F, H, W]` and sequence format `[B, seq_len, dim]` for transformer processing
|
||||
|
||||
### Conditioning & Control
|
||||
|
||||
- **Conditioning** ([`conditioning/`](src/ltx_core/conditioning/)): Tools for preparing and applying various conditioning types (image, video, keyframes)
|
||||
- **Guidance** ([`guidance/`](src/ltx_core/guidance/)): Perturbation system for fine-grained control over attention mechanisms (e.g., skipping specific attention layers)
|
||||
|
||||
### Utilities
|
||||
|
||||
- **Loader** ([`loader/`](src/ltx_core/loader/)): Model loading from `.safetensors`, LoRA fusion, weight remapping, and memory management
|
||||
|
||||
For complete, production-ready pipeline implementations that combine these building blocks, see the [`ltx-pipelines`](../ltx-pipelines/) package.
|
||||
|
||||
---
|
||||
|
||||
# Architecture Overview
|
||||
|
||||
This section provides a deep dive into the internal architecture of the LTX-2 Audio-Video generation model.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
1. [High-Level Architecture](#high-level-architecture)
|
||||
2. [The Transformer](#the-transformer)
|
||||
3. [Video VAE](#video-vae)
|
||||
4. [Audio VAE](#audio-vae)
|
||||
5. [Text Encoding (Gemma)](#text-encoding-gemma)
|
||||
6. [Spatial Upscaler](#spatial-upsampler)
|
||||
7. [Data Flow](#data-flow)
|
||||
|
||||
---
|
||||
|
||||
## High-Level Architecture
|
||||
|
||||
LTX-2 is a **joint Audio-Video diffusion transformer** that processes both modalities simultaneously in a unified architecture. Unlike traditional models that handle video and audio separately, LTX-2 uses cross-modal attention to enable natural synchronization.
|
||||
|
||||
```text
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ INPUT PREPARATION │
|
||||
│ │
|
||||
│ Video Pixels → Video VAE Encoder → Video Latents │
|
||||
│ Audio Waveform → Audio VAE Encoder → Audio Latents │
|
||||
│ Text Prompt → Gemma Encoder → Text Embeddings │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
↓
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ LTX-2 TRANSFORMER (48 Blocks) │
|
||||
│ │
|
||||
│ ┌──────────────┐ ┌──────────────┐ │
|
||||
│ │ Video Stream │ │ Audio Stream │ │
|
||||
│ │ │ │ │ │
|
||||
│ │ Self-Attn │ │ Self-Attn │ │
|
||||
│ │ Cross-Attn │ │ Cross-Attn │ │
|
||||
│ │ │◄────────────►│ │ │
|
||||
│ │ A↔V Cross │ │ A↔V Cross │ │
|
||||
│ │ Feed-Forward │ │ Feed-Forward │ │
|
||||
│ └──────────────┘ └──────────────┘ │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
↓
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ OUTPUT DECODING │
|
||||
│ │
|
||||
│ Video Latents → Video VAE Decoder → Video Pixels │
|
||||
│ Audio Latents → Audio VAE Decoder → Mel Spectrogram │
|
||||
│ Mel Spectrogram → Vocoder → Audio Waveform │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## The Transformer
|
||||
|
||||
The core of LTX-2 is a 48-layer transformer that processes both video and audio tokens simultaneously.
|
||||
|
||||
### Model Structure
|
||||
|
||||
**Source**: [`src/ltx_core/model/transformer/model.py`](src/ltx_core/model/transformer/model.py)
|
||||
|
||||
The `LTXModel` class implements the transformer. It supports both video-only and audio-video generation modes. For actual usage, see the [`ltx-pipelines`](../ltx-pipelines/) package which handles model loading and initialization.
|
||||
|
||||
### Transformer Block Architecture
|
||||
|
||||
**Source**: [`src/ltx_core/model/transformer/transformer.py`](src/ltx_core/model/transformer/transformer.py)
|
||||
|
||||
```text
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ TRANSFORMER BLOCK │
|
||||
│ │
|
||||
│ VIDEO PATH: │
|
||||
│ Input → RMSNorm → AdaLN → Self-Attn (attn1) │
|
||||
│ → RMSNorm → Cross-Attn (attn2, text) │
|
||||
│ → RMSNorm → AdaLN → A↔V Cross-Attn │
|
||||
│ → RMSNorm → AdaLN → Feed-Forward (ff) → Output │
|
||||
│ │
|
||||
│ AUDIO PATH: │
|
||||
│ Input → RMSNorm → AdaLN → Self-Attn (audio_attn1) │
|
||||
│ → RMSNorm → Cross-Attn (audio_attn2, text) │
|
||||
│ → RMSNorm → AdaLN → A↔V Cross-Attn │
|
||||
│ → RMSNorm → AdaLN → Feed-Forward (audio_ff) │
|
||||
│ │
|
||||
│ AdaLN (Adaptive Layer Normalization): │
|
||||
│ - Uses scale_shift_table (6 params) for video/audio │
|
||||
│ - Uses scale_shift_table_a2v_ca (5 params) for A↔V CA │
|
||||
│ - Conditioned on per-token timestep embeddings │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### Perturbations
|
||||
|
||||
The transformer supports [**perturbations**](src/ltx_core/guidance/perturbations.py) that selectively skip attention operations.
|
||||
|
||||
Perturbations allow you to disable specific attention mechanisms during inference, which is useful for guidance techniques like STG (Spatio-Temporal Guidance).
|
||||
|
||||
**Supported Perturbation Types**:
|
||||
|
||||
- `SKIP_VIDEO_SELF_ATTN`: Skip video self-attention
|
||||
- `SKIP_AUDIO_SELF_ATTN`: Skip audio self-attention
|
||||
- `SKIP_A2V_CROSS_ATTN`: Skip audio-to-video cross-attention
|
||||
- `SKIP_V2A_CROSS_ATTN`: Skip video-to-audio cross-attention
|
||||
|
||||
Perturbations are used internally by guidance mechanisms like STG (Spatio-Temporal Guidance). For usage examples, see the [`ltx-pipelines`](../ltx-pipelines/) package.
|
||||
|
||||
---
|
||||
|
||||
## Video VAE
|
||||
|
||||
The Video VAE ([`src/ltx_core/model/video_vae/`](src/ltx_core/model/video_vae/)) encodes video pixels into latent representations and decodes them back.
|
||||
|
||||
### Architecture
|
||||
|
||||
- **Encoder**: Compresses `[B, 3, F, H, W]` pixels → `[B, 128, F', H/32, W/32]` latents
|
||||
- Where `F' = 1 + (F-1)/8` (frame count must satisfy `(F-1) % 8 == 0`)
|
||||
- Example: `[B, 3, 33, 512, 512]` → `[B, 128, 5, 16, 16]`
|
||||
- **Decoder**: Expands `[B, 128, F, H, W]` latents → `[B, 3, F', H*32, W*32]` pixels
|
||||
- Where `F' = 1 + (F-1)*8`
|
||||
- Example: `[B, 128, 5, 16, 16]` → `[B, 3, 33, 512, 512]`
|
||||
|
||||
The Video VAE is used internally by pipelines for encoding video pixels to latents and decoding latents back to pixels. For usage examples, see the [`ltx-pipelines`](../ltx-pipelines/) package.
|
||||
|
||||
---
|
||||
|
||||
## Audio VAE
|
||||
|
||||
The Audio VAE ([`src/ltx_core/model/audio_vae/`](src/ltx_core/model/audio_vae/)) processes audio spectrograms.
|
||||
|
||||
### Audio VAE Architecture
|
||||
|
||||
- **Encoder**: Compresses mel spectrogram `[B, mel_bins, T]` → `[B, 8, T/4, 16]` latents
|
||||
- Temporal downsampling: 4× (`LATENT_DOWNSAMPLE_FACTOR = 4`)
|
||||
- Frequency bins: Fixed 16 mel bins in latent space
|
||||
- Latent channels: 8
|
||||
- **Decoder**: Expands `[B, 8, T, 16]` latents → mel spectrogram `[B, mel_bins, T*4]`
|
||||
- **Vocoder**: Converts mel spectrogram → audio waveform
|
||||
|
||||
**Downsampling**:
|
||||
|
||||
- Temporal: 4× (time steps)
|
||||
- Frequency: Variable (input mel_bins → fixed 16 in latent space)
|
||||
|
||||
The Audio VAE is used internally by pipelines for encoding mel spectrograms to latents and decoding latents back to mel spectrograms. The vocoder converts mel spectrograms to audio waveforms. For usage examples, see the [`ltx-pipelines`](../ltx-pipelines/) package.
|
||||
|
||||
---
|
||||
|
||||
## Text Encoding (Gemma)
|
||||
|
||||
LTX-2 uses **Gemma** (Google's open LLM) as the text encoder, located in [`src/ltx_core/text_encoders/gemma/`](src/ltx_core/text_encoders/gemma/).
|
||||
|
||||
### Text Encoder Architecture
|
||||
|
||||
- **Tokenizer**: Converts text → token IDs
|
||||
- **Gemma Model**: Processes tokens → embeddings
|
||||
- **Text Projection**: Uses `PixArtAlphaTextProjection` to project caption embeddings
|
||||
- Two-layer MLP with GELU (tanh approximation) or SiLU activation
|
||||
- Projects from caption channels (3840) to model dimensions
|
||||
- **Feature Extractor**: Extracts video/audio-specific embeddings
|
||||
- **Separate Encoders**:
|
||||
- `AVEncoder`: For audio-video generation (outputs separate video and audio contexts)
|
||||
- `VideoOnlyEncoder`: For video-only generation
|
||||
|
||||
### System Prompts
|
||||
|
||||
System prompts are also used to enhance user's prompts.
|
||||
|
||||
- **Text-to-Video**: [`gemma_t2v_system_prompt.txt`](src/ltx_core/text_encoders/gemma/encoders/prompts/gemma_t2v_system_prompt.txt)
|
||||
- **Image-to-Video**: [`gemma_i2v_system_prompt.txt`](src/ltx_core/text_encoders/gemma/encoders/prompts/gemma_i2v_system_prompt.txt)
|
||||
|
||||
**Important**: Video and audio receive **different** context embeddings, even from the same prompt. This allows better modality-specific conditioning.
|
||||
|
||||
**Output Format**:
|
||||
|
||||
- Video context: `[B, seq_len, 4096]` - Video-specific text embeddings
|
||||
- Audio context: `[B, seq_len, 2048]` - Audio-specific text embeddings
|
||||
|
||||
The text encoder is used internally by pipelines. For usage examples, see the [`ltx-pipelines`](../ltx-pipelines/) package.
|
||||
|
||||
---
|
||||
|
||||
## Upscaler
|
||||
|
||||
The Upscaler ([`src/ltx_core/model/upsampler/`](src/ltx_core/model/upsampler/)) upsamples latent representations for higher-resolution output.
|
||||
|
||||
The spatial upsampler is used internally by two-stage pipelines (e.g., [`TI2VidTwoStagesPipeline`](../ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py), [`ICLoraPipeline`](../ltx-pipelines/src/ltx_pipelines/ic_lora.py)) to upsample low-resolution latents before final VAE decoding. For usage examples, see the [`ltx-pipelines`](../ltx-pipelines/) package.
|
||||
|
||||
---
|
||||
|
||||
## Data Flow
|
||||
|
||||
### Complete Generation Pipeline
|
||||
|
||||
Here's how all the components work together conceptually ([`src/ltx_core/components/`](src/ltx_core/components/)):
|
||||
|
||||
**Pipeline Steps**:
|
||||
|
||||
1. **Text Encoding**: Text prompt → Gemma encoder → separate video/audio embeddings
|
||||
2. **Latent Initialization**: Initialize noise latents in spatial format `[B, C, F, H, W]`
|
||||
3. **Patchification**: Convert spatial latents to sequence format `[B, seq_len, dim]` for transformer
|
||||
4. **Sigma Schedule**: Generate noise schedule (adapts to token count)
|
||||
5. **Denoising Loop**: Iteratively denoise using transformer predictions
|
||||
- Create Modality inputs with per-token timesteps and RoPE positions
|
||||
- Forward pass through transformer (conditional and unconditional for CFG)
|
||||
- Apply guidance (CFG, STG, etc.)
|
||||
- Update latents using diffusion step (Euler, etc.)
|
||||
6. **Unpatchification**: Convert sequence back to spatial format
|
||||
7. **VAE Decoding**: Decode latents to pixel space (with optional upsampling for two-stage)
|
||||
|
||||
- [`TI2VidTwoStagesPipeline`](../ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py) - Two-stage text-to-video (recommended)
|
||||
- [`ICLoraPipeline`](../ltx-pipelines/src/ltx_pipelines/ic_lora.py) - Video-to-video with IC-LoRA control
|
||||
- [`DistilledPipeline`](../ltx-pipelines/src/ltx_pipelines/distilled.py) - Fast inference with distilled model
|
||||
- [`KeyframeInterpolationPipeline`](../ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py) - Keyframe-based interpolation
|
||||
|
||||
See the [ltx-pipelines README](../ltx-pipelines/README.md) for usage examples.
|
||||
|
||||
## 🔗 Related Projects
|
||||
|
||||
- **[ltx-pipelines](../ltx-pipelines/)** - High-level pipeline implementations for text-to-video, image-to-video, and video-to-video
|
||||
- **[ltx-trainer](../ltx-trainer/)** - Training and fine-tuning tools
|
||||
@@ -0,0 +1,37 @@
|
||||
[project]
|
||||
name = "ltx-core"
|
||||
version = "1.0.0"
|
||||
description = "Core implementation of Lightricks' LTX-2 model"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
dependencies = [
|
||||
"torch~=2.7",
|
||||
"torchaudio",
|
||||
"einops",
|
||||
"numpy",
|
||||
"transformers",
|
||||
"safetensors",
|
||||
"accelerate",
|
||||
"scipy>=1.14",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
xformers = ["xformers"]
|
||||
|
||||
|
||||
[tool.uv.sources]
|
||||
xformers = { index = "pytorch" }
|
||||
|
||||
[[tool.uv.index]]
|
||||
name = "pytorch"
|
||||
url = "https://download.pytorch.org/whl/cu129"
|
||||
explicit = true
|
||||
|
||||
[build-system]
|
||||
requires = ["uv_build>=0.9.8,<0.10.0"]
|
||||
build-backend = "uv_build"
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"scikit-image>=0.25.2",
|
||||
]
|
||||
@@ -0,0 +1,10 @@
|
||||
"""
|
||||
Diffusion pipeline components.
|
||||
Submodules:
|
||||
diffusion_steps - Diffusion stepping algorithms (EulerDiffusionStep)
|
||||
guiders - Guidance strategies (CFGGuider, STGGuider, APG variants)
|
||||
noisers - Noise samplers (GaussianNoiser)
|
||||
patchifiers - Latent patchification (VideoLatentPatchifier, AudioPatchifier)
|
||||
protocols - Protocol definitions (Patchifier, etc.)
|
||||
schedulers - Sigma schedulers (LTX2Scheduler, LinearQuadraticScheduler)
|
||||
"""
|
||||
@@ -0,0 +1,22 @@
|
||||
import torch
|
||||
|
||||
from ltx_core.components.protocols import DiffusionStepProtocol
|
||||
from ltx_core.utils import to_velocity
|
||||
|
||||
|
||||
class EulerDiffusionStep(DiffusionStepProtocol):
|
||||
"""
|
||||
First-order Euler method for diffusion sampling.
|
||||
Takes a single step from the current noise level (sigma) to the next by
|
||||
computing velocity from the denoised prediction and applying: sample + velocity * dt.
|
||||
"""
|
||||
|
||||
def step(
|
||||
self, sample: torch.Tensor, denoised_sample: torch.Tensor, sigmas: torch.Tensor, step_index: int
|
||||
) -> torch.Tensor:
|
||||
sigma = sigmas[step_index]
|
||||
sigma_next = sigmas[step_index + 1]
|
||||
dt = sigma_next - sigma
|
||||
velocity = to_velocity(sample, sigma, denoised_sample)
|
||||
|
||||
return (sample.to(torch.float32) + velocity.to(torch.float32) * dt).to(sample.dtype)
|
||||
@@ -0,0 +1,198 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.components.protocols import GuiderProtocol
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CFGGuider(GuiderProtocol):
|
||||
"""
|
||||
Classifier-free guidance (CFG) guider.
|
||||
Computes the guidance delta as (scale - 1) * (cond - uncond), steering the
|
||||
denoising process toward the conditioned prediction.
|
||||
Attributes:
|
||||
scale: Guidance strength. 1.0 means no guidance, higher values increase
|
||||
adherence to the conditioning.
|
||||
"""
|
||||
|
||||
scale: float
|
||||
|
||||
def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor:
|
||||
return (self.scale - 1) * (cond - uncond)
|
||||
|
||||
def enabled(self) -> bool:
|
||||
return self.scale != 1.0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CFGStarRescalingGuider(GuiderProtocol):
|
||||
"""
|
||||
Calculates the CFG delta between conditioned and unconditioned samples.
|
||||
To minimize offset in the denoising direction and move mostly along the
|
||||
conditioning axis within the distribution, the unconditioned sample is
|
||||
rescaled in accordance with the norm of the conditioned sample.
|
||||
Attributes:
|
||||
scale (float):
|
||||
Global guidance strength. A value of 1.0 corresponds to no extra
|
||||
guidance beyond the base model prediction. Values > 1.0 increase
|
||||
the influence of the conditioned sample relative to the
|
||||
unconditioned one.
|
||||
"""
|
||||
|
||||
scale: float
|
||||
|
||||
def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor:
|
||||
rescaled_neg = projection_coef(cond, uncond) * uncond
|
||||
return (self.scale - 1) * (cond - rescaled_neg)
|
||||
|
||||
def enabled(self) -> bool:
|
||||
return self.scale != 1.0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class STGGuider(GuiderProtocol):
|
||||
"""
|
||||
Calculates the STG delta between conditioned and perturbed denoised samples.
|
||||
Perturbed samples are the result of the denoising process with perturbations,
|
||||
e.g. attentions acting as passthrough for certain layers and modalities.
|
||||
Attributes:
|
||||
scale (float):
|
||||
Global strength of the STG guidance. A value of 0.0 disables the
|
||||
guidance. Larger values increase the correction applied in the
|
||||
direction of (pos_denoised - perturbed_denoised).
|
||||
"""
|
||||
|
||||
scale: float
|
||||
|
||||
def delta(self, pos_denoised: torch.Tensor, perturbed_denoised: torch.Tensor) -> torch.Tensor:
|
||||
return self.scale * (pos_denoised - perturbed_denoised)
|
||||
|
||||
def enabled(self) -> bool:
|
||||
return self.scale != 0.0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LtxAPGGuider(GuiderProtocol):
|
||||
"""
|
||||
Calculates the APG (adaptive projected guidance) delta between conditioned
|
||||
and unconditioned samples.
|
||||
To minimize offset in the denoising direction and move mostly along the
|
||||
conditioning axis within the distribution, the (cond - uncond) delta is
|
||||
decomposed into components parallel and orthogonal to the conditioned
|
||||
sample. The `eta` parameter weights the parallel component, while `scale`
|
||||
is applied to the orthogonal component. Optionally, a norm threshold can
|
||||
be used to suppress guidance when the magnitude of the correction is small.
|
||||
Attributes:
|
||||
scale (float):
|
||||
Strength applied to the component of the guidance that is orthogonal
|
||||
to the conditioned sample. Controls how aggressively we move in
|
||||
directions that change semantics but stay consistent with the
|
||||
conditioning manifold.
|
||||
eta (float):
|
||||
Weight of the component of the guidance that is parallel to the
|
||||
conditioned sample. A value of 1.0 keeps the full parallel
|
||||
component; values in [0, 1] attenuate it, and values > 1.0 amplify
|
||||
motion along the conditioning direction.
|
||||
norm_threshold (float):
|
||||
Minimum L2 norm of the guidance delta below which the guidance
|
||||
can be reduced or ignored (depending on implementation).
|
||||
This is useful for avoiding noisy or unstable updates when the
|
||||
guidance signal is very small.
|
||||
"""
|
||||
|
||||
scale: float
|
||||
eta: float = 1.0
|
||||
norm_threshold: float = 0.0
|
||||
|
||||
def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor:
|
||||
guidance = cond - uncond
|
||||
if self.norm_threshold > 0:
|
||||
ones = torch.ones_like(guidance)
|
||||
guidance_norm = guidance.norm(p=2, dim=[-1, -2, -3], keepdim=True)
|
||||
scale_factor = torch.minimum(ones, self.norm_threshold / guidance_norm)
|
||||
guidance = guidance * scale_factor
|
||||
proj_coeff = projection_coef(guidance, cond)
|
||||
g_parallel = proj_coeff * cond
|
||||
g_orth = guidance - g_parallel
|
||||
g_apg = g_parallel * self.eta + g_orth
|
||||
|
||||
return g_apg * (self.scale - 1)
|
||||
|
||||
def enabled(self) -> bool:
|
||||
return self.scale != 1.0
|
||||
|
||||
|
||||
@dataclass(frozen=False)
|
||||
class LegacyStatefulAPGGuider(GuiderProtocol):
|
||||
"""
|
||||
Calculates the APG (adaptive projected guidance) delta between conditioned
|
||||
and unconditioned samples.
|
||||
To minimize offset in the denoising direction and move mostly along the
|
||||
conditioning axis within the distribution, the (cond - uncond) delta is
|
||||
decomposed into components parallel and orthogonal to the conditioned
|
||||
sample. The `eta` parameter weights the parallel component, while `scale`
|
||||
is applied to the orthogonal component. Optionally, a norm threshold can
|
||||
be used to suppress guidance when the magnitude of the correction is small.
|
||||
Attributes:
|
||||
scale (float):
|
||||
Strength applied to the component of the guidance that is orthogonal
|
||||
to the conditioned sample. Controls how aggressively we move in
|
||||
directions that change semantics but stay consistent with the
|
||||
conditioning manifold.
|
||||
eta (float):
|
||||
Weight of the component of the guidance that is parallel to the
|
||||
conditioned sample. A value of 1.0 keeps the full parallel
|
||||
component; values in [0, 1] attenuate it, and values > 1.0 amplify
|
||||
motion along the conditioning direction.
|
||||
norm_threshold (float):
|
||||
Minimum L2 norm of the guidance delta below which the guidance
|
||||
can be reduced or ignored (depending on implementation).
|
||||
This is useful for avoiding noisy or unstable updates when the
|
||||
guidance signal is very small.
|
||||
momentum (float):
|
||||
Exponential moving-average coefficient for accumulating guidance
|
||||
over time. running_avg = momentum * running_avg + guidance
|
||||
"""
|
||||
|
||||
scale: float
|
||||
eta: float
|
||||
norm_threshold: float = 5.0
|
||||
momentum: float = 0.0
|
||||
# it is user's responsibility not to use same APGGuider for several denoisings or different modalities
|
||||
# in order not to share accumulated average across different denoisings or modalities
|
||||
running_avg: torch.Tensor | None = None
|
||||
|
||||
def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor:
|
||||
guidance = cond - uncond
|
||||
if self.momentum != 0:
|
||||
if self.running_avg is None:
|
||||
self.running_avg = guidance.clone()
|
||||
else:
|
||||
self.running_avg = self.momentum * self.running_avg + guidance
|
||||
guidance = self.running_avg
|
||||
|
||||
if self.norm_threshold > 0:
|
||||
ones = torch.ones_like(guidance)
|
||||
guidance_norm = guidance.norm(p=2, dim=[-1, -2, -3], keepdim=True)
|
||||
scale_factor = torch.minimum(ones, self.norm_threshold / guidance_norm)
|
||||
guidance = guidance * scale_factor
|
||||
|
||||
proj_coeff = projection_coef(guidance, cond)
|
||||
g_parallel = proj_coeff * cond
|
||||
g_orth = guidance - g_parallel
|
||||
g_apg = g_parallel * self.eta + g_orth
|
||||
|
||||
return g_apg * self.scale
|
||||
|
||||
def enabled(self) -> bool:
|
||||
return self.scale != 0.0
|
||||
|
||||
|
||||
def projection_coef(to_project: torch.Tensor, project_onto: torch.Tensor) -> torch.Tensor:
|
||||
batch_size = to_project.shape[0]
|
||||
positive_flat = to_project.reshape(batch_size, -1)
|
||||
negative_flat = project_onto.reshape(batch_size, -1)
|
||||
dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True)
|
||||
squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
|
||||
return dot_product / squared_norm
|
||||
@@ -0,0 +1,35 @@
|
||||
from dataclasses import replace
|
||||
from typing import Protocol
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.types import LatentState
|
||||
|
||||
|
||||
class Noiser(Protocol):
|
||||
"""Protocol for adding noise to a latent state during diffusion."""
|
||||
|
||||
def __call__(self, latent_state: LatentState, noise_scale: float) -> LatentState: ...
|
||||
|
||||
|
||||
class GaussianNoiser(Noiser):
|
||||
"""Adds Gaussian noise to a latent state, scaled by the denoise mask."""
|
||||
|
||||
def __init__(self, generator: torch.Generator):
|
||||
super().__init__()
|
||||
|
||||
self.generator = generator
|
||||
|
||||
def __call__(self, latent_state: LatentState, noise_scale: float = 1.0) -> LatentState:
|
||||
noise = torch.randn(
|
||||
*latent_state.latent.shape,
|
||||
device=latent_state.latent.device,
|
||||
dtype=latent_state.latent.dtype,
|
||||
generator=self.generator,
|
||||
)
|
||||
scaled_mask = latent_state.denoise_mask * noise_scale
|
||||
latent = noise * scaled_mask + latent_state.latent * (1 - scaled_mask)
|
||||
return replace(
|
||||
latent_state,
|
||||
latent=latent.to(latent_state.latent.dtype),
|
||||
)
|
||||
@@ -0,0 +1,348 @@
|
||||
import math
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import einops
|
||||
import torch
|
||||
|
||||
from ltx_core.components.protocols import Patchifier
|
||||
from ltx_core.types import AudioLatentShape, SpatioTemporalScaleFactors, VideoLatentShape
|
||||
|
||||
|
||||
class VideoLatentPatchifier(Patchifier):
|
||||
def __init__(self, patch_size: int):
|
||||
# Patch sizes for video latents.
|
||||
self._patch_size = (
|
||||
1, # temporal dimension
|
||||
patch_size, # height dimension
|
||||
patch_size, # width dimension
|
||||
)
|
||||
|
||||
@property
|
||||
def patch_size(self) -> Tuple[int, int, int]:
|
||||
return self._patch_size
|
||||
|
||||
def get_token_count(self, tgt_shape: VideoLatentShape) -> int:
|
||||
return math.prod(tgt_shape.to_torch_shape()[2:]) // math.prod(self._patch_size)
|
||||
|
||||
def patchify(
|
||||
self,
|
||||
latents: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
latents = einops.rearrange(
|
||||
latents,
|
||||
"b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)",
|
||||
p1=self._patch_size[0],
|
||||
p2=self._patch_size[1],
|
||||
p3=self._patch_size[2],
|
||||
)
|
||||
|
||||
return latents
|
||||
|
||||
def unpatchify(
|
||||
self,
|
||||
latents: torch.Tensor,
|
||||
output_shape: VideoLatentShape,
|
||||
) -> torch.Tensor:
|
||||
assert self._patch_size[0] == 1, "Temporal patch size must be 1 for symmetric patchifier"
|
||||
|
||||
patch_grid_frames = output_shape.frames // self._patch_size[0]
|
||||
patch_grid_height = output_shape.height // self._patch_size[1]
|
||||
patch_grid_width = output_shape.width // self._patch_size[2]
|
||||
|
||||
latents = einops.rearrange(
|
||||
latents,
|
||||
"b (f h w) (c p q) -> b c f (h p) (w q)",
|
||||
f=patch_grid_frames,
|
||||
h=patch_grid_height,
|
||||
w=patch_grid_width,
|
||||
p=self._patch_size[1],
|
||||
q=self._patch_size[2],
|
||||
)
|
||||
|
||||
return latents
|
||||
|
||||
def get_patch_grid_bounds(
|
||||
self,
|
||||
output_shape: AudioLatentShape | VideoLatentShape,
|
||||
device: Optional[torch.device] = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Return the per-dimension bounds [inclusive start, exclusive end) for every
|
||||
patch produced by `patchify`. The bounds are expressed in the original
|
||||
video grid coordinates: frame/time, height, and width.
|
||||
The resulting tensor is shaped `[batch_size, 3, num_patches, 2]`, where:
|
||||
- axis 1 (size 3) enumerates (frame/time, height, width) dimensions
|
||||
- axis 3 (size 2) stores `[start, end)` indices within each dimension
|
||||
Args:
|
||||
output_shape: Video grid description containing frames, height, and width.
|
||||
device: Device of the latent tensor.
|
||||
"""
|
||||
if not isinstance(output_shape, VideoLatentShape):
|
||||
raise ValueError("VideoLatentPatchifier expects VideoLatentShape when computing coordinates")
|
||||
|
||||
frames = output_shape.frames
|
||||
height = output_shape.height
|
||||
width = output_shape.width
|
||||
batch_size = output_shape.batch
|
||||
|
||||
# Validate inputs to ensure positive dimensions
|
||||
assert frames > 0, f"frames must be positive, got {frames}"
|
||||
assert height > 0, f"height must be positive, got {height}"
|
||||
assert width > 0, f"width must be positive, got {width}"
|
||||
assert batch_size > 0, f"batch_size must be positive, got {batch_size}"
|
||||
|
||||
# Generate grid coordinates for each dimension (frame, height, width)
|
||||
# We use torch.arange to create the starting coordinates for each patch.
|
||||
# indexing='ij' ensures the dimensions are in the order (frame, height, width).
|
||||
grid_coords = torch.meshgrid(
|
||||
torch.arange(start=0, end=frames, step=self._patch_size[0], device=device),
|
||||
torch.arange(start=0, end=height, step=self._patch_size[1], device=device),
|
||||
torch.arange(start=0, end=width, step=self._patch_size[2], device=device),
|
||||
indexing="ij",
|
||||
)
|
||||
|
||||
# Stack the grid coordinates to create the start coordinates tensor.
|
||||
# Shape becomes (3, grid_f, grid_h, grid_w)
|
||||
patch_starts = torch.stack(grid_coords, dim=0)
|
||||
|
||||
# Create a tensor containing the size of a single patch:
|
||||
# (frame_patch_size, height_patch_size, width_patch_size).
|
||||
# Reshape to (3, 1, 1, 1) to enable broadcasting when adding to the start coordinates.
|
||||
patch_size_delta = torch.tensor(
|
||||
self._patch_size,
|
||||
device=patch_starts.device,
|
||||
dtype=patch_starts.dtype,
|
||||
).view(3, 1, 1, 1)
|
||||
|
||||
# Calculate end coordinates: start + patch_size
|
||||
# Shape becomes (3, grid_f, grid_h, grid_w)
|
||||
patch_ends = patch_starts + patch_size_delta
|
||||
|
||||
# Stack start and end coordinates together along the last dimension
|
||||
# Shape becomes (3, grid_f, grid_h, grid_w, 2), where the last dimension is [start, end]
|
||||
latent_coords = torch.stack((patch_starts, patch_ends), dim=-1)
|
||||
|
||||
# Broadcast to batch size and flatten all spatial/temporal dimensions into one sequence.
|
||||
# Final Shape: (batch_size, 3, num_patches, 2)
|
||||
latent_coords = einops.repeat(
|
||||
latent_coords,
|
||||
"c f h w bounds -> b c (f h w) bounds",
|
||||
b=batch_size,
|
||||
bounds=2,
|
||||
)
|
||||
|
||||
return latent_coords
|
||||
|
||||
|
||||
def get_pixel_coords(
|
||||
latent_coords: torch.Tensor,
|
||||
scale_factors: SpatioTemporalScaleFactors,
|
||||
causal_fix: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Map latent-space `[start, end)` coordinates to their pixel-space equivalents by scaling
|
||||
each axis (frame/time, height, width) with the corresponding VAE downsampling factors.
|
||||
Optionally compensate for causal encoding that keeps the first frame at unit temporal scale.
|
||||
Args:
|
||||
latent_coords: Tensor of latent bounds shaped `(batch, 3, num_patches, 2)`.
|
||||
scale_factors: SpatioTemporalScaleFactors tuple `(temporal, height, width)` with integer scale factors applied
|
||||
per axis.
|
||||
causal_fix: When True, rewrites the temporal axis of the first frame so causal VAEs
|
||||
that treat frame zero differently still yield non-negative timestamps.
|
||||
"""
|
||||
# Broadcast the VAE scale factors so they align with the `(batch, axis, patch, bound)` layout.
|
||||
broadcast_shape = [1] * latent_coords.ndim
|
||||
broadcast_shape[1] = -1 # axis dimension corresponds to (frame/time, height, width)
|
||||
scale_tensor = torch.tensor(scale_factors, device=latent_coords.device).view(*broadcast_shape)
|
||||
|
||||
# Apply per-axis scaling to convert latent bounds into pixel-space coordinates.
|
||||
pixel_coords = latent_coords * scale_tensor
|
||||
|
||||
if causal_fix:
|
||||
# VAE temporal stride for the very first frame is 1 instead of `scale_factors[0]`.
|
||||
# Shift and clamp to keep the first-frame timestamps causal and non-negative.
|
||||
pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + 1 - scale_factors[0]).clamp(min=0)
|
||||
|
||||
return pixel_coords
|
||||
|
||||
|
||||
class AudioPatchifier(Patchifier):
|
||||
def __init__(
|
||||
self,
|
||||
patch_size: int,
|
||||
sample_rate: int = 16000,
|
||||
hop_length: int = 160,
|
||||
audio_latent_downsample_factor: int = 4,
|
||||
is_causal: bool = True,
|
||||
shift: int = 0,
|
||||
):
|
||||
"""
|
||||
Patchifier tailored for spectrogram/audio latents.
|
||||
Args:
|
||||
patch_size: Number of mel bins combined into a single patch. This
|
||||
controls the resolution along the frequency axis.
|
||||
sample_rate: Original waveform sampling rate. Used to map latent
|
||||
indices back to seconds so downstream consumers can align audio
|
||||
and video cues.
|
||||
hop_length: Window hop length used for the spectrogram. Determines
|
||||
how many real-time samples separate two consecutive latent frames.
|
||||
audio_latent_downsample_factor: Ratio between spectrogram frames and
|
||||
latent frames; compensates for additional downsampling inside the
|
||||
VAE encoder.
|
||||
is_causal: When True, timing is shifted to account for causal
|
||||
receptive fields so timestamps do not peek into the future.
|
||||
shift: Integer offset applied to the latent indices. Enables
|
||||
constructing overlapping windows from the same latent sequence.
|
||||
"""
|
||||
self.hop_length = hop_length
|
||||
self.sample_rate = sample_rate
|
||||
self.audio_latent_downsample_factor = audio_latent_downsample_factor
|
||||
self.is_causal = is_causal
|
||||
self.shift = shift
|
||||
self._patch_size = (1, patch_size, patch_size)
|
||||
|
||||
@property
|
||||
def patch_size(self) -> Tuple[int, int, int]:
|
||||
return self._patch_size
|
||||
|
||||
def get_token_count(self, tgt_shape: AudioLatentShape) -> int:
|
||||
return tgt_shape.frames
|
||||
|
||||
def _get_audio_latent_time_in_sec(
|
||||
self,
|
||||
start_latent: int,
|
||||
end_latent: int,
|
||||
dtype: torch.dtype,
|
||||
device: Optional[torch.device] = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Converts latent indices into real-time seconds while honoring causal
|
||||
offsets and the configured hop length.
|
||||
Args:
|
||||
start_latent: Inclusive start index inside the latent sequence. This
|
||||
sets the first timestamp returned.
|
||||
end_latent: Exclusive end index. Determines how many timestamps get
|
||||
generated.
|
||||
dtype: Floating-point dtype used for the returned tensor, allowing
|
||||
callers to control precision.
|
||||
device: Target device for the timestamp tensor. When omitted the
|
||||
computation occurs on CPU to avoid surprising GPU allocations.
|
||||
"""
|
||||
if device is None:
|
||||
device = torch.device("cpu")
|
||||
|
||||
audio_latent_frame = torch.arange(start_latent, end_latent, dtype=dtype, device=device)
|
||||
|
||||
audio_mel_frame = audio_latent_frame * self.audio_latent_downsample_factor
|
||||
|
||||
if self.is_causal:
|
||||
# Frame offset for causal alignment.
|
||||
# The "+1" ensures the timestamp corresponds to the first sample that is fully available.
|
||||
causal_offset = 1
|
||||
audio_mel_frame = (audio_mel_frame + causal_offset - self.audio_latent_downsample_factor).clip(min=0)
|
||||
|
||||
return audio_mel_frame * self.hop_length / self.sample_rate
|
||||
|
||||
def _compute_audio_timings(
|
||||
self,
|
||||
batch_size: int,
|
||||
num_steps: int,
|
||||
device: Optional[torch.device] = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Builds a `(B, 1, T, 2)` tensor containing timestamps for each latent frame.
|
||||
This helper method underpins `get_patch_grid_bounds` for the audio patchifier.
|
||||
Args:
|
||||
batch_size: Number of sequences to broadcast the timings over.
|
||||
num_steps: Number of latent frames (time steps) to convert into timestamps.
|
||||
device: Device on which the resulting tensor should reside.
|
||||
"""
|
||||
resolved_device = device
|
||||
if resolved_device is None:
|
||||
resolved_device = torch.device("cpu")
|
||||
|
||||
start_timings = self._get_audio_latent_time_in_sec(
|
||||
self.shift,
|
||||
num_steps + self.shift,
|
||||
torch.float32,
|
||||
resolved_device,
|
||||
)
|
||||
start_timings = start_timings.unsqueeze(0).expand(batch_size, -1).unsqueeze(1)
|
||||
|
||||
end_timings = self._get_audio_latent_time_in_sec(
|
||||
self.shift + 1,
|
||||
num_steps + self.shift + 1,
|
||||
torch.float32,
|
||||
resolved_device,
|
||||
)
|
||||
end_timings = end_timings.unsqueeze(0).expand(batch_size, -1).unsqueeze(1)
|
||||
|
||||
return torch.stack([start_timings, end_timings], dim=-1)
|
||||
|
||||
def patchify(
|
||||
self,
|
||||
audio_latents: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Flattens the audio latent tensor along time. Use `get_patch_grid_bounds`
|
||||
to derive timestamps for each latent frame based on the configured hop
|
||||
length and downsampling.
|
||||
Args:
|
||||
audio_latents: Latent tensor to patchify.
|
||||
Returns:
|
||||
Flattened patch tokens tensor. Use `get_patch_grid_bounds` to compute the
|
||||
corresponding timing metadata when needed.
|
||||
"""
|
||||
audio_latents = einops.rearrange(
|
||||
audio_latents,
|
||||
"b c t f -> b t (c f)",
|
||||
)
|
||||
|
||||
return audio_latents
|
||||
|
||||
def unpatchify(
|
||||
self,
|
||||
audio_latents: torch.Tensor,
|
||||
output_shape: AudioLatentShape,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Restores the `(B, C, T, F)` spectrogram tensor from flattened patches.
|
||||
Use `get_patch_grid_bounds` to recompute the timestamps that describe each
|
||||
frame's position in real time.
|
||||
Args:
|
||||
audio_latents: Latent tensor to unpatchify.
|
||||
output_shape: Shape of the unpatched output tensor.
|
||||
Returns:
|
||||
Unpatched latent tensor. Use `get_patch_grid_bounds` to compute the timing
|
||||
metadata associated with the restored latents.
|
||||
"""
|
||||
# audio_latents shape: (batch, time, freq * channels)
|
||||
audio_latents = einops.rearrange(
|
||||
audio_latents,
|
||||
"b t (c f) -> b c t f",
|
||||
c=output_shape.channels,
|
||||
f=output_shape.mel_bins,
|
||||
)
|
||||
|
||||
return audio_latents
|
||||
|
||||
def get_patch_grid_bounds(
|
||||
self,
|
||||
output_shape: AudioLatentShape | VideoLatentShape,
|
||||
device: Optional[torch.device] = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Return the temporal bounds `[inclusive start, exclusive end)` for every
|
||||
patch emitted by `patchify`. For audio this corresponds to timestamps in
|
||||
seconds aligned with the original spectrogram grid.
|
||||
The returned tensor has shape `[batch_size, 1, time_steps, 2]`, where:
|
||||
- axis 1 (size 1) represents the temporal dimension
|
||||
- axis 3 (size 2) stores the `[start, end)` timestamps per patch
|
||||
Args:
|
||||
output_shape: Audio grid specification describing the number of time steps.
|
||||
device: Target device for the returned tensor.
|
||||
"""
|
||||
if not isinstance(output_shape, AudioLatentShape):
|
||||
raise ValueError("AudioPatchifier expects AudioLatentShape when computing coordinates")
|
||||
|
||||
return self._compute_audio_timings(output_shape.batch, output_shape.frames, device)
|
||||
@@ -0,0 +1,101 @@
|
||||
from typing import Protocol, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.types import AudioLatentShape, VideoLatentShape
|
||||
|
||||
|
||||
class Patchifier(Protocol):
|
||||
"""
|
||||
Protocol for patchifiers that convert latent tensors into patches and assemble them back.
|
||||
"""
|
||||
|
||||
def patchify(
|
||||
self,
|
||||
latents: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
...
|
||||
"""
|
||||
Convert latent tensors into flattened patch tokens.
|
||||
Args:
|
||||
latents: Latent tensor to patchify.
|
||||
Returns:
|
||||
Flattened patch tokens tensor.
|
||||
"""
|
||||
|
||||
def unpatchify(
|
||||
self,
|
||||
latents: torch.Tensor,
|
||||
output_shape: AudioLatentShape | VideoLatentShape,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Converts latent tensors between spatio-temporal formats and flattened sequence representations.
|
||||
Args:
|
||||
latents: Patch tokens that must be rearranged back into the latent grid constructed by `patchify`.
|
||||
output_shape: Shape of the output tensor. Note that output_shape is either AudioLatentShape or
|
||||
VideoLatentShape.
|
||||
Returns:
|
||||
Dense latent tensor restored from the flattened representation.
|
||||
"""
|
||||
|
||||
@property
|
||||
def patch_size(self) -> Tuple[int, int, int]:
|
||||
...
|
||||
"""
|
||||
Returns the patch size as a tuple of (temporal, height, width) dimensions
|
||||
"""
|
||||
|
||||
def get_patch_grid_bounds(
|
||||
self,
|
||||
output_shape: AudioLatentShape | VideoLatentShape,
|
||||
device: torch.device | None = None,
|
||||
) -> torch.Tensor:
|
||||
...
|
||||
"""
|
||||
Compute metadata describing where each latent patch resides within the
|
||||
grid specified by `output_shape`.
|
||||
Args:
|
||||
output_shape: Target grid layout for the patches.
|
||||
device: Target device for the returned tensor.
|
||||
Returns:
|
||||
Tensor containing patch coordinate metadata such as spatial or temporal intervals.
|
||||
"""
|
||||
|
||||
|
||||
class SchedulerProtocol(Protocol):
|
||||
"""
|
||||
Protocol for schedulers that provide a sigmas schedule tensor for a
|
||||
given number of steps. Device is cpu.
|
||||
"""
|
||||
|
||||
def execute(self, steps: int, **kwargs) -> torch.FloatTensor: ...
|
||||
|
||||
|
||||
class GuiderProtocol(Protocol):
|
||||
"""
|
||||
Protocol for guiders that compute a delta tensor given conditioning inputs.
|
||||
The returned delta should be added to the conditional output (cond), enabling
|
||||
multiple guiders to be chained together by accumulating their deltas.
|
||||
"""
|
||||
|
||||
scale: float
|
||||
|
||||
def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor: ...
|
||||
|
||||
def enabled(self) -> bool:
|
||||
"""
|
||||
Returns whether the corresponding perturbation is enabled. E.g. for CFG, this should return False if the scale
|
||||
is 1.0.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class DiffusionStepProtocol(Protocol):
|
||||
"""
|
||||
Protocol for diffusion steps that provide a next sample tensor for a given current sample tensor,
|
||||
current denoised sample tensor, and sigmas tensor.
|
||||
"""
|
||||
|
||||
def step(
|
||||
self, sample: torch.Tensor, denoised_sample: torch.Tensor, sigmas: torch.Tensor, step_index: int
|
||||
) -> torch.Tensor: ...
|
||||
@@ -0,0 +1,129 @@
|
||||
import math
|
||||
from functools import lru_cache
|
||||
|
||||
import numpy
|
||||
import scipy
|
||||
import torch
|
||||
|
||||
from ltx_core.components.protocols import SchedulerProtocol
|
||||
|
||||
BASE_SHIFT_ANCHOR = 1024
|
||||
MAX_SHIFT_ANCHOR = 4096
|
||||
|
||||
|
||||
class LTX2Scheduler(SchedulerProtocol):
|
||||
"""
|
||||
Default scheduler for LTX-2 diffusion sampling.
|
||||
Generates a sigma schedule with token-count-dependent shifting and optional
|
||||
stretching to a terminal value.
|
||||
"""
|
||||
|
||||
def execute(
|
||||
self,
|
||||
steps: int,
|
||||
latent: torch.Tensor | None = None,
|
||||
max_shift: float = 2.05,
|
||||
base_shift: float = 0.95,
|
||||
stretch: bool = True,
|
||||
terminal: float = 0.1,
|
||||
**_kwargs,
|
||||
) -> torch.FloatTensor:
|
||||
tokens = math.prod(latent.shape[2:]) if latent is not None else MAX_SHIFT_ANCHOR
|
||||
sigmas = torch.linspace(1.0, 0.0, steps + 1)
|
||||
|
||||
x1 = BASE_SHIFT_ANCHOR
|
||||
x2 = MAX_SHIFT_ANCHOR
|
||||
mm = (max_shift - base_shift) / (x2 - x1)
|
||||
b = base_shift - mm * x1
|
||||
sigma_shift = (tokens) * mm + b
|
||||
|
||||
power = 1
|
||||
sigmas = torch.where(
|
||||
sigmas != 0,
|
||||
math.exp(sigma_shift) / (math.exp(sigma_shift) + (1 / sigmas - 1) ** power),
|
||||
0,
|
||||
)
|
||||
|
||||
# Stretch sigmas so that its final value matches the given terminal value.
|
||||
if stretch:
|
||||
non_zero_mask = sigmas != 0
|
||||
non_zero_sigmas = sigmas[non_zero_mask]
|
||||
one_minus_z = 1.0 - non_zero_sigmas
|
||||
scale_factor = one_minus_z[-1] / (1.0 - terminal)
|
||||
stretched = 1.0 - (one_minus_z / scale_factor)
|
||||
sigmas[non_zero_mask] = stretched
|
||||
|
||||
return sigmas.to(torch.float32)
|
||||
|
||||
|
||||
class LinearQuadraticScheduler(SchedulerProtocol):
|
||||
"""
|
||||
Scheduler with linear steps followed by quadratic steps.
|
||||
Produces a sigma schedule that transitions linearly up to a threshold,
|
||||
then follows a quadratic curve for the remaining steps.
|
||||
"""
|
||||
|
||||
def execute(
|
||||
self, steps: int, threshold_noise: float = 0.025, linear_steps: int | None = None, **_kwargs
|
||||
) -> torch.FloatTensor:
|
||||
if steps == 1:
|
||||
return torch.FloatTensor([1.0, 0.0])
|
||||
|
||||
if linear_steps is None:
|
||||
linear_steps = steps // 2
|
||||
linear_sigma_schedule = [i * threshold_noise / linear_steps for i in range(linear_steps)]
|
||||
threshold_noise_step_diff = linear_steps - threshold_noise * steps
|
||||
quadratic_steps = steps - linear_steps
|
||||
quadratic_sigma_schedule = []
|
||||
if quadratic_steps > 0:
|
||||
quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2)
|
||||
linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps**2)
|
||||
const = quadratic_coef * (linear_steps**2)
|
||||
quadratic_sigma_schedule = [
|
||||
quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, steps)
|
||||
]
|
||||
sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule + [1.0]
|
||||
sigma_schedule = [1.0 - x for x in sigma_schedule]
|
||||
return torch.FloatTensor(sigma_schedule)
|
||||
|
||||
|
||||
class BetaScheduler(SchedulerProtocol):
|
||||
"""
|
||||
Scheduler using a beta distribution to sample timesteps.
|
||||
Based on: https://arxiv.org/abs/2407.12173
|
||||
"""
|
||||
|
||||
shift = 2.37
|
||||
timesteps_length = 10000
|
||||
|
||||
def execute(self, steps: int, alpha: float = 0.6, beta: float = 0.6) -> torch.FloatTensor:
|
||||
"""
|
||||
Execute the beta scheduler.
|
||||
Args:
|
||||
steps: The number of steps to execute the scheduler for.
|
||||
alpha: The alpha parameter for the beta distribution.
|
||||
beta: The beta parameter for the beta distribution.
|
||||
Warnings:
|
||||
The number of steps within `sigmas` theoretically might be less than `steps+1`,
|
||||
because of the deduplication of the identical timesteps
|
||||
Returns:
|
||||
A tensor of sigmas.
|
||||
"""
|
||||
model_sampling_sigmas = _precalculate_model_sampling_sigmas(self.shift, self.timesteps_length)
|
||||
total_timesteps = len(model_sampling_sigmas) - 1
|
||||
ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
|
||||
ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps).tolist()
|
||||
ts = list(dict.fromkeys(ts))
|
||||
|
||||
sigmas = [float(model_sampling_sigmas[int(t)]) for t in ts] + [0.0]
|
||||
return torch.FloatTensor(sigmas)
|
||||
|
||||
|
||||
@lru_cache(maxsize=5)
|
||||
def _precalculate_model_sampling_sigmas(shift: float, timesteps_length: int) -> torch.Tensor:
|
||||
timesteps = torch.arange(1, timesteps_length + 1, 1) / timesteps_length
|
||||
return torch.Tensor([flux_time_shift(shift, 1.0, t) for t in timesteps])
|
||||
|
||||
|
||||
def flux_time_shift(mu: float, sigma: float, t: float) -> float:
|
||||
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
|
||||
@@ -0,0 +1,12 @@
|
||||
"""Conditioning utilities: latent state, tools, and conditioning types."""
|
||||
|
||||
from ltx_core.conditioning.exceptions import ConditioningError
|
||||
from ltx_core.conditioning.item import ConditioningItem
|
||||
from ltx_core.conditioning.types import VideoConditionByKeyframeIndex, VideoConditionByLatentIndex
|
||||
|
||||
__all__ = [
|
||||
"ConditioningError",
|
||||
"ConditioningItem",
|
||||
"VideoConditionByKeyframeIndex",
|
||||
"VideoConditionByLatentIndex",
|
||||
]
|
||||
@@ -0,0 +1,4 @@
|
||||
class ConditioningError(Exception):
|
||||
"""
|
||||
Class for conditioning-related errors.
|
||||
"""
|
||||
@@ -0,0 +1,20 @@
|
||||
from typing import Protocol
|
||||
|
||||
from ltx_core.tools import LatentTools
|
||||
from ltx_core.types import LatentState
|
||||
|
||||
|
||||
class ConditioningItem(Protocol):
|
||||
"""Protocol for conditioning items that modify latent state during diffusion."""
|
||||
|
||||
def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
|
||||
"""
|
||||
Apply the conditioning to the latent state.
|
||||
Args:
|
||||
latent_state: The latent state to apply the conditioning to. This is state always patchified.
|
||||
Returns:
|
||||
The latent state after the conditioning has been applied.
|
||||
IMPORTANT: If the conditioning needs to add extra tokens to the latent, it should add them to the end of the
|
||||
latent.
|
||||
"""
|
||||
...
|
||||
@@ -0,0 +1,9 @@
|
||||
"""Conditioning type implementations."""
|
||||
|
||||
from ltx_core.conditioning.types.keyframe_cond import VideoConditionByKeyframeIndex
|
||||
from ltx_core.conditioning.types.latent_cond import VideoConditionByLatentIndex
|
||||
|
||||
__all__ = [
|
||||
"VideoConditionByKeyframeIndex",
|
||||
"VideoConditionByLatentIndex",
|
||||
]
|
||||
@@ -0,0 +1,53 @@
|
||||
import torch
|
||||
|
||||
from ltx_core.components.patchifiers import get_pixel_coords
|
||||
from ltx_core.conditioning.item import ConditioningItem
|
||||
from ltx_core.tools import VideoLatentTools
|
||||
from ltx_core.types import LatentState, VideoLatentShape
|
||||
|
||||
|
||||
class VideoConditionByKeyframeIndex(ConditioningItem):
|
||||
"""
|
||||
Conditions video generation on keyframe latents at a specific frame index.
|
||||
Appends keyframe tokens to the latent state with positions offset by frame_idx,
|
||||
and sets denoise strength according to the strength parameter.
|
||||
"""
|
||||
|
||||
def __init__(self, keyframes: torch.Tensor, frame_idx: int, strength: float):
|
||||
self.keyframes = keyframes
|
||||
self.frame_idx = frame_idx
|
||||
self.strength = strength
|
||||
|
||||
def apply_to(
|
||||
self,
|
||||
latent_state: LatentState,
|
||||
latent_tools: VideoLatentTools,
|
||||
) -> LatentState:
|
||||
tokens = latent_tools.patchifier.patchify(self.keyframes)
|
||||
latent_coords = latent_tools.patchifier.get_patch_grid_bounds(
|
||||
output_shape=VideoLatentShape.from_torch_shape(self.keyframes.shape),
|
||||
device=self.keyframes.device,
|
||||
)
|
||||
positions = get_pixel_coords(
|
||||
latent_coords=latent_coords,
|
||||
scale_factors=latent_tools.scale_factors,
|
||||
causal_fix=latent_tools.causal_fix if self.frame_idx == 0 else False,
|
||||
)
|
||||
|
||||
positions[:, 0, ...] += self.frame_idx
|
||||
positions = positions.to(dtype=torch.float32)
|
||||
positions[:, 0, ...] /= latent_tools.fps
|
||||
|
||||
denoise_mask = torch.full(
|
||||
size=(*tokens.shape[:2], 1),
|
||||
fill_value=1.0 - self.strength,
|
||||
device=self.keyframes.device,
|
||||
dtype=self.keyframes.dtype,
|
||||
)
|
||||
|
||||
return LatentState(
|
||||
latent=torch.cat([latent_state.latent, tokens], dim=1),
|
||||
denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
|
||||
positions=torch.cat([latent_state.positions, positions], dim=2),
|
||||
clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
|
||||
)
|
||||
@@ -0,0 +1,44 @@
|
||||
import torch
|
||||
|
||||
from ltx_core.conditioning.exceptions import ConditioningError
|
||||
from ltx_core.conditioning.item import ConditioningItem
|
||||
from ltx_core.tools import LatentTools
|
||||
from ltx_core.types import LatentState
|
||||
|
||||
|
||||
class VideoConditionByLatentIndex(ConditioningItem):
|
||||
"""
|
||||
Conditions video generation by injecting latents at a specific latent frame index.
|
||||
Replaces tokens in the latent state at positions corresponding to latent_idx,
|
||||
and sets denoise strength according to the strength parameter.
|
||||
"""
|
||||
|
||||
def __init__(self, latent: torch.Tensor, strength: float, latent_idx: int):
|
||||
self.latent = latent
|
||||
self.strength = strength
|
||||
self.latent_idx = latent_idx
|
||||
|
||||
def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
|
||||
cond_batch, cond_channels, _, cond_height, cond_width = self.latent.shape
|
||||
tgt_batch, tgt_channels, tgt_frames, tgt_height, tgt_width = latent_tools.target_shape.to_torch_shape()
|
||||
|
||||
if (cond_batch, cond_channels, cond_height, cond_width) != (tgt_batch, tgt_channels, tgt_height, tgt_width):
|
||||
raise ConditioningError(
|
||||
f"Can't apply image conditioning item to latent with shape {latent_tools.target_shape}, expected "
|
||||
f"shape is ({tgt_batch}, {tgt_channels}, {tgt_frames}, {tgt_height}, {tgt_width}). Make sure "
|
||||
"the image and latent have the same spatial shape."
|
||||
)
|
||||
|
||||
tokens = latent_tools.patchifier.patchify(self.latent)
|
||||
start_token = latent_tools.patchifier.get_token_count(
|
||||
latent_tools.target_shape._replace(frames=self.latent_idx)
|
||||
)
|
||||
stop_token = start_token + tokens.shape[1]
|
||||
|
||||
latent_state = latent_state.clone()
|
||||
|
||||
latent_state.latent[:, start_token:stop_token] = tokens
|
||||
latent_state.clean_latent[:, start_token:stop_token] = tokens
|
||||
latent_state.denoise_mask[:, start_token:stop_token] = 1.0 - self.strength
|
||||
|
||||
return latent_state
|
||||
@@ -0,0 +1,15 @@
|
||||
"""Guidance and perturbation utilities for attention manipulation."""
|
||||
|
||||
from ltx_core.guidance.perturbations import (
|
||||
BatchedPerturbationConfig,
|
||||
Perturbation,
|
||||
PerturbationConfig,
|
||||
PerturbationType,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"BatchedPerturbationConfig",
|
||||
"Perturbation",
|
||||
"PerturbationConfig",
|
||||
"PerturbationType",
|
||||
]
|
||||
@@ -0,0 +1,79 @@
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
|
||||
import torch
|
||||
from torch._prims_common import DeviceLikeType
|
||||
|
||||
|
||||
class PerturbationType(Enum):
|
||||
"""Types of attention perturbations for STG (Spatio-Temporal Guidance)."""
|
||||
|
||||
SKIP_A2V_CROSS_ATTN = "skip_a2v_cross_attn"
|
||||
SKIP_V2A_CROSS_ATTN = "skip_v2a_cross_attn"
|
||||
SKIP_VIDEO_SELF_ATTN = "skip_video_self_attn"
|
||||
SKIP_AUDIO_SELF_ATTN = "skip_audio_self_attn"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Perturbation:
|
||||
"""A single perturbation specifying which attention type to skip and in which blocks."""
|
||||
|
||||
type: PerturbationType
|
||||
blocks: list[int] | None # None means all blocks
|
||||
|
||||
def is_perturbed(self, perturbation_type: PerturbationType, block: int) -> bool:
|
||||
if self.type != perturbation_type:
|
||||
return False
|
||||
|
||||
if self.blocks is None:
|
||||
return True
|
||||
|
||||
return block in self.blocks
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PerturbationConfig:
|
||||
"""Configuration holding a list of perturbations for a single sample."""
|
||||
|
||||
perturbations: list[Perturbation] | None
|
||||
|
||||
def is_perturbed(self, perturbation_type: PerturbationType, block: int) -> bool:
|
||||
if self.perturbations is None:
|
||||
return False
|
||||
|
||||
return any(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
|
||||
|
||||
@staticmethod
|
||||
def empty() -> "PerturbationConfig":
|
||||
return PerturbationConfig([])
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BatchedPerturbationConfig:
|
||||
"""Perturbation configurations for a batch, with utilities for generating attention masks."""
|
||||
|
||||
perturbations: list[PerturbationConfig]
|
||||
|
||||
def mask(
|
||||
self, perturbation_type: PerturbationType, block: int, device: DeviceLikeType, dtype: torch.dtype
|
||||
) -> torch.Tensor:
|
||||
mask = torch.ones((len(self.perturbations),), device=device, dtype=dtype)
|
||||
for batch_idx, perturbation in enumerate(self.perturbations):
|
||||
if perturbation.is_perturbed(perturbation_type, block):
|
||||
mask[batch_idx] = 0
|
||||
|
||||
return mask
|
||||
|
||||
def mask_like(self, perturbation_type: PerturbationType, block: int, values: torch.Tensor) -> torch.Tensor:
|
||||
mask = self.mask(perturbation_type, block, values.device, values.dtype)
|
||||
return mask.view(mask.numel(), *([1] * len(values.shape[1:])))
|
||||
|
||||
def any_in_batch(self, perturbation_type: PerturbationType, block: int) -> bool:
|
||||
return any(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
|
||||
|
||||
def all_in_batch(self, perturbation_type: PerturbationType, block: int) -> bool:
|
||||
return all(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
|
||||
|
||||
@staticmethod
|
||||
def empty(batch_size: int) -> "BatchedPerturbationConfig":
|
||||
return BatchedPerturbationConfig([PerturbationConfig.empty() for _ in range(batch_size)])
|
||||
@@ -0,0 +1,48 @@
|
||||
"""Loader utilities for model weights, LoRAs, and safetensor operations."""
|
||||
|
||||
from ltx_core.loader.fuse_loras import apply_loras
|
||||
from ltx_core.loader.module_ops import ModuleOps
|
||||
from ltx_core.loader.primitives import (
|
||||
LoRAAdaptableProtocol,
|
||||
LoraPathStrengthAndSDOps,
|
||||
LoraStateDictWithStrength,
|
||||
ModelBuilderProtocol,
|
||||
StateDict,
|
||||
StateDictLoader,
|
||||
)
|
||||
from ltx_core.loader.registry import DummyRegistry, Registry, StateDictRegistry
|
||||
from ltx_core.loader.sd_ops import (
|
||||
LTXV_LORA_COMFY_RENAMING_MAP,
|
||||
ContentMatching,
|
||||
ContentReplacement,
|
||||
KeyValueOperation,
|
||||
KeyValueOperationResult,
|
||||
SDKeyValueOperation,
|
||||
SDOps,
|
||||
)
|
||||
from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader, SafetensorsStateDictLoader
|
||||
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
|
||||
|
||||
__all__ = [
|
||||
"LTXV_LORA_COMFY_RENAMING_MAP",
|
||||
"ContentMatching",
|
||||
"ContentReplacement",
|
||||
"DummyRegistry",
|
||||
"KeyValueOperation",
|
||||
"KeyValueOperationResult",
|
||||
"LoRAAdaptableProtocol",
|
||||
"LoraPathStrengthAndSDOps",
|
||||
"LoraStateDictWithStrength",
|
||||
"ModelBuilderProtocol",
|
||||
"ModuleOps",
|
||||
"Registry",
|
||||
"SDKeyValueOperation",
|
||||
"SDOps",
|
||||
"SafetensorsModelStateDictLoader",
|
||||
"SafetensorsStateDictLoader",
|
||||
"SingleGPUModelBuilder",
|
||||
"StateDict",
|
||||
"StateDictLoader",
|
||||
"StateDictRegistry",
|
||||
"apply_loras",
|
||||
]
|
||||
@@ -0,0 +1,100 @@
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from ltx_core.loader.kernels import fused_add_round_kernel
|
||||
from ltx_core.loader.primitives import LoraStateDictWithStrength, StateDict
|
||||
|
||||
BLOCK_SIZE = 1024
|
||||
|
||||
|
||||
def fused_add_round_launch(target_weight: torch.Tensor, original_weight: torch.Tensor, seed: int) -> torch.Tensor:
|
||||
if original_weight.dtype == torch.float8_e4m3fn:
|
||||
exponent_bits, mantissa_bits, exponent_bias = 4, 3, 7
|
||||
elif original_weight.dtype == torch.float8_e5m2:
|
||||
exponent_bits, mantissa_bits, exponent_bias = 5, 2, 15 # noqa: F841
|
||||
else:
|
||||
raise ValueError("Unsupported dtype")
|
||||
|
||||
if target_weight.dtype != torch.bfloat16:
|
||||
raise ValueError("target_weight dtype must be bfloat16")
|
||||
|
||||
# Calculate grid and block sizes
|
||||
n_elements = original_weight.numel()
|
||||
grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
|
||||
|
||||
# Launch kernel
|
||||
fused_add_round_kernel[grid](
|
||||
original_weight,
|
||||
target_weight,
|
||||
seed,
|
||||
n_elements,
|
||||
exponent_bias,
|
||||
mantissa_bits,
|
||||
BLOCK_SIZE,
|
||||
)
|
||||
return target_weight
|
||||
|
||||
|
||||
def calculate_weight_float8_(target_weights: torch.Tensor, original_weights: torch.Tensor) -> torch.Tensor:
|
||||
result = fused_add_round_launch(target_weights, original_weights, seed=0).to(target_weights.dtype)
|
||||
target_weights.copy_(result, non_blocking=True)
|
||||
return target_weights
|
||||
|
||||
|
||||
def _prepare_deltas(
|
||||
lora_sd_and_strengths: list[LoraStateDictWithStrength], key: str, dtype: torch.dtype, device: torch.device
|
||||
) -> torch.Tensor | None:
|
||||
deltas = []
|
||||
prefix = key[: -len(".weight")]
|
||||
key_a = f"{prefix}.lora_A.weight"
|
||||
key_b = f"{prefix}.lora_B.weight"
|
||||
for lsd, coef in lora_sd_and_strengths:
|
||||
if key_a not in lsd.sd or key_b not in lsd.sd:
|
||||
continue
|
||||
product = torch.matmul(lsd.sd[key_b] * coef, lsd.sd[key_a])
|
||||
deltas.append(product.to(dtype=dtype, device=device))
|
||||
if len(deltas) == 0:
|
||||
return None
|
||||
elif len(deltas) == 1:
|
||||
return deltas[0]
|
||||
return torch.sum(torch.stack(deltas, dim=0), dim=0)
|
||||
|
||||
|
||||
def apply_loras(
|
||||
model_sd: StateDict,
|
||||
lora_sd_and_strengths: list[LoraStateDictWithStrength],
|
||||
dtype: torch.dtype,
|
||||
destination_sd: StateDict | None = None,
|
||||
) -> StateDict:
|
||||
sd = {}
|
||||
if destination_sd is not None:
|
||||
sd = destination_sd.sd
|
||||
size = 0
|
||||
device = torch.device("meta")
|
||||
inner_dtypes = set()
|
||||
for key, weight in model_sd.sd.items():
|
||||
if weight is None:
|
||||
continue
|
||||
device = weight.device
|
||||
target_dtype = dtype if dtype is not None else weight.dtype
|
||||
deltas_dtype = target_dtype if target_dtype not in [torch.float8_e4m3fn, torch.float8_e5m2] else torch.bfloat16
|
||||
deltas = _prepare_deltas(lora_sd_and_strengths, key, deltas_dtype, device)
|
||||
if deltas is None:
|
||||
if key in sd:
|
||||
continue
|
||||
deltas = weight.clone().to(dtype=target_dtype, device=device)
|
||||
elif weight.dtype == torch.float8_e4m3fn:
|
||||
if str(device).startswith("cuda"):
|
||||
deltas = calculate_weight_float8_(deltas, weight)
|
||||
else:
|
||||
deltas.add_(weight.to(dtype=deltas.dtype, device=device))
|
||||
elif weight.dtype == torch.bfloat16:
|
||||
deltas.add_(weight)
|
||||
else:
|
||||
raise ValueError(f"Unsupported dtype: {weight.dtype}")
|
||||
sd[key] = deltas.to(dtype=target_dtype)
|
||||
inner_dtypes.add(target_dtype)
|
||||
size += deltas.nbytes
|
||||
if destination_sd is not None:
|
||||
return destination_sd
|
||||
return StateDict(sd, device, size, inner_dtypes)
|
||||
@@ -0,0 +1,72 @@
|
||||
# ruff: noqa: ANN001, ANN201, ERA001, N803, N806
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
|
||||
@triton.jit
|
||||
def fused_add_round_kernel(
|
||||
x_ptr,
|
||||
output_ptr, # contents will be added to the output
|
||||
seed,
|
||||
n_elements,
|
||||
EXPONENT_BIAS,
|
||||
MANTISSA_BITS,
|
||||
BLOCK_SIZE: tl.constexpr,
|
||||
):
|
||||
"""
|
||||
A kernel to upcast 8bit quantized weights to bfloat16 with stochastic rounding
|
||||
and add them to bfloat16 output weights. Might be used to upcast original model weights
|
||||
and to further add them to precalculated deltas coming from LoRAs.
|
||||
"""
|
||||
# Get program ID and compute offsets
|
||||
pid = tl.program_id(axis=0)
|
||||
block_start = pid * BLOCK_SIZE
|
||||
offsets = block_start + tl.arange(0, BLOCK_SIZE)
|
||||
mask = offsets < n_elements
|
||||
|
||||
# Load data
|
||||
x = tl.load(x_ptr + offsets, mask=mask)
|
||||
rand_vals = tl.rand(seed, offsets) - 0.5
|
||||
|
||||
x = tl.cast(x, tl.float16)
|
||||
delta = tl.load(output_ptr + offsets, mask=mask)
|
||||
delta = tl.cast(delta, tl.float16)
|
||||
x = x + delta
|
||||
|
||||
x_bits = tl.cast(x, tl.int16, bitcast=True)
|
||||
|
||||
# Calculate the exponent. Unbiased fp16 exponent is ((x_bits & 0x7C00) >> 10) - 15 for
|
||||
# normal numbers and -14 for subnormals.
|
||||
fp16_exponent_bits = (x_bits & 0x7C00) >> 10
|
||||
fp16_normals = fp16_exponent_bits > 0
|
||||
fp16_exponent = tl.where(fp16_normals, fp16_exponent_bits - 15, -14)
|
||||
|
||||
# Add the target dtype's exponent bias and clamp to the target dtype's exponent range.
|
||||
exponent = fp16_exponent + EXPONENT_BIAS
|
||||
MAX_EXPONENT = 2 * EXPONENT_BIAS + 1
|
||||
exponent = tl.where(exponent > MAX_EXPONENT, MAX_EXPONENT, exponent)
|
||||
exponent = tl.where(exponent < 0, 0, exponent)
|
||||
|
||||
# Normal ULP exponent, expressed as an fp16 exponent field:
|
||||
# (exponent - EXPONENT_BIAS - MANTISSA_BITS) + 15
|
||||
# Simplifies to: fp16_exponent - MANTISSA_BITS + 15
|
||||
# See https://en.wikipedia.org/wiki/Unit_in_the_last_place
|
||||
eps_exp = tl.maximum(0, tl.minimum(31, exponent - EXPONENT_BIAS - MANTISSA_BITS + 15))
|
||||
|
||||
# Calculate epsilon in the target dtype
|
||||
eps_normal = tl.cast(tl.cast(eps_exp << 10, tl.int16), tl.float16, bitcast=True)
|
||||
|
||||
# Subnormal ULP: 2^(1 - EXPONENT_BIAS - MANTISSA_BITS) ->
|
||||
# fp16 exponent bits: (1 - EXPONENT_BIAS - MANTISSA_BITS) + 15 =
|
||||
# 16 - EXPONENT_BIAS - MANTISSA_BITS
|
||||
eps_subnormal = tl.cast(tl.cast((16 - EXPONENT_BIAS - MANTISSA_BITS) << 10, tl.int16), tl.float16, bitcast=True)
|
||||
eps = tl.where(exponent > 0, eps_normal, eps_subnormal)
|
||||
|
||||
# Apply zero mask to epsilon
|
||||
eps = tl.where(x == 0, 0.0, eps)
|
||||
|
||||
# Apply stochastic rounding
|
||||
output = tl.cast(x + rand_vals * eps, tl.bfloat16)
|
||||
|
||||
# Store the result
|
||||
tl.store(output_ptr + offsets, output, mask=mask)
|
||||
@@ -0,0 +1,14 @@
|
||||
from typing import Callable, NamedTuple
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class ModuleOps(NamedTuple):
|
||||
"""
|
||||
Defines a named operation for matching and mutating PyTorch modules.
|
||||
Used to selectively transform modules in a model (e.g., replacing layers with quantized versions).
|
||||
"""
|
||||
|
||||
name: str
|
||||
matcher: Callable[[torch.nn.Module], bool]
|
||||
mutator: Callable[[torch.nn.Module], torch.nn.Module]
|
||||
@@ -0,0 +1,109 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import NamedTuple, Protocol
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.loader.module_ops import ModuleOps
|
||||
from ltx_core.loader.sd_ops import SDOps
|
||||
from ltx_core.model.model_protocol import ModelType
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StateDict:
|
||||
"""
|
||||
Immutable container for a PyTorch state dictionary.
|
||||
Contains:
|
||||
- sd: Dictionary of tensors (weights, buffers, etc.)
|
||||
- device: Device where tensors are stored
|
||||
- size: Total memory footprint in bytes
|
||||
- dtype: Set of tensor dtypes present
|
||||
"""
|
||||
|
||||
sd: dict
|
||||
device: torch.device
|
||||
size: int
|
||||
dtype: set[torch.dtype]
|
||||
|
||||
def footprint(self) -> tuple[int, torch.device]:
|
||||
return self.size, self.device
|
||||
|
||||
|
||||
class StateDictLoader(Protocol):
|
||||
"""
|
||||
Protocol for loading state dictionaries from various sources.
|
||||
Implementations must provide:
|
||||
- metadata: Extract model metadata from a single path
|
||||
- load: Load state dict from path(s) and apply SDOps transformations
|
||||
"""
|
||||
|
||||
def metadata(self, path: str) -> dict:
|
||||
"""
|
||||
Load metadata from path
|
||||
"""
|
||||
|
||||
def load(self, path: str | list[str], sd_ops: SDOps | None = None, device: torch.device | None = None) -> StateDict:
|
||||
"""
|
||||
Load state dict from path or paths (for sharded model storage) and apply sd_ops
|
||||
"""
|
||||
|
||||
|
||||
class ModelBuilderProtocol(Protocol[ModelType]):
|
||||
"""
|
||||
Protocol for building PyTorch models from configuration dictionaries.
|
||||
Implementations must provide:
|
||||
- meta_model: Create a model from configuration dictionary and apply module operations
|
||||
- build: Create and initialize a model from state dictionary and apply dtype transformations
|
||||
"""
|
||||
|
||||
def meta_model(self, config: dict, module_ops: list[ModuleOps] | None = None) -> ModelType:
|
||||
"""
|
||||
Create a model on the meta device from a configuration dictionary.
|
||||
This decouples model creation from weight loading, allowing the model
|
||||
architecture to be instantiated without allocating memory for parameters.
|
||||
Args:
|
||||
config: Model configuration dictionary.
|
||||
module_ops: Optional list of module operations to apply (e.g., quantization).
|
||||
Returns:
|
||||
Model instance on meta device (no actual memory allocated for parameters).
|
||||
"""
|
||||
...
|
||||
|
||||
def build(self, dtype: torch.dtype | None = None) -> ModelType:
|
||||
"""
|
||||
Build the model
|
||||
Args:
|
||||
dtype: Target dtype for the model, if None, uses the dtype of the model_path model
|
||||
Returns:
|
||||
Model instance
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class LoRAAdaptableProtocol(Protocol):
|
||||
"""
|
||||
Protocol for models that can be adapted with LoRAs.
|
||||
Implementations must provide:
|
||||
- lora: Add a LoRA to the model
|
||||
"""
|
||||
|
||||
def lora(self, lora_path: str, strength: float) -> "LoRAAdaptableProtocol":
|
||||
pass
|
||||
|
||||
|
||||
class LoraPathStrengthAndSDOps(NamedTuple):
|
||||
"""
|
||||
Tuple containing a LoRA path, strength, and SDOps for applying to the LoRA state dict.
|
||||
"""
|
||||
|
||||
path: str
|
||||
strength: float
|
||||
sd_ops: SDOps
|
||||
|
||||
|
||||
class LoraStateDictWithStrength(NamedTuple):
|
||||
"""
|
||||
Tuple containing a LoRA state dict and strength for applying to the model.
|
||||
"""
|
||||
|
||||
state_dict: StateDict
|
||||
strength: float
|
||||
@@ -0,0 +1,84 @@
|
||||
import hashlib
|
||||
import threading
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Protocol
|
||||
|
||||
from ltx_core.loader.primitives import StateDict
|
||||
from ltx_core.loader.sd_ops import SDOps
|
||||
|
||||
|
||||
class Registry(Protocol):
|
||||
"""
|
||||
Protocol for managing state dictionaries in a registry.
|
||||
It is used to store state dictionaries and reuse them later without loading them again.
|
||||
Implementations must provide:
|
||||
- add: Add a state dictionary to the registry
|
||||
- pop: Remove a state dictionary from the registry
|
||||
- get: Retrieve a state dictionary from the registry
|
||||
- clear: Clear all state dictionaries from the registry
|
||||
"""
|
||||
|
||||
def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> None: ...
|
||||
|
||||
def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: ...
|
||||
|
||||
def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: ...
|
||||
|
||||
def clear(self) -> None: ...
|
||||
|
||||
|
||||
class DummyRegistry(Registry):
|
||||
"""
|
||||
Dummy registry that does not store state dictionaries.
|
||||
"""
|
||||
|
||||
def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> None:
|
||||
pass
|
||||
|
||||
def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
|
||||
pass
|
||||
|
||||
def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
|
||||
pass
|
||||
|
||||
def clear(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class StateDictRegistry(Registry):
|
||||
"""
|
||||
Registry that stores state dictionaries in a dictionary.
|
||||
"""
|
||||
|
||||
_state_dicts: dict[str, StateDict] = field(default_factory=dict)
|
||||
_lock: threading.Lock = field(default_factory=threading.Lock)
|
||||
|
||||
def _generate_id(self, paths: list[str], sd_ops: SDOps) -> str:
|
||||
m = hashlib.sha256()
|
||||
parts = [str(Path(p).resolve()) for p in paths]
|
||||
if sd_ops is not None:
|
||||
parts.append(sd_ops.name)
|
||||
m.update("\0".join(parts).encode("utf-8"))
|
||||
return m.hexdigest()
|
||||
|
||||
def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> str:
|
||||
sd_id = self._generate_id(paths, sd_ops)
|
||||
with self._lock:
|
||||
if sd_id in self._state_dicts:
|
||||
raise ValueError(f"State dict retrieved from {paths} with {sd_ops} already added, check with get first")
|
||||
self._state_dicts[sd_id] = state_dict
|
||||
return sd_id
|
||||
|
||||
def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
|
||||
with self._lock:
|
||||
return self._state_dicts.pop(self._generate_id(paths, sd_ops), None)
|
||||
|
||||
def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
|
||||
with self._lock:
|
||||
return self._state_dicts.get(self._generate_id(paths, sd_ops), None)
|
||||
|
||||
def clear(self) -> None:
|
||||
with self._lock:
|
||||
self._state_dicts.clear()
|
||||
@@ -0,0 +1,127 @@
|
||||
from dataclasses import dataclass, replace
|
||||
from typing import NamedTuple, Protocol
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ContentReplacement:
|
||||
"""
|
||||
Represents a content replacement operation.
|
||||
Used to replace a specific content with a replacement in a state dict key.
|
||||
"""
|
||||
|
||||
content: str
|
||||
replacement: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ContentMatching:
|
||||
"""
|
||||
Represents a content matching operation.
|
||||
Used to match a specific prefix and suffix in a state dict key.
|
||||
"""
|
||||
|
||||
prefix: str = ""
|
||||
suffix: str = ""
|
||||
|
||||
|
||||
class KeyValueOperationResult(NamedTuple):
|
||||
"""
|
||||
Represents the result of a key-value operation.
|
||||
Contains the new key and value after the operation has been applied.
|
||||
"""
|
||||
|
||||
new_key: str
|
||||
new_value: torch.Tensor
|
||||
|
||||
|
||||
class KeyValueOperation(Protocol):
|
||||
"""
|
||||
Protocol for key-value operations.
|
||||
Used to apply operations to a specific key and value in a state dict.
|
||||
"""
|
||||
|
||||
def __call__(self, tensor_key: str, tensor_value: torch.Tensor) -> list[KeyValueOperationResult]: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SDKeyValueOperation:
|
||||
"""
|
||||
Represents a key-value operation.
|
||||
Used to apply operations to a specific key and value in a state dict.
|
||||
"""
|
||||
|
||||
key_matcher: ContentMatching
|
||||
kv_operation: KeyValueOperation
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SDOps:
|
||||
"""Immutable class representing state dict key operations."""
|
||||
|
||||
name: str
|
||||
mapping: tuple[
|
||||
ContentReplacement | ContentMatching | SDKeyValueOperation, ...
|
||||
] = () # Immutable tuple of (key, value) pairs
|
||||
|
||||
def with_replacement(self, content: str, replacement: str) -> "SDOps":
|
||||
"""Create a new SDOps instance with the specified replacement added to the mapping."""
|
||||
|
||||
new_mapping = (*self.mapping, ContentReplacement(content, replacement))
|
||||
return replace(self, mapping=new_mapping)
|
||||
|
||||
def with_matching(self, prefix: str = "", suffix: str = "") -> "SDOps":
|
||||
"""Create a new SDOps instance with the specified prefix and suffix matching added to the mapping."""
|
||||
|
||||
new_mapping = (*self.mapping, ContentMatching(prefix, suffix))
|
||||
return replace(self, mapping=new_mapping)
|
||||
|
||||
def with_kv_operation(
|
||||
self,
|
||||
operation: KeyValueOperation,
|
||||
key_prefix: str = "",
|
||||
key_suffix: str = "",
|
||||
) -> "SDOps":
|
||||
"""Create a new SDOps instance with the specified value operation added to the mapping."""
|
||||
key_matcher = ContentMatching(key_prefix, key_suffix)
|
||||
sd_kv_operation = SDKeyValueOperation(key_matcher, operation)
|
||||
new_mapping = (*self.mapping, sd_kv_operation)
|
||||
return replace(self, mapping=new_mapping)
|
||||
|
||||
def apply_to_key(self, key: str) -> str | None:
|
||||
"""Apply the mapping to the given name."""
|
||||
matchers = [content for content in self.mapping if isinstance(content, ContentMatching)]
|
||||
valid = any(key.startswith(f.prefix) and key.endswith(f.suffix) for f in matchers)
|
||||
if not valid:
|
||||
return None
|
||||
|
||||
for replacement in self.mapping:
|
||||
if not isinstance(replacement, ContentReplacement):
|
||||
continue
|
||||
if replacement.content in key:
|
||||
key = key.replace(replacement.content, replacement.replacement)
|
||||
return key
|
||||
|
||||
def apply_to_key_value(self, key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
|
||||
"""Apply the value operation to the given name and associated value."""
|
||||
for operation in self.mapping:
|
||||
if not isinstance(operation, SDKeyValueOperation):
|
||||
continue
|
||||
if key.startswith(operation.key_matcher.prefix) and key.endswith(operation.key_matcher.suffix):
|
||||
return operation.kv_operation(key, value)
|
||||
return [KeyValueOperationResult(key, value)]
|
||||
|
||||
|
||||
# Predefined SDOps instances
|
||||
LTXV_LORA_COMFY_RENAMING_MAP = (
|
||||
SDOps("LTXV_LORA_COMFY_PREFIX_MAP").with_matching().with_replacement("diffusion_model.", "")
|
||||
)
|
||||
|
||||
LTXV_LORA_COMFY_TARGET_MAP = (
|
||||
SDOps("LTXV_LORA_COMFY_TARGET_MAP")
|
||||
.with_matching()
|
||||
.with_replacement("diffusion_model.", "")
|
||||
.with_replacement(".lora_A.weight", ".weight")
|
||||
.with_replacement(".lora_B.weight", ".weight")
|
||||
)
|
||||
@@ -0,0 +1,63 @@
|
||||
import json
|
||||
|
||||
import safetensors
|
||||
import torch
|
||||
|
||||
from ltx_core.loader.primitives import StateDict, StateDictLoader
|
||||
from ltx_core.loader.sd_ops import SDOps
|
||||
|
||||
|
||||
class SafetensorsStateDictLoader(StateDictLoader):
|
||||
"""
|
||||
Loads weights from safetensors files without metadata support.
|
||||
Use this for loading raw weight files. For model files that include
|
||||
configuration metadata, use SafetensorsModelStateDictLoader instead.
|
||||
"""
|
||||
|
||||
def metadata(self, path: str) -> dict:
|
||||
raise NotImplementedError("Not implemented")
|
||||
|
||||
def load(self, path: str | list[str], sd_ops: SDOps, device: torch.device | None = None) -> StateDict:
|
||||
"""
|
||||
Load state dict from path or paths (for sharded model storage) and apply sd_ops
|
||||
"""
|
||||
sd = {}
|
||||
size = 0
|
||||
dtype = set()
|
||||
device = device or torch.device("cpu")
|
||||
model_paths = path if isinstance(path, list) else [path]
|
||||
for shard_path in model_paths:
|
||||
with safetensors.safe_open(shard_path, framework="pt", device=str(device)) as f:
|
||||
safetensor_keys = f.keys()
|
||||
for name in safetensor_keys:
|
||||
expected_name = name if sd_ops is None else sd_ops.apply_to_key(name)
|
||||
if expected_name is None:
|
||||
continue
|
||||
value = f.get_tensor(name).to(device=device, non_blocking=True, copy=False)
|
||||
key_value_pairs = ((expected_name, value),)
|
||||
if sd_ops is not None:
|
||||
key_value_pairs = sd_ops.apply_to_key_value(expected_name, value)
|
||||
for key, value in key_value_pairs:
|
||||
size += value.nbytes
|
||||
dtype.add(value.dtype)
|
||||
sd[key] = value
|
||||
|
||||
return StateDict(sd=sd, device=device, size=size, dtype=dtype)
|
||||
|
||||
|
||||
class SafetensorsModelStateDictLoader(StateDictLoader):
|
||||
"""
|
||||
Loads weights and configuration metadata from safetensors model files.
|
||||
Unlike SafetensorsStateDictLoader, this loader can read model configuration
|
||||
from the safetensors file metadata via the metadata() method.
|
||||
"""
|
||||
|
||||
def __init__(self, weight_loader: SafetensorsStateDictLoader | None = None):
|
||||
self.weight_loader = weight_loader if weight_loader is not None else SafetensorsStateDictLoader()
|
||||
|
||||
def metadata(self, path: str) -> dict:
|
||||
with safetensors.safe_open(path, framework="pt") as f:
|
||||
return json.loads(f.metadata()["config"])
|
||||
|
||||
def load(self, path: str | list[str], sd_ops: SDOps | None = None, device: torch.device | None = None) -> StateDict:
|
||||
return self.weight_loader.load(path, sd_ops, device)
|
||||
@@ -0,0 +1,101 @@
|
||||
import logging
|
||||
from dataclasses import dataclass, field, replace
|
||||
from typing import Generic
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.loader.fuse_loras import apply_loras
|
||||
from ltx_core.loader.module_ops import ModuleOps
|
||||
from ltx_core.loader.primitives import (
|
||||
LoRAAdaptableProtocol,
|
||||
LoraPathStrengthAndSDOps,
|
||||
LoraStateDictWithStrength,
|
||||
ModelBuilderProtocol,
|
||||
StateDict,
|
||||
StateDictLoader,
|
||||
)
|
||||
from ltx_core.loader.registry import DummyRegistry, Registry
|
||||
from ltx_core.loader.sd_ops import SDOps
|
||||
from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader
|
||||
from ltx_core.model.model_protocol import ModelConfigurator, ModelType
|
||||
|
||||
logger: logging.Logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType], LoRAAdaptableProtocol):
|
||||
"""
|
||||
Builder for PyTorch models residing on a single GPU.
|
||||
"""
|
||||
|
||||
model_class_configurator: type[ModelConfigurator[ModelType]]
|
||||
model_path: str | tuple[str, ...]
|
||||
model_sd_ops: SDOps | None = None
|
||||
module_ops: tuple[ModuleOps, ...] = field(default_factory=tuple)
|
||||
loras: tuple[LoraPathStrengthAndSDOps, ...] = field(default_factory=tuple)
|
||||
model_loader: StateDictLoader = field(default_factory=SafetensorsModelStateDictLoader)
|
||||
registry: Registry = field(default_factory=DummyRegistry)
|
||||
|
||||
def lora(self, lora_path: str, strength: float = 1.0, sd_ops: SDOps | None = None) -> "SingleGPUModelBuilder":
|
||||
return replace(self, loras=(*self.loras, LoraPathStrengthAndSDOps(lora_path, strength, sd_ops)))
|
||||
|
||||
def model_config(self) -> dict:
|
||||
first_shard_path = self.model_path[0] if isinstance(self.model_path, tuple) else self.model_path
|
||||
return self.model_loader.metadata(first_shard_path)
|
||||
|
||||
def meta_model(self, config: dict, module_ops: tuple[ModuleOps, ...]) -> ModelType:
|
||||
with torch.device("meta"):
|
||||
model = self.model_class_configurator.from_config(config)
|
||||
for module_op in module_ops:
|
||||
if module_op.matcher(model):
|
||||
model = module_op.mutator(model)
|
||||
return model
|
||||
|
||||
def load_sd(
|
||||
self, paths: list[str], registry: Registry, device: torch.device | None, sd_ops: SDOps | None = None
|
||||
) -> StateDict:
|
||||
state_dict = registry.get(paths, sd_ops)
|
||||
if state_dict is None:
|
||||
state_dict = self.model_loader.load(paths, sd_ops=sd_ops, device=device)
|
||||
registry.add(paths, sd_ops=sd_ops, state_dict=state_dict)
|
||||
return state_dict
|
||||
|
||||
def _return_model(self, meta_model: ModelType, device: torch.device) -> ModelType:
|
||||
uninitialized_params = [name for name, param in meta_model.named_parameters() if str(param.device) == "meta"]
|
||||
uninitialized_buffers = [name for name, buffer in meta_model.named_buffers() if str(buffer.device) == "meta"]
|
||||
if uninitialized_params or uninitialized_buffers:
|
||||
logger.warning(f"Uninitialized parameters or buffers: {uninitialized_params + uninitialized_buffers}")
|
||||
return meta_model
|
||||
retval = meta_model.to(device)
|
||||
return retval
|
||||
|
||||
def build(self, device: torch.device | None = None, dtype: torch.dtype | None = None) -> ModelType:
|
||||
device = torch.device("cuda") if device is None else device
|
||||
config = self.model_config()
|
||||
meta_model = self.meta_model(config, self.module_ops)
|
||||
model_paths = self.model_path if isinstance(self.model_path, tuple) else [self.model_path]
|
||||
model_state_dict = self.load_sd(model_paths, sd_ops=self.model_sd_ops, registry=self.registry, device=device)
|
||||
|
||||
lora_strengths = [lora.strength for lora in self.loras]
|
||||
if not lora_strengths or (min(lora_strengths) == 0 and max(lora_strengths) == 0):
|
||||
sd = model_state_dict.sd
|
||||
if dtype is not None:
|
||||
sd = {key: value.to(dtype=dtype) for key, value in model_state_dict.sd.items()}
|
||||
meta_model.load_state_dict(sd, strict=False, assign=True)
|
||||
return self._return_model(meta_model, device)
|
||||
|
||||
lora_state_dicts = [
|
||||
self.load_sd([lora.path], sd_ops=lora.sd_ops, registry=self.registry, device=device) for lora in self.loras
|
||||
]
|
||||
lora_sd_and_strengths = [
|
||||
LoraStateDictWithStrength(sd, strength)
|
||||
for sd, strength in zip(lora_state_dicts, lora_strengths, strict=True)
|
||||
]
|
||||
final_sd = apply_loras(
|
||||
model_sd=model_state_dict,
|
||||
lora_sd_and_strengths=lora_sd_and_strengths,
|
||||
dtype=dtype,
|
||||
destination_sd=model_state_dict if isinstance(self.registry, DummyRegistry) else None,
|
||||
)
|
||||
meta_model.load_state_dict(final_sd.sd, strict=False, assign=True)
|
||||
return self._return_model(meta_model, device)
|
||||
@@ -0,0 +1,8 @@
|
||||
"""Model definitions for LTX-2."""
|
||||
|
||||
from ltx_core.model.model_protocol import ModelConfigurator, ModelType
|
||||
|
||||
__all__ = [
|
||||
"ModelConfigurator",
|
||||
"ModelType",
|
||||
]
|
||||
@@ -0,0 +1,27 @@
|
||||
"""Audio VAE model components."""
|
||||
|
||||
from ltx_core.model.audio_vae.audio_vae import AudioDecoder, AudioEncoder, decode_audio
|
||||
from ltx_core.model.audio_vae.model_configurator import (
|
||||
AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
|
||||
AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
|
||||
VOCODER_COMFY_KEYS_FILTER,
|
||||
AudioDecoderConfigurator,
|
||||
AudioEncoderConfigurator,
|
||||
VocoderConfigurator,
|
||||
)
|
||||
from ltx_core.model.audio_vae.ops import AudioProcessor
|
||||
from ltx_core.model.audio_vae.vocoder import Vocoder
|
||||
|
||||
__all__ = [
|
||||
"AUDIO_VAE_DECODER_COMFY_KEYS_FILTER",
|
||||
"AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER",
|
||||
"VOCODER_COMFY_KEYS_FILTER",
|
||||
"AudioDecoder",
|
||||
"AudioDecoderConfigurator",
|
||||
"AudioEncoder",
|
||||
"AudioEncoderConfigurator",
|
||||
"AudioProcessor",
|
||||
"Vocoder",
|
||||
"VocoderConfigurator",
|
||||
"decode_audio",
|
||||
]
|
||||
@@ -0,0 +1,71 @@
|
||||
from enum import Enum
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.model.common.normalization import NormType, build_normalization_layer
|
||||
|
||||
|
||||
class AttentionType(Enum):
|
||||
"""Enum for specifying the attention mechanism type."""
|
||||
|
||||
VANILLA = "vanilla"
|
||||
LINEAR = "linear"
|
||||
NONE = "none"
|
||||
|
||||
|
||||
class AttnBlock(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
norm_type: NormType = NormType.GROUP,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
|
||||
self.norm = build_normalization_layer(in_channels, normtype=norm_type)
|
||||
self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
h_ = x
|
||||
h_ = self.norm(h_)
|
||||
q = self.q(h_)
|
||||
k = self.k(h_)
|
||||
v = self.v(h_)
|
||||
|
||||
# compute attention
|
||||
b, c, h, w = q.shape
|
||||
q = q.reshape(b, c, h * w).contiguous()
|
||||
q = q.permute(0, 2, 1).contiguous() # b,hw,c
|
||||
k = k.reshape(b, c, h * w).contiguous() # b,c,hw
|
||||
w_ = torch.bmm(q, k).contiguous() # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
|
||||
w_ = w_ * (int(c) ** (-0.5))
|
||||
w_ = torch.nn.functional.softmax(w_, dim=2)
|
||||
|
||||
# attend to values
|
||||
v = v.reshape(b, c, h * w).contiguous()
|
||||
w_ = w_.permute(0, 2, 1).contiguous() # b,hw,hw (first hw of k, second of q)
|
||||
h_ = torch.bmm(v, w_).contiguous() # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
|
||||
h_ = h_.reshape(b, c, h, w).contiguous()
|
||||
|
||||
h_ = self.proj_out(h_)
|
||||
|
||||
return x + h_
|
||||
|
||||
|
||||
def make_attn(
|
||||
in_channels: int,
|
||||
attn_type: AttentionType = AttentionType.VANILLA,
|
||||
norm_type: NormType = NormType.GROUP,
|
||||
) -> torch.nn.Module:
|
||||
match attn_type:
|
||||
case AttentionType.VANILLA:
|
||||
return AttnBlock(in_channels, norm_type=norm_type)
|
||||
case AttentionType.NONE:
|
||||
return torch.nn.Identity()
|
||||
case AttentionType.LINEAR:
|
||||
raise NotImplementedError(f"Attention type {attn_type.value} is not supported yet.")
|
||||
case _:
|
||||
raise ValueError(f"Unknown attention type: {attn_type}")
|
||||
@@ -0,0 +1,480 @@
|
||||
from typing import Set, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from ltx_core.components.patchifiers import AudioPatchifier
|
||||
from ltx_core.model.audio_vae.attention import AttentionType, make_attn
|
||||
from ltx_core.model.audio_vae.causal_conv_2d import make_conv2d
|
||||
from ltx_core.model.audio_vae.causality_axis import CausalityAxis
|
||||
from ltx_core.model.audio_vae.downsample import build_downsampling_path
|
||||
from ltx_core.model.audio_vae.ops import PerChannelStatistics
|
||||
from ltx_core.model.audio_vae.resnet import ResnetBlock
|
||||
from ltx_core.model.audio_vae.upsample import build_upsampling_path
|
||||
from ltx_core.model.audio_vae.vocoder import Vocoder
|
||||
from ltx_core.model.common.normalization import NormType, build_normalization_layer
|
||||
from ltx_core.types import AudioLatentShape
|
||||
|
||||
LATENT_DOWNSAMPLE_FACTOR = 4
|
||||
|
||||
|
||||
def build_mid_block(
|
||||
channels: int,
|
||||
temb_channels: int,
|
||||
dropout: float,
|
||||
norm_type: NormType,
|
||||
causality_axis: CausalityAxis,
|
||||
attn_type: AttentionType,
|
||||
add_attention: bool,
|
||||
) -> torch.nn.Module:
|
||||
"""Build the middle block with two ResNet blocks and optional attention."""
|
||||
mid = torch.nn.Module()
|
||||
mid.block_1 = ResnetBlock(
|
||||
in_channels=channels,
|
||||
out_channels=channels,
|
||||
temb_channels=temb_channels,
|
||||
dropout=dropout,
|
||||
norm_type=norm_type,
|
||||
causality_axis=causality_axis,
|
||||
)
|
||||
mid.attn_1 = make_attn(channels, attn_type=attn_type, norm_type=norm_type) if add_attention else torch.nn.Identity()
|
||||
mid.block_2 = ResnetBlock(
|
||||
in_channels=channels,
|
||||
out_channels=channels,
|
||||
temb_channels=temb_channels,
|
||||
dropout=dropout,
|
||||
norm_type=norm_type,
|
||||
causality_axis=causality_axis,
|
||||
)
|
||||
return mid
|
||||
|
||||
|
||||
def run_mid_block(mid: torch.nn.Module, features: torch.Tensor) -> torch.Tensor:
|
||||
"""Run features through the middle block."""
|
||||
features = mid.block_1(features, temb=None)
|
||||
features = mid.attn_1(features)
|
||||
return mid.block_2(features, temb=None)
|
||||
|
||||
|
||||
class AudioEncoder(torch.nn.Module):
|
||||
"""
|
||||
Encoder that compresses audio spectrograms into latent representations.
|
||||
The encoder uses a series of downsampling blocks with residual connections,
|
||||
attention mechanisms, and configurable causal convolutions.
|
||||
"""
|
||||
|
||||
def __init__( # noqa: PLR0913
|
||||
self,
|
||||
*,
|
||||
ch: int,
|
||||
ch_mult: Tuple[int, ...] = (1, 2, 4, 8),
|
||||
num_res_blocks: int,
|
||||
attn_resolutions: Set[int],
|
||||
dropout: float = 0.0,
|
||||
resamp_with_conv: bool = True,
|
||||
in_channels: int,
|
||||
resolution: int,
|
||||
z_channels: int,
|
||||
double_z: bool = True,
|
||||
attn_type: AttentionType = AttentionType.VANILLA,
|
||||
mid_block_add_attention: bool = True,
|
||||
norm_type: NormType = NormType.GROUP,
|
||||
causality_axis: CausalityAxis = CausalityAxis.WIDTH,
|
||||
sample_rate: int = 16000,
|
||||
mel_hop_length: int = 160,
|
||||
n_fft: int = 1024,
|
||||
is_causal: bool = True,
|
||||
mel_bins: int = 64,
|
||||
**_ignore_kwargs,
|
||||
) -> None:
|
||||
"""
|
||||
Initialize the Encoder.
|
||||
Args:
|
||||
Arguments are configuration parameters, loaded from the audio VAE checkpoint config
|
||||
(audio_vae.model.params.ddconfig):
|
||||
ch: Base number of feature channels used in the first convolution layer.
|
||||
ch_mult: Multiplicative factors for the number of channels at each resolution level.
|
||||
num_res_blocks: Number of residual blocks to use at each resolution level.
|
||||
attn_resolutions: Spatial resolutions (e.g., in time/frequency) at which to apply attention.
|
||||
resolution: Input spatial resolution of the spectrogram (height, width).
|
||||
z_channels: Number of channels in the latent representation.
|
||||
norm_type: Normalization layer type to use within the network (e.g., group, batch).
|
||||
causality_axis: Axis along which convolutions should be causal (e.g., time axis).
|
||||
sample_rate: Audio sample rate in Hz for the input signals.
|
||||
mel_hop_length: Hop length used when computing the mel spectrogram.
|
||||
n_fft: FFT size used to compute the spectrogram.
|
||||
mel_bins: Number of mel-frequency bins in the input spectrogram.
|
||||
in_channels: Number of channels in the input spectrogram tensor.
|
||||
double_z: If True, predict both mean and log-variance (doubling latent channels).
|
||||
is_causal: If True, use causal convolutions suitable for streaming setups.
|
||||
dropout: Dropout probability used in residual and mid blocks.
|
||||
attn_type: Type of attention mechanism to use in attention blocks.
|
||||
resamp_with_conv: If True, perform resolution changes using strided convolutions.
|
||||
mid_block_add_attention: If True, add an attention block in the mid-level of the encoder.
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.per_channel_statistics = PerChannelStatistics(latent_channels=ch)
|
||||
self.sample_rate = sample_rate
|
||||
self.mel_hop_length = mel_hop_length
|
||||
self.n_fft = n_fft
|
||||
self.is_causal = is_causal
|
||||
self.mel_bins = mel_bins
|
||||
|
||||
self.patchifier = AudioPatchifier(
|
||||
patch_size=1,
|
||||
audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR,
|
||||
sample_rate=sample_rate,
|
||||
hop_length=mel_hop_length,
|
||||
is_causal=is_causal,
|
||||
)
|
||||
|
||||
self.ch = ch
|
||||
self.temb_ch = 0
|
||||
self.num_resolutions = len(ch_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.resolution = resolution
|
||||
self.in_channels = in_channels
|
||||
self.z_channels = z_channels
|
||||
self.double_z = double_z
|
||||
self.norm_type = norm_type
|
||||
self.causality_axis = causality_axis
|
||||
self.attn_type = attn_type
|
||||
|
||||
# downsampling
|
||||
self.conv_in = make_conv2d(
|
||||
in_channels,
|
||||
self.ch,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
causality_axis=self.causality_axis,
|
||||
)
|
||||
|
||||
self.non_linearity = torch.nn.SiLU()
|
||||
|
||||
self.down, block_in = build_downsampling_path(
|
||||
ch=ch,
|
||||
ch_mult=ch_mult,
|
||||
num_resolutions=self.num_resolutions,
|
||||
num_res_blocks=num_res_blocks,
|
||||
resolution=resolution,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout,
|
||||
norm_type=self.norm_type,
|
||||
causality_axis=self.causality_axis,
|
||||
attn_type=self.attn_type,
|
||||
attn_resolutions=attn_resolutions,
|
||||
resamp_with_conv=resamp_with_conv,
|
||||
)
|
||||
|
||||
self.mid = build_mid_block(
|
||||
channels=block_in,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout,
|
||||
norm_type=self.norm_type,
|
||||
causality_axis=self.causality_axis,
|
||||
attn_type=self.attn_type,
|
||||
add_attention=mid_block_add_attention,
|
||||
)
|
||||
|
||||
self.norm_out = build_normalization_layer(block_in, normtype=self.norm_type)
|
||||
self.conv_out = make_conv2d(
|
||||
block_in,
|
||||
2 * z_channels if double_z else z_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
causality_axis=self.causality_axis,
|
||||
)
|
||||
|
||||
def forward(self, spectrogram: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Encode audio spectrogram into latent representations.
|
||||
Args:
|
||||
spectrogram: Input spectrogram of shape (batch, channels, time, frequency)
|
||||
Returns:
|
||||
Encoded latent representation of shape (batch, channels, frames, mel_bins)
|
||||
"""
|
||||
h = self.conv_in(spectrogram)
|
||||
h = self._run_downsampling_path(h)
|
||||
h = run_mid_block(self.mid, h)
|
||||
h = self._finalize_output(h)
|
||||
|
||||
return self._normalize_latents(h)
|
||||
|
||||
def _run_downsampling_path(self, h: torch.Tensor) -> torch.Tensor:
|
||||
for level in range(self.num_resolutions):
|
||||
stage = self.down[level]
|
||||
for block_idx in range(self.num_res_blocks):
|
||||
h = stage.block[block_idx](h, temb=None)
|
||||
if stage.attn:
|
||||
h = stage.attn[block_idx](h)
|
||||
|
||||
if level != self.num_resolutions - 1:
|
||||
h = stage.downsample(h)
|
||||
|
||||
return h
|
||||
|
||||
def _finalize_output(self, h: torch.Tensor) -> torch.Tensor:
|
||||
h = self.norm_out(h)
|
||||
h = self.non_linearity(h)
|
||||
return self.conv_out(h)
|
||||
|
||||
def _normalize_latents(self, latent_output: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Normalize encoder latents using per-channel statistics.
|
||||
When the encoder is configured with ``double_z=True``, the final
|
||||
convolution produces twice the number of latent channels, typically
|
||||
interpreted as two concatenated tensors along the channel dimension
|
||||
(e.g., mean and variance or other auxiliary parameters).
|
||||
This method intentionally uses only the first half of the channels
|
||||
(the "mean" component) as input to the patchifier and normalization
|
||||
logic. The remaining channels are left unchanged by this method and
|
||||
are expected to be consumed elsewhere in the VAE pipeline.
|
||||
If ``double_z=False``, the encoder output already contains only the
|
||||
mean latents and the chunking operation simply returns that tensor.
|
||||
"""
|
||||
means = torch.chunk(latent_output, 2, dim=1)[0]
|
||||
latent_shape = AudioLatentShape(
|
||||
batch=means.shape[0],
|
||||
channels=means.shape[1],
|
||||
frames=means.shape[2],
|
||||
mel_bins=means.shape[3],
|
||||
)
|
||||
latent_patched = self.patchifier.patchify(means)
|
||||
latent_normalized = self.per_channel_statistics.normalize(latent_patched)
|
||||
return self.patchifier.unpatchify(latent_normalized, latent_shape)
|
||||
|
||||
|
||||
class AudioDecoder(torch.nn.Module):
|
||||
"""
|
||||
Symmetric decoder that reconstructs audio spectrograms from latent features.
|
||||
The decoder mirrors the encoder structure with configurable channel multipliers,
|
||||
attention resolutions, and causal convolutions.
|
||||
"""
|
||||
|
||||
def __init__( # noqa: PLR0913
|
||||
self,
|
||||
*,
|
||||
ch: int,
|
||||
out_ch: int,
|
||||
ch_mult: Tuple[int, ...] = (1, 2, 4, 8),
|
||||
num_res_blocks: int,
|
||||
attn_resolutions: Set[int],
|
||||
resolution: int,
|
||||
z_channels: int,
|
||||
norm_type: NormType = NormType.GROUP,
|
||||
causality_axis: CausalityAxis = CausalityAxis.WIDTH,
|
||||
dropout: float = 0.0,
|
||||
mid_block_add_attention: bool = True,
|
||||
sample_rate: int = 16000,
|
||||
mel_hop_length: int = 160,
|
||||
is_causal: bool = True,
|
||||
mel_bins: int | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Initialize the Decoder.
|
||||
Args:
|
||||
Arguments are configuration parameters, loaded from the audio VAE checkpoint config
|
||||
(audio_vae.model.params.ddconfig):
|
||||
- ch, out_ch, ch_mult, num_res_blocks, attn_resolutions
|
||||
- resolution, z_channels
|
||||
- norm_type, causality_axis
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Internal behavioural defaults that are not driven by the checkpoint.
|
||||
resamp_with_conv = True
|
||||
attn_type = AttentionType.VANILLA
|
||||
|
||||
# Per-channel statistics for denormalizing latents
|
||||
self.per_channel_statistics = PerChannelStatistics(latent_channels=ch)
|
||||
self.sample_rate = sample_rate
|
||||
self.mel_hop_length = mel_hop_length
|
||||
self.is_causal = is_causal
|
||||
self.mel_bins = mel_bins
|
||||
self.patchifier = AudioPatchifier(
|
||||
patch_size=1,
|
||||
audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR,
|
||||
sample_rate=sample_rate,
|
||||
hop_length=mel_hop_length,
|
||||
is_causal=is_causal,
|
||||
)
|
||||
|
||||
self.ch = ch
|
||||
self.temb_ch = 0
|
||||
self.num_resolutions = len(ch_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.resolution = resolution
|
||||
self.out_ch = out_ch
|
||||
self.give_pre_end = False
|
||||
self.tanh_out = False
|
||||
self.norm_type = norm_type
|
||||
self.z_channels = z_channels
|
||||
self.channel_multipliers = ch_mult
|
||||
self.attn_resolutions = attn_resolutions
|
||||
self.causality_axis = causality_axis
|
||||
self.attn_type = attn_type
|
||||
|
||||
base_block_channels = ch * self.channel_multipliers[-1]
|
||||
base_resolution = resolution // (2 ** (self.num_resolutions - 1))
|
||||
self.z_shape = (1, z_channels, base_resolution, base_resolution)
|
||||
|
||||
self.conv_in = make_conv2d(
|
||||
z_channels, base_block_channels, kernel_size=3, stride=1, causality_axis=self.causality_axis
|
||||
)
|
||||
self.non_linearity = torch.nn.SiLU()
|
||||
self.mid = build_mid_block(
|
||||
channels=base_block_channels,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout,
|
||||
norm_type=self.norm_type,
|
||||
causality_axis=self.causality_axis,
|
||||
attn_type=self.attn_type,
|
||||
add_attention=mid_block_add_attention,
|
||||
)
|
||||
self.up, final_block_channels = build_upsampling_path(
|
||||
ch=ch,
|
||||
ch_mult=ch_mult,
|
||||
num_resolutions=self.num_resolutions,
|
||||
num_res_blocks=num_res_blocks,
|
||||
resolution=resolution,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout,
|
||||
norm_type=self.norm_type,
|
||||
causality_axis=self.causality_axis,
|
||||
attn_type=self.attn_type,
|
||||
attn_resolutions=attn_resolutions,
|
||||
resamp_with_conv=resamp_with_conv,
|
||||
initial_block_channels=base_block_channels,
|
||||
)
|
||||
|
||||
self.norm_out = build_normalization_layer(final_block_channels, normtype=self.norm_type)
|
||||
self.conv_out = make_conv2d(
|
||||
final_block_channels, out_ch, kernel_size=3, stride=1, causality_axis=self.causality_axis
|
||||
)
|
||||
|
||||
def forward(self, sample: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Decode latent features back to audio spectrograms.
|
||||
Args:
|
||||
sample: Encoded latent representation of shape (batch, channels, frames, mel_bins)
|
||||
Returns:
|
||||
Reconstructed audio spectrogram of shape (batch, channels, time, frequency)
|
||||
"""
|
||||
sample, target_shape = self._denormalize_latents(sample)
|
||||
|
||||
h = self.conv_in(sample)
|
||||
h = run_mid_block(self.mid, h)
|
||||
h = self._run_upsampling_path(h)
|
||||
h = self._finalize_output(h)
|
||||
|
||||
return self._adjust_output_shape(h, target_shape)
|
||||
|
||||
def _denormalize_latents(self, sample: torch.Tensor) -> tuple[torch.Tensor, AudioLatentShape]:
|
||||
latent_shape = AudioLatentShape(
|
||||
batch=sample.shape[0],
|
||||
channels=sample.shape[1],
|
||||
frames=sample.shape[2],
|
||||
mel_bins=sample.shape[3],
|
||||
)
|
||||
|
||||
sample_patched = self.patchifier.patchify(sample)
|
||||
sample_denormalized = self.per_channel_statistics.un_normalize(sample_patched)
|
||||
sample = self.patchifier.unpatchify(sample_denormalized, latent_shape)
|
||||
|
||||
target_frames = latent_shape.frames * LATENT_DOWNSAMPLE_FACTOR
|
||||
if self.causality_axis != CausalityAxis.NONE:
|
||||
target_frames = max(target_frames - (LATENT_DOWNSAMPLE_FACTOR - 1), 1)
|
||||
|
||||
target_shape = AudioLatentShape(
|
||||
batch=latent_shape.batch,
|
||||
channels=self.out_ch,
|
||||
frames=target_frames,
|
||||
mel_bins=self.mel_bins if self.mel_bins is not None else latent_shape.mel_bins,
|
||||
)
|
||||
|
||||
return sample, target_shape
|
||||
|
||||
def _adjust_output_shape(
|
||||
self,
|
||||
decoded_output: torch.Tensor,
|
||||
target_shape: AudioLatentShape,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Adjust output shape to match target dimensions for variable-length audio.
|
||||
This function handles the common case where decoded audio spectrograms need to be
|
||||
resized to match a specific target shape.
|
||||
Args:
|
||||
decoded_output: Tensor of shape (batch, channels, time, frequency)
|
||||
target_shape: AudioLatentShape describing (batch, channels, time, mel bins)
|
||||
Returns:
|
||||
Tensor adjusted to match target_shape exactly
|
||||
"""
|
||||
# Current output shape: (batch, channels, time, frequency)
|
||||
_, _, current_time, current_freq = decoded_output.shape
|
||||
target_channels = target_shape.channels
|
||||
target_time = target_shape.frames
|
||||
target_freq = target_shape.mel_bins
|
||||
|
||||
# Step 1: Crop first to avoid exceeding target dimensions
|
||||
decoded_output = decoded_output[
|
||||
:, :target_channels, : min(current_time, target_time), : min(current_freq, target_freq)
|
||||
]
|
||||
|
||||
# Step 2: Calculate padding needed for time and frequency dimensions
|
||||
time_padding_needed = target_time - decoded_output.shape[2]
|
||||
freq_padding_needed = target_freq - decoded_output.shape[3]
|
||||
|
||||
# Step 3: Apply padding if needed
|
||||
if time_padding_needed > 0 or freq_padding_needed > 0:
|
||||
# PyTorch padding format: (pad_left, pad_right, pad_top, pad_bottom)
|
||||
# For audio: pad_left/right = frequency, pad_top/bottom = time
|
||||
padding = (
|
||||
0,
|
||||
max(freq_padding_needed, 0), # frequency padding (left, right)
|
||||
0,
|
||||
max(time_padding_needed, 0), # time padding (top, bottom)
|
||||
)
|
||||
decoded_output = F.pad(decoded_output, padding)
|
||||
|
||||
# Step 4: Final safety crop to ensure exact target shape
|
||||
decoded_output = decoded_output[:, :target_channels, :target_time, :target_freq]
|
||||
|
||||
return decoded_output
|
||||
|
||||
def _run_upsampling_path(self, h: torch.Tensor) -> torch.Tensor:
|
||||
for level in reversed(range(self.num_resolutions)):
|
||||
stage = self.up[level]
|
||||
for block_idx, block in enumerate(stage.block):
|
||||
h = block(h, temb=None)
|
||||
if stage.attn:
|
||||
h = stage.attn[block_idx](h)
|
||||
|
||||
if level != 0 and hasattr(stage, "upsample"):
|
||||
h = stage.upsample(h)
|
||||
|
||||
return h
|
||||
|
||||
def _finalize_output(self, h: torch.Tensor) -> torch.Tensor:
|
||||
if self.give_pre_end:
|
||||
return h
|
||||
|
||||
h = self.norm_out(h)
|
||||
h = self.non_linearity(h)
|
||||
h = self.conv_out(h)
|
||||
return torch.tanh(h) if self.tanh_out else h
|
||||
|
||||
|
||||
def decode_audio(latent: torch.Tensor, audio_decoder: "AudioDecoder", vocoder: "Vocoder") -> torch.Tensor:
|
||||
"""
|
||||
Decode an audio latent representation using the provided audio decoder and vocoder.
|
||||
Args:
|
||||
latent: Input audio latent tensor.
|
||||
audio_decoder: Model to decode the latent to waveform features.
|
||||
vocoder: Model to convert decoded features to audio waveform.
|
||||
Returns:
|
||||
Decoded audio as a float tensor.
|
||||
"""
|
||||
decoded_audio = audio_decoder(latent)
|
||||
decoded_audio = vocoder(decoded_audio).squeeze(0).float()
|
||||
return decoded_audio
|
||||
@@ -0,0 +1,110 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from ltx_core.model.audio_vae.causality_axis import CausalityAxis
|
||||
|
||||
|
||||
class CausalConv2d(torch.nn.Module):
|
||||
"""
|
||||
A causal 2D convolution.
|
||||
This layer ensures that the output at time `t` only depends on inputs
|
||||
at time `t` and earlier. It achieves this by applying asymmetric padding
|
||||
to the time dimension (width) before the convolution.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: int | tuple[int, int],
|
||||
stride: int = 1,
|
||||
dilation: int | tuple[int, int] = 1,
|
||||
groups: int = 1,
|
||||
bias: bool = True,
|
||||
causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.causality_axis = causality_axis
|
||||
|
||||
# Ensure kernel_size and dilation are tuples
|
||||
kernel_size = torch.nn.modules.utils._pair(kernel_size)
|
||||
dilation = torch.nn.modules.utils._pair(dilation)
|
||||
|
||||
# Calculate padding dimensions
|
||||
pad_h = (kernel_size[0] - 1) * dilation[0]
|
||||
pad_w = (kernel_size[1] - 1) * dilation[1]
|
||||
|
||||
# The padding tuple for F.pad is (pad_left, pad_right, pad_top, pad_bottom)
|
||||
match self.causality_axis:
|
||||
case CausalityAxis.NONE:
|
||||
self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2)
|
||||
case CausalityAxis.WIDTH | CausalityAxis.WIDTH_COMPATIBILITY:
|
||||
self.padding = (pad_w, 0, pad_h // 2, pad_h - pad_h // 2)
|
||||
case CausalityAxis.HEIGHT:
|
||||
self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h, 0)
|
||||
case _:
|
||||
raise ValueError(f"Invalid causality_axis: {causality_axis}")
|
||||
|
||||
# The internal convolution layer uses no padding, as we handle it manually
|
||||
self.conv = torch.nn.Conv2d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
padding=0,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
# Apply causal padding before convolution
|
||||
x = F.pad(x, self.padding)
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
def make_conv2d(
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: int | tuple[int, int],
|
||||
stride: int = 1,
|
||||
padding: tuple[int, int, int, int] | None = None,
|
||||
dilation: int = 1,
|
||||
groups: int = 1,
|
||||
bias: bool = True,
|
||||
causality_axis: CausalityAxis | None = None,
|
||||
) -> torch.nn.Module:
|
||||
"""
|
||||
Create a 2D convolution layer that can be either causal or non-causal.
|
||||
Args:
|
||||
in_channels: Number of input channels
|
||||
out_channels: Number of output channels
|
||||
kernel_size: Size of the convolution kernel
|
||||
stride: Convolution stride
|
||||
padding: Padding (if None, will be calculated based on causal flag)
|
||||
dilation: Dilation rate
|
||||
groups: Number of groups for grouped convolution
|
||||
bias: Whether to use bias
|
||||
causality_axis: Dimension along which to apply causality.
|
||||
Returns:
|
||||
Either a regular Conv2d or CausalConv2d layer
|
||||
"""
|
||||
if causality_axis is not None:
|
||||
# For causal convolution, padding is handled internally by CausalConv2d
|
||||
return CausalConv2d(in_channels, out_channels, kernel_size, stride, dilation, groups, bias, causality_axis)
|
||||
else:
|
||||
# For non-causal convolution, use symmetric padding if not specified
|
||||
if padding is None:
|
||||
padding = kernel_size // 2 if isinstance(kernel_size, int) else tuple(k // 2 for k in kernel_size)
|
||||
|
||||
return torch.nn.Conv2d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
groups,
|
||||
bias,
|
||||
)
|
||||
@@ -0,0 +1,10 @@
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class CausalityAxis(Enum):
|
||||
"""Enum for specifying the causality axis in causal convolutions."""
|
||||
|
||||
NONE = None
|
||||
WIDTH = "width"
|
||||
HEIGHT = "height"
|
||||
WIDTH_COMPATIBILITY = "width-compatibility"
|
||||
@@ -0,0 +1,110 @@
|
||||
from typing import Set, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.model.audio_vae.attention import AttentionType, make_attn
|
||||
from ltx_core.model.audio_vae.causality_axis import CausalityAxis
|
||||
from ltx_core.model.audio_vae.resnet import ResnetBlock
|
||||
from ltx_core.model.common.normalization import NormType
|
||||
|
||||
|
||||
class Downsample(torch.nn.Module):
|
||||
"""
|
||||
A downsampling layer that can use either a strided convolution
|
||||
or average pooling. Supports standard and causal padding for the
|
||||
convolutional mode.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
with_conv: bool,
|
||||
causality_axis: CausalityAxis = CausalityAxis.WIDTH,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.with_conv = with_conv
|
||||
self.causality_axis = causality_axis
|
||||
|
||||
if self.causality_axis != CausalityAxis.NONE and not self.with_conv:
|
||||
raise ValueError("causality is only supported when `with_conv=True`.")
|
||||
|
||||
if self.with_conv:
|
||||
# Do time downsampling here
|
||||
# no asymmetric padding in torch conv, must do it ourselves
|
||||
self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if self.with_conv:
|
||||
# Padding tuple is in the order: (left, right, top, bottom).
|
||||
match self.causality_axis:
|
||||
case CausalityAxis.NONE:
|
||||
pad = (0, 1, 0, 1)
|
||||
case CausalityAxis.WIDTH:
|
||||
pad = (2, 0, 0, 1)
|
||||
case CausalityAxis.HEIGHT:
|
||||
pad = (0, 1, 2, 0)
|
||||
case CausalityAxis.WIDTH_COMPATIBILITY:
|
||||
pad = (1, 0, 0, 1)
|
||||
case _:
|
||||
raise ValueError(f"Invalid causality_axis: {self.causality_axis}")
|
||||
|
||||
x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
|
||||
x = self.conv(x)
|
||||
else:
|
||||
# This branch is only taken if with_conv=False, which implies causality_axis is NONE.
|
||||
x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def build_downsampling_path( # noqa: PLR0913
|
||||
*,
|
||||
ch: int,
|
||||
ch_mult: Tuple[int, ...],
|
||||
num_resolutions: int,
|
||||
num_res_blocks: int,
|
||||
resolution: int,
|
||||
temb_channels: int,
|
||||
dropout: float,
|
||||
norm_type: NormType,
|
||||
causality_axis: CausalityAxis,
|
||||
attn_type: AttentionType,
|
||||
attn_resolutions: Set[int],
|
||||
resamp_with_conv: bool,
|
||||
) -> tuple[torch.nn.ModuleList, int]:
|
||||
"""Build the downsampling path with residual blocks, attention, and downsampling layers."""
|
||||
down_modules = torch.nn.ModuleList()
|
||||
curr_res = resolution
|
||||
in_ch_mult = (1, *tuple(ch_mult))
|
||||
block_in = ch
|
||||
|
||||
for i_level in range(num_resolutions):
|
||||
block = torch.nn.ModuleList()
|
||||
attn = torch.nn.ModuleList()
|
||||
block_in = ch * in_ch_mult[i_level]
|
||||
block_out = ch * ch_mult[i_level]
|
||||
|
||||
for _ in range(num_res_blocks):
|
||||
block.append(
|
||||
ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_out,
|
||||
temb_channels=temb_channels,
|
||||
dropout=dropout,
|
||||
norm_type=norm_type,
|
||||
causality_axis=causality_axis,
|
||||
)
|
||||
)
|
||||
block_in = block_out
|
||||
if curr_res in attn_resolutions:
|
||||
attn.append(make_attn(block_in, attn_type=attn_type, norm_type=norm_type))
|
||||
|
||||
down = torch.nn.Module()
|
||||
down.block = block
|
||||
down.attn = attn
|
||||
if i_level != num_resolutions - 1:
|
||||
down.downsample = Downsample(block_in, resamp_with_conv, causality_axis=causality_axis)
|
||||
curr_res = curr_res // 2
|
||||
down_modules.append(down)
|
||||
|
||||
return down_modules, block_in
|
||||
@@ -0,0 +1,123 @@
|
||||
from ltx_core.loader.sd_ops import SDOps
|
||||
from ltx_core.model.audio_vae.attention import AttentionType
|
||||
from ltx_core.model.audio_vae.audio_vae import AudioDecoder, AudioEncoder
|
||||
from ltx_core.model.audio_vae.causality_axis import CausalityAxis
|
||||
from ltx_core.model.audio_vae.vocoder import Vocoder
|
||||
from ltx_core.model.common.normalization import NormType
|
||||
from ltx_core.model.model_protocol import ModelConfigurator
|
||||
|
||||
|
||||
class VocoderConfigurator(ModelConfigurator[Vocoder]):
|
||||
@classmethod
|
||||
def from_config(cls: type[Vocoder], config: dict) -> Vocoder:
|
||||
config = config.get("vocoder", {})
|
||||
return Vocoder(
|
||||
resblock_kernel_sizes=config.get("resblock_kernel_sizes", [3, 7, 11]),
|
||||
upsample_rates=config.get("upsample_rates", [6, 5, 2, 2, 2]),
|
||||
upsample_kernel_sizes=config.get("upsample_kernel_sizes", [16, 15, 8, 4, 4]),
|
||||
resblock_dilation_sizes=config.get("resblock_dilation_sizes", [[1, 3, 5], [1, 3, 5], [1, 3, 5]]),
|
||||
upsample_initial_channel=config.get("upsample_initial_channel", 1024),
|
||||
stereo=config.get("stereo", True),
|
||||
resblock=config.get("resblock", "1"),
|
||||
output_sample_rate=config.get("output_sample_rate", 24000),
|
||||
)
|
||||
|
||||
|
||||
VOCODER_COMFY_KEYS_FILTER = (
|
||||
SDOps("VOCODER_COMFY_KEYS_FILTER").with_matching(prefix="vocoder.").with_replacement("vocoder.", "")
|
||||
)
|
||||
|
||||
|
||||
class AudioDecoderConfigurator(ModelConfigurator[AudioDecoder]):
|
||||
@classmethod
|
||||
def from_config(cls: type[AudioDecoder], config: dict) -> AudioDecoder:
|
||||
audio_vae_cfg = config.get("audio_vae", {})
|
||||
model_cfg = audio_vae_cfg.get("model", {})
|
||||
model_params = model_cfg.get("params", {})
|
||||
ddconfig = model_params.get("ddconfig", {})
|
||||
preprocessing_cfg = audio_vae_cfg.get("preprocessing", {})
|
||||
stft_cfg = preprocessing_cfg.get("stft", {})
|
||||
mel_cfg = preprocessing_cfg.get("mel", {})
|
||||
variables_cfg = audio_vae_cfg.get("variables", {})
|
||||
|
||||
sample_rate = model_params.get("sampling_rate", 16000)
|
||||
mel_hop_length = stft_cfg.get("hop_length", 160)
|
||||
is_causal = stft_cfg.get("causal", True)
|
||||
mel_bins = ddconfig.get("mel_bins") or mel_cfg.get("n_mel_channels") or variables_cfg.get("mel_bins")
|
||||
|
||||
return AudioDecoder(
|
||||
ch=ddconfig.get("ch", 128),
|
||||
out_ch=ddconfig.get("out_ch", 2),
|
||||
ch_mult=tuple(ddconfig.get("ch_mult", (1, 2, 4))),
|
||||
num_res_blocks=ddconfig.get("num_res_blocks", 2),
|
||||
attn_resolutions=ddconfig.get("attn_resolutions", {8, 16, 32}),
|
||||
resolution=ddconfig.get("resolution", 256),
|
||||
z_channels=ddconfig.get("z_channels", 8),
|
||||
norm_type=NormType(ddconfig.get("norm_type", "pixel")),
|
||||
causality_axis=CausalityAxis(ddconfig.get("causality_axis", "height")),
|
||||
dropout=ddconfig.get("dropout", 0.0),
|
||||
mid_block_add_attention=ddconfig.get("mid_block_add_attention", True),
|
||||
sample_rate=sample_rate,
|
||||
mel_hop_length=mel_hop_length,
|
||||
is_causal=is_causal,
|
||||
mel_bins=mel_bins,
|
||||
)
|
||||
|
||||
|
||||
class AudioEncoderConfigurator(ModelConfigurator[AudioEncoder]):
|
||||
@classmethod
|
||||
def from_config(cls: type[AudioEncoder], config: dict) -> AudioEncoder:
|
||||
audio_vae_cfg = config.get("audio_vae", {})
|
||||
model_cfg = audio_vae_cfg.get("model", {})
|
||||
model_params = model_cfg.get("params", {})
|
||||
ddconfig = model_params.get("ddconfig", {})
|
||||
preprocessing_cfg = audio_vae_cfg.get("preprocessing", {})
|
||||
stft_cfg = preprocessing_cfg.get("stft", {})
|
||||
mel_cfg = preprocessing_cfg.get("mel", {})
|
||||
variables_cfg = audio_vae_cfg.get("variables", {})
|
||||
|
||||
sample_rate = model_params.get("sampling_rate", 16000)
|
||||
mel_hop_length = stft_cfg.get("hop_length", 160)
|
||||
n_fft = stft_cfg.get("filter_length", 1024)
|
||||
is_causal = stft_cfg.get("causal", True)
|
||||
mel_bins = ddconfig.get("mel_bins") or mel_cfg.get("n_mel_channels") or variables_cfg.get("mel_bins")
|
||||
|
||||
return AudioEncoder(
|
||||
ch=ddconfig.get("ch", 128),
|
||||
ch_mult=tuple(ddconfig.get("ch_mult", (1, 2, 4))),
|
||||
num_res_blocks=ddconfig.get("num_res_blocks", 2),
|
||||
attn_resolutions=ddconfig.get("attn_resolutions", {8, 16, 32}),
|
||||
resolution=ddconfig.get("resolution", 256),
|
||||
z_channels=ddconfig.get("z_channels", 8),
|
||||
double_z=ddconfig.get("double_z", True),
|
||||
dropout=ddconfig.get("dropout", 0.0),
|
||||
resamp_with_conv=ddconfig.get("resamp_with_conv", True),
|
||||
in_channels=ddconfig.get("in_channels", 2),
|
||||
attn_type=AttentionType(ddconfig.get("attn_type", "vanilla")),
|
||||
mid_block_add_attention=ddconfig.get("mid_block_add_attention", True),
|
||||
norm_type=NormType(ddconfig.get("norm_type", "pixel")),
|
||||
causality_axis=CausalityAxis(ddconfig.get("causality_axis", "height")),
|
||||
sample_rate=sample_rate,
|
||||
mel_hop_length=mel_hop_length,
|
||||
n_fft=n_fft,
|
||||
is_causal=is_causal,
|
||||
mel_bins=mel_bins,
|
||||
)
|
||||
|
||||
|
||||
AUDIO_VAE_DECODER_COMFY_KEYS_FILTER = (
|
||||
SDOps("AUDIO_VAE_DECODER_COMFY_KEYS_FILTER")
|
||||
.with_matching(prefix="audio_vae.decoder.")
|
||||
.with_matching(prefix="audio_vae.per_channel_statistics.")
|
||||
.with_replacement("audio_vae.decoder.", "")
|
||||
.with_replacement("audio_vae.per_channel_statistics.", "per_channel_statistics.")
|
||||
)
|
||||
|
||||
|
||||
AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER = (
|
||||
SDOps("AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER")
|
||||
.with_matching(prefix="audio_vae.encoder.")
|
||||
.with_matching(prefix="audio_vae.per_channel_statistics.")
|
||||
.with_replacement("audio_vae.encoder.", "")
|
||||
.with_replacement("audio_vae.per_channel_statistics.", "per_channel_statistics.")
|
||||
)
|
||||
@@ -0,0 +1,76 @@
|
||||
import torch
|
||||
import torchaudio
|
||||
from torch import nn
|
||||
|
||||
|
||||
class AudioProcessor(nn.Module):
|
||||
"""Converts audio waveforms to log-mel spectrograms with optional resampling."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sample_rate: int,
|
||||
mel_bins: int,
|
||||
mel_hop_length: int,
|
||||
n_fft: int,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.sample_rate = sample_rate
|
||||
self.mel_transform = torchaudio.transforms.MelSpectrogram(
|
||||
sample_rate=sample_rate,
|
||||
n_fft=n_fft,
|
||||
win_length=n_fft,
|
||||
hop_length=mel_hop_length,
|
||||
f_min=0.0,
|
||||
f_max=sample_rate / 2.0,
|
||||
n_mels=mel_bins,
|
||||
window_fn=torch.hann_window,
|
||||
center=True,
|
||||
pad_mode="reflect",
|
||||
power=1.0,
|
||||
mel_scale="slaney",
|
||||
norm="slaney",
|
||||
)
|
||||
|
||||
def resample_waveform(
|
||||
self,
|
||||
waveform: torch.Tensor,
|
||||
source_rate: int,
|
||||
target_rate: int,
|
||||
) -> torch.Tensor:
|
||||
"""Resample waveform to target sample rate if needed."""
|
||||
if source_rate == target_rate:
|
||||
return waveform
|
||||
resampled = torchaudio.functional.resample(waveform, source_rate, target_rate)
|
||||
return resampled.to(device=waveform.device, dtype=waveform.dtype)
|
||||
|
||||
def waveform_to_mel(
|
||||
self,
|
||||
waveform: torch.Tensor,
|
||||
waveform_sample_rate: int,
|
||||
) -> torch.Tensor:
|
||||
"""Convert waveform to log-mel spectrogram [batch, channels, time, n_mels]."""
|
||||
waveform = self.resample_waveform(waveform, waveform_sample_rate, self.sample_rate)
|
||||
|
||||
mel = self.mel_transform(waveform)
|
||||
mel = torch.log(torch.clamp(mel, min=1e-5))
|
||||
|
||||
mel = mel.to(device=waveform.device, dtype=waveform.dtype)
|
||||
return mel.permute(0, 1, 3, 2).contiguous()
|
||||
|
||||
|
||||
class PerChannelStatistics(nn.Module):
|
||||
"""
|
||||
Per-channel statistics for normalizing and denormalizing the latent representation.
|
||||
This statics is computed over the entire dataset and stored in model's checkpoint under AudioVAE state_dict.
|
||||
"""
|
||||
|
||||
def __init__(self, latent_channels: int = 128) -> None:
|
||||
super().__init__()
|
||||
self.register_buffer("std-of-means", torch.empty(latent_channels))
|
||||
self.register_buffer("mean-of-means", torch.empty(latent_channels))
|
||||
|
||||
def un_normalize(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return (x * self.get_buffer("std-of-means").to(x)) + self.get_buffer("mean-of-means").to(x)
|
||||
|
||||
def normalize(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return (x - self.get_buffer("mean-of-means").to(x)) / self.get_buffer("std-of-means").to(x)
|
||||
@@ -0,0 +1,176 @@
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.model.audio_vae.causal_conv_2d import make_conv2d
|
||||
from ltx_core.model.audio_vae.causality_axis import CausalityAxis
|
||||
from ltx_core.model.common.normalization import NormType, build_normalization_layer
|
||||
|
||||
LRELU_SLOPE = 0.1
|
||||
|
||||
|
||||
class ResBlock1(torch.nn.Module):
|
||||
def __init__(self, channels: int, kernel_size: int = 3, dilation: Tuple[int, int, int] = (1, 3, 5)):
|
||||
super(ResBlock1, self).__init__()
|
||||
self.convs1 = torch.nn.ModuleList(
|
||||
[
|
||||
torch.nn.Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=dilation[0],
|
||||
padding="same",
|
||||
),
|
||||
torch.nn.Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=dilation[1],
|
||||
padding="same",
|
||||
),
|
||||
torch.nn.Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=dilation[2],
|
||||
padding="same",
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
self.convs2 = torch.nn.ModuleList(
|
||||
[
|
||||
torch.nn.Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=1,
|
||||
padding="same",
|
||||
),
|
||||
torch.nn.Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=1,
|
||||
padding="same",
|
||||
),
|
||||
torch.nn.Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=1,
|
||||
padding="same",
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
for conv1, conv2 in zip(self.convs1, self.convs2, strict=True):
|
||||
xt = torch.nn.functional.leaky_relu(x, LRELU_SLOPE)
|
||||
xt = conv1(xt)
|
||||
xt = torch.nn.functional.leaky_relu(xt, LRELU_SLOPE)
|
||||
xt = conv2(xt)
|
||||
x = xt + x
|
||||
return x
|
||||
|
||||
|
||||
class ResBlock2(torch.nn.Module):
|
||||
def __init__(self, channels: int, kernel_size: int = 3, dilation: Tuple[int, int] = (1, 3)):
|
||||
super(ResBlock2, self).__init__()
|
||||
self.convs = torch.nn.ModuleList(
|
||||
[
|
||||
torch.nn.Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=dilation[0],
|
||||
padding="same",
|
||||
),
|
||||
torch.nn.Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=dilation[1],
|
||||
padding="same",
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
for conv in self.convs:
|
||||
xt = torch.nn.functional.leaky_relu(x, LRELU_SLOPE)
|
||||
xt = conv(xt)
|
||||
x = xt + x
|
||||
return x
|
||||
|
||||
|
||||
class ResnetBlock(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
in_channels: int,
|
||||
out_channels: int | None = None,
|
||||
conv_shortcut: bool = False,
|
||||
dropout: float = 0.0,
|
||||
temb_channels: int = 512,
|
||||
norm_type: NormType = NormType.GROUP,
|
||||
causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.causality_axis = causality_axis
|
||||
|
||||
if self.causality_axis != CausalityAxis.NONE and norm_type == NormType.GROUP:
|
||||
raise ValueError("Causal ResnetBlock with GroupNorm is not supported.")
|
||||
self.in_channels = in_channels
|
||||
out_channels = in_channels if out_channels is None else out_channels
|
||||
self.out_channels = out_channels
|
||||
self.use_conv_shortcut = conv_shortcut
|
||||
|
||||
self.norm1 = build_normalization_layer(in_channels, normtype=norm_type)
|
||||
self.non_linearity = torch.nn.SiLU()
|
||||
self.conv1 = make_conv2d(in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis)
|
||||
if temb_channels > 0:
|
||||
self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
|
||||
self.norm2 = build_normalization_layer(out_channels, normtype=norm_type)
|
||||
self.dropout = torch.nn.Dropout(dropout)
|
||||
self.conv2 = make_conv2d(out_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis)
|
||||
if self.in_channels != self.out_channels:
|
||||
if self.use_conv_shortcut:
|
||||
self.conv_shortcut = make_conv2d(
|
||||
in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis
|
||||
)
|
||||
else:
|
||||
self.nin_shortcut = make_conv2d(
|
||||
in_channels, out_channels, kernel_size=1, stride=1, causality_axis=causality_axis
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
temb: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
h = x
|
||||
h = self.norm1(h)
|
||||
h = self.non_linearity(h)
|
||||
h = self.conv1(h)
|
||||
|
||||
if temb is not None:
|
||||
h = h + self.temb_proj(self.non_linearity(temb))[:, :, None, None]
|
||||
|
||||
h = self.norm2(h)
|
||||
h = self.non_linearity(h)
|
||||
h = self.dropout(h)
|
||||
h = self.conv2(h)
|
||||
|
||||
if self.in_channels != self.out_channels:
|
||||
x = self.conv_shortcut(x) if self.use_conv_shortcut else self.nin_shortcut(x)
|
||||
|
||||
return x + h
|
||||
@@ -0,0 +1,106 @@
|
||||
from typing import Set, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.model.audio_vae.attention import AttentionType, make_attn
|
||||
from ltx_core.model.audio_vae.causal_conv_2d import make_conv2d
|
||||
from ltx_core.model.audio_vae.causality_axis import CausalityAxis
|
||||
from ltx_core.model.audio_vae.resnet import ResnetBlock
|
||||
from ltx_core.model.common.normalization import NormType
|
||||
|
||||
|
||||
class Upsample(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
with_conv: bool,
|
||||
causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.with_conv = with_conv
|
||||
self.causality_axis = causality_axis
|
||||
if self.with_conv:
|
||||
self.conv = make_conv2d(in_channels, in_channels, kernel_size=3, stride=1, causality_axis=causality_axis)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
|
||||
if self.with_conv:
|
||||
x = self.conv(x)
|
||||
# Drop FIRST element in the causal axis to undo encoder's padding, while keeping the length 1 + 2 * n.
|
||||
# For example, if the input is [0, 1, 2], after interpolation, the output is [0, 0, 1, 1, 2, 2].
|
||||
# The causal convolution will pad the first element as [-, -, 0, 0, 1, 1, 2, 2],
|
||||
# So the output elements rely on the following windows:
|
||||
# 0: [-,-,0]
|
||||
# 1: [-,0,0]
|
||||
# 2: [0,0,1]
|
||||
# 3: [0,1,1]
|
||||
# 4: [1,1,2]
|
||||
# 5: [1,2,2]
|
||||
# Notice that the first and second elements in the output rely only on the first element in the input,
|
||||
# while all other elements rely on two elements in the input.
|
||||
# So we can drop the first element to undo the padding (rather than the last element).
|
||||
# This is a no-op for non-causal convolutions.
|
||||
match self.causality_axis:
|
||||
case CausalityAxis.NONE:
|
||||
pass # x remains unchanged
|
||||
case CausalityAxis.HEIGHT:
|
||||
x = x[:, :, 1:, :]
|
||||
case CausalityAxis.WIDTH:
|
||||
x = x[:, :, :, 1:]
|
||||
case CausalityAxis.WIDTH_COMPATIBILITY:
|
||||
pass # x remains unchanged
|
||||
case _:
|
||||
raise ValueError(f"Invalid causality_axis: {self.causality_axis}")
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def build_upsampling_path( # noqa: PLR0913
|
||||
*,
|
||||
ch: int,
|
||||
ch_mult: Tuple[int, ...],
|
||||
num_resolutions: int,
|
||||
num_res_blocks: int,
|
||||
resolution: int,
|
||||
temb_channels: int,
|
||||
dropout: float,
|
||||
norm_type: NormType,
|
||||
causality_axis: CausalityAxis,
|
||||
attn_type: AttentionType,
|
||||
attn_resolutions: Set[int],
|
||||
resamp_with_conv: bool,
|
||||
initial_block_channels: int,
|
||||
) -> tuple[torch.nn.ModuleList, int]:
|
||||
"""Build the upsampling path with residual blocks, attention, and upsampling layers."""
|
||||
up_modules = torch.nn.ModuleList()
|
||||
block_in = initial_block_channels
|
||||
curr_res = resolution // (2 ** (num_resolutions - 1))
|
||||
|
||||
for level in reversed(range(num_resolutions)):
|
||||
stage = torch.nn.Module()
|
||||
stage.block = torch.nn.ModuleList()
|
||||
stage.attn = torch.nn.ModuleList()
|
||||
block_out = ch * ch_mult[level]
|
||||
|
||||
for _ in range(num_res_blocks + 1):
|
||||
stage.block.append(
|
||||
ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_out,
|
||||
temb_channels=temb_channels,
|
||||
dropout=dropout,
|
||||
norm_type=norm_type,
|
||||
causality_axis=causality_axis,
|
||||
)
|
||||
)
|
||||
block_in = block_out
|
||||
if curr_res in attn_resolutions:
|
||||
stage.attn.append(make_attn(block_in, attn_type=attn_type, norm_type=norm_type))
|
||||
|
||||
if level != 0:
|
||||
stage.upsample = Upsample(block_in, resamp_with_conv, causality_axis=causality_axis)
|
||||
curr_res *= 2
|
||||
|
||||
up_modules.insert(0, stage)
|
||||
|
||||
return up_modules, block_in
|
||||
@@ -0,0 +1,123 @@
|
||||
import math
|
||||
from typing import List
|
||||
|
||||
import einops
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from ltx_core.model.audio_vae.resnet import LRELU_SLOPE, ResBlock1, ResBlock2
|
||||
|
||||
|
||||
class Vocoder(torch.nn.Module):
|
||||
"""
|
||||
Vocoder model for synthesizing audio from Mel spectrograms.
|
||||
Args:
|
||||
resblock_kernel_sizes: List of kernel sizes for the residual blocks.
|
||||
This value is read from the checkpoint at `config.vocoder.resblock_kernel_sizes`.
|
||||
upsample_rates: List of upsampling rates.
|
||||
This value is read from the checkpoint at `config.vocoder.upsample_rates`.
|
||||
upsample_kernel_sizes: List of kernel sizes for the upsampling layers.
|
||||
This value is read from the checkpoint at `config.vocoder.upsample_kernel_sizes`.
|
||||
resblock_dilation_sizes: List of dilation sizes for the residual blocks.
|
||||
This value is read from the checkpoint at `config.vocoder.resblock_dilation_sizes`.
|
||||
upsample_initial_channel: Initial number of channels for the upsampling layers.
|
||||
This value is read from the checkpoint at `config.vocoder.upsample_initial_channel`.
|
||||
stereo: Whether to use stereo output.
|
||||
This value is read from the checkpoint at `config.vocoder.stereo`.
|
||||
resblock: Type of residual block to use.
|
||||
This value is read from the checkpoint at `config.vocoder.resblock`.
|
||||
output_sample_rate: Waveform sample rate.
|
||||
This value is read from the checkpoint at `config.vocoder.output_sample_rate`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
resblock_kernel_sizes: List[int] | None = None,
|
||||
upsample_rates: List[int] | None = None,
|
||||
upsample_kernel_sizes: List[int] | None = None,
|
||||
resblock_dilation_sizes: List[List[int]] | None = None,
|
||||
upsample_initial_channel: int = 1024,
|
||||
stereo: bool = True,
|
||||
resblock: str = "1",
|
||||
output_sample_rate: int = 24000,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# Initialize default values if not provided. Note that mutable default values are not supported.
|
||||
if resblock_kernel_sizes is None:
|
||||
resblock_kernel_sizes = [3, 7, 11]
|
||||
if upsample_rates is None:
|
||||
upsample_rates = [6, 5, 2, 2, 2]
|
||||
if upsample_kernel_sizes is None:
|
||||
upsample_kernel_sizes = [16, 15, 8, 4, 4]
|
||||
if resblock_dilation_sizes is None:
|
||||
resblock_dilation_sizes = [[1, 3, 5], [1, 3, 5], [1, 3, 5]]
|
||||
|
||||
self.output_sample_rate = output_sample_rate
|
||||
self.num_kernels = len(resblock_kernel_sizes)
|
||||
self.num_upsamples = len(upsample_rates)
|
||||
in_channels = 128 if stereo else 64
|
||||
self.conv_pre = nn.Conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3)
|
||||
resblock_class = ResBlock1 if resblock == "1" else ResBlock2
|
||||
|
||||
self.ups = nn.ModuleList()
|
||||
for i, (stride, kernel_size) in enumerate(zip(upsample_rates, upsample_kernel_sizes, strict=True)):
|
||||
self.ups.append(
|
||||
nn.ConvTranspose1d(
|
||||
upsample_initial_channel // (2**i),
|
||||
upsample_initial_channel // (2 ** (i + 1)),
|
||||
kernel_size,
|
||||
stride,
|
||||
padding=(kernel_size - stride) // 2,
|
||||
)
|
||||
)
|
||||
|
||||
self.resblocks = nn.ModuleList()
|
||||
for i, _ in enumerate(self.ups):
|
||||
ch = upsample_initial_channel // (2 ** (i + 1))
|
||||
for kernel_size, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes, strict=True):
|
||||
self.resblocks.append(resblock_class(ch, kernel_size, dilations))
|
||||
|
||||
out_channels = 2 if stereo else 1
|
||||
final_channels = upsample_initial_channel // (2**self.num_upsamples)
|
||||
self.conv_post = nn.Conv1d(final_channels, out_channels, 7, 1, padding=3)
|
||||
|
||||
self.upsample_factor = math.prod(layer.stride[0] for layer in self.ups)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Forward pass of the vocoder.
|
||||
Args:
|
||||
x: Input Mel spectrogram tensor. Can be either:
|
||||
- 3D: (batch_size, time, mel_bins) for mono
|
||||
- 4D: (batch_size, 2, time, mel_bins) for stereo
|
||||
Returns:
|
||||
Audio waveform tensor of shape (batch_size, out_channels, audio_length)
|
||||
"""
|
||||
x = x.transpose(2, 3) # (batch, channels, time, mel_bins) -> (batch, channels, mel_bins, time)
|
||||
|
||||
if x.dim() == 4: # stereo
|
||||
assert x.shape[1] == 2, "Input must have 2 channels for stereo"
|
||||
x = einops.rearrange(x, "b s c t -> b (s c) t")
|
||||
|
||||
x = self.conv_pre(x)
|
||||
|
||||
for i in range(self.num_upsamples):
|
||||
x = F.leaky_relu(x, LRELU_SLOPE)
|
||||
x = self.ups[i](x)
|
||||
start = i * self.num_kernels
|
||||
end = start + self.num_kernels
|
||||
|
||||
# Evaluate all resblocks with the same input tensor so they can run
|
||||
# independently (and thus in parallel on accelerator hardware) before
|
||||
# aggregating their outputs via mean.
|
||||
block_outputs = torch.stack(
|
||||
[self.resblocks[idx](x) for idx in range(start, end)],
|
||||
dim=0,
|
||||
)
|
||||
|
||||
x = block_outputs.mean(dim=0)
|
||||
|
||||
x = self.conv_post(F.leaky_relu(x))
|
||||
return torch.tanh(x)
|
||||
@@ -0,0 +1,9 @@
|
||||
"""Common model utilities."""
|
||||
|
||||
from ltx_core.model.common.normalization import NormType, PixelNorm, build_normalization_layer
|
||||
|
||||
__all__ = [
|
||||
"NormType",
|
||||
"PixelNorm",
|
||||
"build_normalization_layer",
|
||||
]
|
||||
@@ -0,0 +1,59 @@
|
||||
from enum import Enum
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class NormType(Enum):
|
||||
"""Normalization layer types: GROUP (GroupNorm) or PIXEL (per-location RMS norm)."""
|
||||
|
||||
GROUP = "group"
|
||||
PIXEL = "pixel"
|
||||
|
||||
|
||||
class PixelNorm(nn.Module):
|
||||
"""
|
||||
Per-pixel (per-location) RMS normalization layer.
|
||||
For each element along the chosen dimension, this layer normalizes the tensor
|
||||
by the root-mean-square of its values across that dimension:
|
||||
y = x / sqrt(mean(x^2, dim=dim, keepdim=True) + eps)
|
||||
"""
|
||||
|
||||
def __init__(self, dim: int = 1, eps: float = 1e-8) -> None:
|
||||
"""
|
||||
Args:
|
||||
dim: Dimension along which to compute the RMS (typically channels).
|
||||
eps: Small constant added for numerical stability.
|
||||
"""
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Apply RMS normalization along the configured dimension.
|
||||
"""
|
||||
# Compute mean of squared values along `dim`, keep dimensions for broadcasting.
|
||||
mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True)
|
||||
# Normalize by the root-mean-square (RMS).
|
||||
rms = torch.sqrt(mean_sq + self.eps)
|
||||
return x / rms
|
||||
|
||||
|
||||
def build_normalization_layer(
|
||||
in_channels: int, *, num_groups: int = 32, normtype: NormType = NormType.GROUP
|
||||
) -> nn.Module:
|
||||
"""
|
||||
Create a normalization layer based on the normalization type.
|
||||
Args:
|
||||
in_channels: Number of input channels
|
||||
num_groups: Number of groups for group normalization
|
||||
normtype: Type of normalization: "group" or "pixel"
|
||||
Returns:
|
||||
A normalization layer
|
||||
"""
|
||||
if normtype == NormType.GROUP:
|
||||
return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
|
||||
if normtype == NormType.PIXEL:
|
||||
return PixelNorm(dim=1, eps=1e-6)
|
||||
raise ValueError(f"Invalid normalization type: {normtype}")
|
||||
@@ -0,0 +1,10 @@
|
||||
from typing import Protocol, TypeVar
|
||||
|
||||
ModelType = TypeVar("ModelType")
|
||||
|
||||
|
||||
class ModelConfigurator(Protocol[ModelType]):
|
||||
"""Protocol for model loader classes that instantiates models from a configuration dictionary."""
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict) -> ModelType: ...
|
||||
@@ -0,0 +1,24 @@
|
||||
"""Transformer model components."""
|
||||
|
||||
from ltx_core.model.transformer.modality import Modality
|
||||
from ltx_core.model.transformer.model import LTXModel, X0Model
|
||||
from ltx_core.model.transformer.model_configurator import (
|
||||
LTXV_MODEL_COMFY_RENAMING_MAP,
|
||||
LTXV_MODEL_COMFY_RENAMING_WITH_TRANSFORMER_LINEAR_DOWNCAST_MAP,
|
||||
UPCAST_DURING_INFERENCE,
|
||||
LTXModelConfigurator,
|
||||
LTXVideoOnlyModelConfigurator,
|
||||
UpcastWithStochasticRounding,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"LTXV_MODEL_COMFY_RENAMING_MAP",
|
||||
"LTXV_MODEL_COMFY_RENAMING_WITH_TRANSFORMER_LINEAR_DOWNCAST_MAP",
|
||||
"UPCAST_DURING_INFERENCE",
|
||||
"LTXModel",
|
||||
"LTXModelConfigurator",
|
||||
"LTXVideoOnlyModelConfigurator",
|
||||
"Modality",
|
||||
"UpcastWithStochasticRounding",
|
||||
"X0Model",
|
||||
]
|
||||
@@ -0,0 +1,34 @@
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.model.transformer.timestep_embedding import PixArtAlphaCombinedTimestepSizeEmbeddings
|
||||
|
||||
|
||||
class AdaLayerNormSingle(torch.nn.Module):
|
||||
r"""
|
||||
Norm layer adaptive layer norm single (adaLN-single).
|
||||
As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3).
|
||||
Parameters:
|
||||
embedding_dim (`int`): The size of each embedding vector.
|
||||
use_additional_conditions (`bool`): To use additional conditions for normalization or not.
|
||||
"""
|
||||
|
||||
def __init__(self, embedding_dim: int, embedding_coefficient: int = 6):
|
||||
super().__init__()
|
||||
|
||||
self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings(
|
||||
embedding_dim,
|
||||
size_emb_dim=embedding_dim // 3,
|
||||
)
|
||||
|
||||
self.silu = torch.nn.SiLU()
|
||||
self.linear = torch.nn.Linear(embedding_dim, embedding_coefficient * embedding_dim, bias=True)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
timestep: torch.Tensor,
|
||||
hidden_dtype: Optional[torch.dtype] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype)
|
||||
return self.linear(self.silu(embedded_timestep)), embedded_timestep
|
||||
@@ -0,0 +1,195 @@
|
||||
from enum import Enum
|
||||
from typing import Protocol
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.model.transformer.rope import LTXRopeType, apply_rotary_emb
|
||||
|
||||
memory_efficient_attention = None
|
||||
flash_attn_interface = None
|
||||
try:
|
||||
from xformers.ops import memory_efficient_attention
|
||||
except ImportError:
|
||||
memory_efficient_attention = None
|
||||
try:
|
||||
# FlashAttention3 and XFormersAttention cannot be used together
|
||||
if memory_efficient_attention is None:
|
||||
import flash_attn_interface
|
||||
except ImportError:
|
||||
flash_attn_interface = None
|
||||
|
||||
|
||||
class AttentionCallable(Protocol):
|
||||
def __call__(
|
||||
self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
|
||||
) -> torch.Tensor: ...
|
||||
|
||||
|
||||
class PytorchAttention(AttentionCallable):
|
||||
def __call__(
|
||||
self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
|
||||
) -> torch.Tensor:
|
||||
b, _, dim_head = q.shape
|
||||
dim_head //= heads
|
||||
q, k, v = (t.view(b, -1, heads, dim_head).transpose(1, 2) for t in (q, k, v))
|
||||
|
||||
if mask is not None:
|
||||
# add a batch dimension if there isn't already one
|
||||
if mask.ndim == 2:
|
||||
mask = mask.unsqueeze(0)
|
||||
# add a heads dimension if there isn't already one
|
||||
if mask.ndim == 3:
|
||||
mask = mask.unsqueeze(1)
|
||||
|
||||
out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
|
||||
out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
|
||||
return out
|
||||
|
||||
|
||||
class XFormersAttention(AttentionCallable):
|
||||
def __call__(
|
||||
self,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
heads: int,
|
||||
mask: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
if memory_efficient_attention is None:
|
||||
raise RuntimeError("XFormersAttention was selected but `xformers` is not installed.")
|
||||
|
||||
b, _, dim_head = q.shape
|
||||
dim_head //= heads
|
||||
|
||||
# xformers expects [B, M, H, K]
|
||||
q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v))
|
||||
|
||||
if mask is not None:
|
||||
# add a singleton batch dimension
|
||||
if mask.ndim == 2:
|
||||
mask = mask.unsqueeze(0)
|
||||
# add a singleton heads dimension
|
||||
if mask.ndim == 3:
|
||||
mask = mask.unsqueeze(1)
|
||||
# pad to a multiple of 8
|
||||
pad = 8 - mask.shape[-1] % 8
|
||||
# the xformers docs says that it's allowed to have a mask of shape (1, Nq, Nk)
|
||||
# but when using separated heads, the shape has to be (B, H, Nq, Nk)
|
||||
# in flux, this matrix ends up being over 1GB
|
||||
# here, we create a mask with the same batch/head size as the input mask (potentially singleton or full)
|
||||
mask_out = torch.empty(
|
||||
[mask.shape[0], mask.shape[1], q.shape[1], mask.shape[-1] + pad], dtype=q.dtype, device=q.device
|
||||
)
|
||||
|
||||
mask_out[..., : mask.shape[-1]] = mask
|
||||
# doesn't this remove the padding again??
|
||||
mask = mask_out[..., : mask.shape[-1]]
|
||||
mask = mask.expand(b, heads, -1, -1)
|
||||
|
||||
out = memory_efficient_attention(q.to(v.dtype), k.to(v.dtype), v, attn_bias=mask, p=0.0)
|
||||
out = out.reshape(b, -1, heads * dim_head)
|
||||
return out
|
||||
|
||||
|
||||
class FlashAttention3(AttentionCallable):
|
||||
def __call__(
|
||||
self,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
heads: int,
|
||||
mask: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
if flash_attn_interface is None:
|
||||
raise RuntimeError("FlashAttention3 was selected but `FlashAttention3` is not installed.")
|
||||
|
||||
b, _, dim_head = q.shape
|
||||
dim_head //= heads
|
||||
|
||||
q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v))
|
||||
|
||||
if mask is not None:
|
||||
raise NotImplementedError("Mask is not supported for FlashAttention3")
|
||||
|
||||
out = flash_attn_interface.flash_attn_func(q.to(v.dtype), k.to(v.dtype), v)
|
||||
out = out.reshape(b, -1, heads * dim_head)
|
||||
return out
|
||||
|
||||
|
||||
class AttentionFunction(Enum):
|
||||
PYTORCH = "pytorch"
|
||||
XFORMERS = "xformers"
|
||||
FLASH_ATTENTION_3 = "flash_attention_3"
|
||||
DEFAULT = "default"
|
||||
|
||||
def __call__(
|
||||
self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
|
||||
) -> torch.Tensor:
|
||||
if self is AttentionFunction.PYTORCH:
|
||||
return PytorchAttention()(q, k, v, heads, mask)
|
||||
elif self is AttentionFunction.XFORMERS:
|
||||
return XFormersAttention()(q, k, v, heads, mask)
|
||||
elif self is AttentionFunction.FLASH_ATTENTION_3:
|
||||
return FlashAttention3()(q, k, v, heads, mask)
|
||||
else:
|
||||
# Default behavior: XFormers if installed else - PyTorch
|
||||
return (
|
||||
XFormersAttention()(q, k, v, heads, mask)
|
||||
if memory_efficient_attention is not None
|
||||
else PytorchAttention()(q, k, v, heads, mask)
|
||||
)
|
||||
|
||||
|
||||
class Attention(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
query_dim: int,
|
||||
context_dim: int | None = None,
|
||||
heads: int = 8,
|
||||
dim_head: int = 64,
|
||||
norm_eps: float = 1e-6,
|
||||
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
|
||||
attention_function: AttentionCallable | AttentionFunction = AttentionFunction.DEFAULT,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.rope_type = rope_type
|
||||
self.attention_function = attention_function
|
||||
|
||||
inner_dim = dim_head * heads
|
||||
context_dim = query_dim if context_dim is None else context_dim
|
||||
|
||||
self.heads = heads
|
||||
self.dim_head = dim_head
|
||||
|
||||
self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
|
||||
self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
|
||||
|
||||
self.to_q = torch.nn.Linear(query_dim, inner_dim, bias=True)
|
||||
self.to_k = torch.nn.Linear(context_dim, inner_dim, bias=True)
|
||||
self.to_v = torch.nn.Linear(context_dim, inner_dim, bias=True)
|
||||
|
||||
self.to_out = torch.nn.Sequential(torch.nn.Linear(inner_dim, query_dim, bias=True), torch.nn.Identity())
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
context: torch.Tensor | None = None,
|
||||
mask: torch.Tensor | None = None,
|
||||
pe: torch.Tensor | None = None,
|
||||
k_pe: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
q = self.to_q(x)
|
||||
context = x if context is None else context
|
||||
k = self.to_k(context)
|
||||
v = self.to_v(context)
|
||||
|
||||
q = self.q_norm(q)
|
||||
k = self.k_norm(k)
|
||||
|
||||
if pe is not None:
|
||||
q = apply_rotary_emb(q, pe, self.rope_type)
|
||||
k = apply_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type)
|
||||
|
||||
# attention_function can be an enum *or* a custom callable
|
||||
out = self.attention_function(q, k, v, self.heads, mask)
|
||||
return self.to_out(out)
|
||||
@@ -0,0 +1,15 @@
|
||||
import torch
|
||||
|
||||
from ltx_core.model.transformer.gelu_approx import GELUApprox
|
||||
|
||||
|
||||
class FeedForward(torch.nn.Module):
|
||||
def __init__(self, dim: int, dim_out: int, mult: int = 4) -> None:
|
||||
super().__init__()
|
||||
inner_dim = int(dim * mult)
|
||||
project_in = GELUApprox(dim, inner_dim)
|
||||
|
||||
self.net = torch.nn.Sequential(project_in, torch.nn.Identity(), torch.nn.Linear(inner_dim, dim_out))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.net(x)
|
||||
@@ -0,0 +1,10 @@
|
||||
import torch
|
||||
|
||||
|
||||
class GELUApprox(torch.nn.Module):
|
||||
def __init__(self, dim_in: int, dim_out: int) -> None:
|
||||
super().__init__()
|
||||
self.proj = torch.nn.Linear(dim_in, dim_out)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return torch.nn.functional.gelu(self.proj(x), approximate="tanh")
|
||||
@@ -0,0 +1,23 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Modality:
|
||||
"""
|
||||
Input data for a single modality (video or audio) in the transformer.
|
||||
Bundles the latent tokens, timestep embeddings, positional information,
|
||||
and text conditioning context for processing by the diffusion transformer.
|
||||
"""
|
||||
|
||||
latent: (
|
||||
torch.Tensor
|
||||
) # Shape: (B, T, D) where B is the batch size, T is the number of tokens, and D is input dimension
|
||||
timesteps: torch.Tensor # Shape: (B, T) where T is the number of timesteps
|
||||
positions: (
|
||||
torch.Tensor
|
||||
) # Shape: (B, 3, T) for video, where 3 is the number of dimensions and T is the number of tokens
|
||||
context: torch.Tensor
|
||||
enabled: bool = True
|
||||
context_mask: torch.Tensor | None = None
|
||||
@@ -0,0 +1,468 @@
|
||||
from enum import Enum
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.guidance.perturbations import BatchedPerturbationConfig
|
||||
from ltx_core.model.transformer.adaln import AdaLayerNormSingle
|
||||
from ltx_core.model.transformer.attention import AttentionCallable, AttentionFunction
|
||||
from ltx_core.model.transformer.modality import Modality
|
||||
from ltx_core.model.transformer.rope import LTXRopeType
|
||||
from ltx_core.model.transformer.text_projection import PixArtAlphaTextProjection
|
||||
from ltx_core.model.transformer.transformer import BasicAVTransformerBlock, TransformerConfig
|
||||
from ltx_core.model.transformer.transformer_args import (
|
||||
MultiModalTransformerArgsPreprocessor,
|
||||
TransformerArgs,
|
||||
TransformerArgsPreprocessor,
|
||||
)
|
||||
from ltx_core.utils import to_denoised
|
||||
|
||||
|
||||
class LTXModelType(Enum):
|
||||
AudioVideo = "ltx av model"
|
||||
VideoOnly = "ltx video only model"
|
||||
AudioOnly = "ltx audio only model"
|
||||
|
||||
def is_video_enabled(self) -> bool:
|
||||
return self in (LTXModelType.AudioVideo, LTXModelType.VideoOnly)
|
||||
|
||||
def is_audio_enabled(self) -> bool:
|
||||
return self in (LTXModelType.AudioVideo, LTXModelType.AudioOnly)
|
||||
|
||||
|
||||
class LTXModel(torch.nn.Module):
|
||||
"""
|
||||
LTX model transformer implementation.
|
||||
This class implements the transformer blocks for the LTX model.
|
||||
"""
|
||||
|
||||
def __init__( # noqa: PLR0913
|
||||
self,
|
||||
*,
|
||||
model_type: LTXModelType = LTXModelType.AudioVideo,
|
||||
num_attention_heads: int = 32,
|
||||
attention_head_dim: int = 128,
|
||||
in_channels: int = 128,
|
||||
out_channels: int = 128,
|
||||
num_layers: int = 48,
|
||||
cross_attention_dim: int = 4096,
|
||||
norm_eps: float = 1e-06,
|
||||
attention_type: AttentionFunction | AttentionCallable = AttentionFunction.DEFAULT,
|
||||
caption_channels: int = 3840,
|
||||
positional_embedding_theta: float = 10000.0,
|
||||
positional_embedding_max_pos: list[int] | None = None,
|
||||
timestep_scale_multiplier: int = 1000,
|
||||
use_middle_indices_grid: bool = True,
|
||||
audio_num_attention_heads: int = 32,
|
||||
audio_attention_head_dim: int = 64,
|
||||
audio_in_channels: int = 128,
|
||||
audio_out_channels: int = 128,
|
||||
audio_cross_attention_dim: int = 2048,
|
||||
audio_positional_embedding_max_pos: list[int] | None = None,
|
||||
av_ca_timestep_scale_multiplier: int = 1,
|
||||
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
|
||||
double_precision_rope: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self._enable_gradient_checkpointing = False
|
||||
self.use_middle_indices_grid = use_middle_indices_grid
|
||||
self.rope_type = rope_type
|
||||
self.double_precision_rope = double_precision_rope
|
||||
self.timestep_scale_multiplier = timestep_scale_multiplier
|
||||
self.positional_embedding_theta = positional_embedding_theta
|
||||
self.model_type = model_type
|
||||
cross_pe_max_pos = None
|
||||
if model_type.is_video_enabled():
|
||||
if positional_embedding_max_pos is None:
|
||||
positional_embedding_max_pos = [20, 2048, 2048]
|
||||
self.positional_embedding_max_pos = positional_embedding_max_pos
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.inner_dim = num_attention_heads * attention_head_dim
|
||||
self._init_video(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
caption_channels=caption_channels,
|
||||
norm_eps=norm_eps,
|
||||
)
|
||||
|
||||
if model_type.is_audio_enabled():
|
||||
if audio_positional_embedding_max_pos is None:
|
||||
audio_positional_embedding_max_pos = [20]
|
||||
self.audio_positional_embedding_max_pos = audio_positional_embedding_max_pos
|
||||
self.audio_num_attention_heads = audio_num_attention_heads
|
||||
self.audio_inner_dim = self.audio_num_attention_heads * audio_attention_head_dim
|
||||
self._init_audio(
|
||||
in_channels=audio_in_channels,
|
||||
out_channels=audio_out_channels,
|
||||
caption_channels=caption_channels,
|
||||
norm_eps=norm_eps,
|
||||
)
|
||||
|
||||
if model_type.is_video_enabled() and model_type.is_audio_enabled():
|
||||
cross_pe_max_pos = max(self.positional_embedding_max_pos[0], self.audio_positional_embedding_max_pos[0])
|
||||
self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier
|
||||
self.audio_cross_attention_dim = audio_cross_attention_dim
|
||||
self._init_audio_video(num_scale_shift_values=4)
|
||||
|
||||
self._init_preprocessors(cross_pe_max_pos)
|
||||
# Initialize transformer blocks
|
||||
self._init_transformer_blocks(
|
||||
num_layers=num_layers,
|
||||
attention_head_dim=attention_head_dim if model_type.is_video_enabled() else 0,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
audio_attention_head_dim=audio_attention_head_dim if model_type.is_audio_enabled() else 0,
|
||||
audio_cross_attention_dim=audio_cross_attention_dim,
|
||||
norm_eps=norm_eps,
|
||||
attention_type=attention_type,
|
||||
)
|
||||
|
||||
def _init_video(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
caption_channels: int,
|
||||
norm_eps: float,
|
||||
) -> None:
|
||||
"""Initialize video-specific components."""
|
||||
# Video input components
|
||||
self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True)
|
||||
|
||||
self.adaln_single = AdaLayerNormSingle(self.inner_dim)
|
||||
|
||||
# Video caption projection
|
||||
self.caption_projection = PixArtAlphaTextProjection(
|
||||
in_features=caption_channels,
|
||||
hidden_size=self.inner_dim,
|
||||
)
|
||||
|
||||
# Video output components
|
||||
self.scale_shift_table = torch.nn.Parameter(torch.empty(2, self.inner_dim))
|
||||
self.norm_out = torch.nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=norm_eps)
|
||||
self.proj_out = torch.nn.Linear(self.inner_dim, out_channels)
|
||||
|
||||
def _init_audio(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
caption_channels: int,
|
||||
norm_eps: float,
|
||||
) -> None:
|
||||
"""Initialize audio-specific components."""
|
||||
|
||||
# Audio input components
|
||||
self.audio_patchify_proj = torch.nn.Linear(in_channels, self.audio_inner_dim, bias=True)
|
||||
|
||||
self.audio_adaln_single = AdaLayerNormSingle(
|
||||
self.audio_inner_dim,
|
||||
)
|
||||
|
||||
# Audio caption projection
|
||||
self.audio_caption_projection = PixArtAlphaTextProjection(
|
||||
in_features=caption_channels,
|
||||
hidden_size=self.audio_inner_dim,
|
||||
)
|
||||
|
||||
# Audio output components
|
||||
self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(2, self.audio_inner_dim))
|
||||
self.audio_norm_out = torch.nn.LayerNorm(self.audio_inner_dim, elementwise_affine=False, eps=norm_eps)
|
||||
self.audio_proj_out = torch.nn.Linear(self.audio_inner_dim, out_channels)
|
||||
|
||||
def _init_audio_video(
|
||||
self,
|
||||
num_scale_shift_values: int,
|
||||
) -> None:
|
||||
"""Initialize audio-video cross-attention components."""
|
||||
self.av_ca_video_scale_shift_adaln_single = AdaLayerNormSingle(
|
||||
self.inner_dim,
|
||||
embedding_coefficient=num_scale_shift_values,
|
||||
)
|
||||
|
||||
self.av_ca_audio_scale_shift_adaln_single = AdaLayerNormSingle(
|
||||
self.audio_inner_dim,
|
||||
embedding_coefficient=num_scale_shift_values,
|
||||
)
|
||||
|
||||
self.av_ca_a2v_gate_adaln_single = AdaLayerNormSingle(
|
||||
self.inner_dim,
|
||||
embedding_coefficient=1,
|
||||
)
|
||||
|
||||
self.av_ca_v2a_gate_adaln_single = AdaLayerNormSingle(
|
||||
self.audio_inner_dim,
|
||||
embedding_coefficient=1,
|
||||
)
|
||||
|
||||
def _init_preprocessors(
|
||||
self,
|
||||
cross_pe_max_pos: int | None = None,
|
||||
) -> None:
|
||||
"""Initialize preprocessors for LTX."""
|
||||
|
||||
if self.model_type.is_video_enabled() and self.model_type.is_audio_enabled():
|
||||
self.video_args_preprocessor = MultiModalTransformerArgsPreprocessor(
|
||||
patchify_proj=self.patchify_proj,
|
||||
adaln=self.adaln_single,
|
||||
caption_projection=self.caption_projection,
|
||||
cross_scale_shift_adaln=self.av_ca_video_scale_shift_adaln_single,
|
||||
cross_gate_adaln=self.av_ca_a2v_gate_adaln_single,
|
||||
inner_dim=self.inner_dim,
|
||||
max_pos=self.positional_embedding_max_pos,
|
||||
num_attention_heads=self.num_attention_heads,
|
||||
cross_pe_max_pos=cross_pe_max_pos,
|
||||
use_middle_indices_grid=self.use_middle_indices_grid,
|
||||
audio_cross_attention_dim=self.audio_cross_attention_dim,
|
||||
timestep_scale_multiplier=self.timestep_scale_multiplier,
|
||||
double_precision_rope=self.double_precision_rope,
|
||||
positional_embedding_theta=self.positional_embedding_theta,
|
||||
rope_type=self.rope_type,
|
||||
av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier,
|
||||
)
|
||||
self.audio_args_preprocessor = MultiModalTransformerArgsPreprocessor(
|
||||
patchify_proj=self.audio_patchify_proj,
|
||||
adaln=self.audio_adaln_single,
|
||||
caption_projection=self.audio_caption_projection,
|
||||
cross_scale_shift_adaln=self.av_ca_audio_scale_shift_adaln_single,
|
||||
cross_gate_adaln=self.av_ca_v2a_gate_adaln_single,
|
||||
inner_dim=self.audio_inner_dim,
|
||||
max_pos=self.audio_positional_embedding_max_pos,
|
||||
num_attention_heads=self.audio_num_attention_heads,
|
||||
cross_pe_max_pos=cross_pe_max_pos,
|
||||
use_middle_indices_grid=self.use_middle_indices_grid,
|
||||
audio_cross_attention_dim=self.audio_cross_attention_dim,
|
||||
timestep_scale_multiplier=self.timestep_scale_multiplier,
|
||||
double_precision_rope=self.double_precision_rope,
|
||||
positional_embedding_theta=self.positional_embedding_theta,
|
||||
rope_type=self.rope_type,
|
||||
av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier,
|
||||
)
|
||||
elif self.model_type.is_video_enabled():
|
||||
self.video_args_preprocessor = TransformerArgsPreprocessor(
|
||||
patchify_proj=self.patchify_proj,
|
||||
adaln=self.adaln_single,
|
||||
caption_projection=self.caption_projection,
|
||||
inner_dim=self.inner_dim,
|
||||
max_pos=self.positional_embedding_max_pos,
|
||||
num_attention_heads=self.num_attention_heads,
|
||||
use_middle_indices_grid=self.use_middle_indices_grid,
|
||||
timestep_scale_multiplier=self.timestep_scale_multiplier,
|
||||
double_precision_rope=self.double_precision_rope,
|
||||
positional_embedding_theta=self.positional_embedding_theta,
|
||||
rope_type=self.rope_type,
|
||||
)
|
||||
elif self.model_type.is_audio_enabled():
|
||||
self.audio_args_preprocessor = TransformerArgsPreprocessor(
|
||||
patchify_proj=self.audio_patchify_proj,
|
||||
adaln=self.audio_adaln_single,
|
||||
caption_projection=self.audio_caption_projection,
|
||||
inner_dim=self.audio_inner_dim,
|
||||
max_pos=self.audio_positional_embedding_max_pos,
|
||||
num_attention_heads=self.audio_num_attention_heads,
|
||||
use_middle_indices_grid=self.use_middle_indices_grid,
|
||||
timestep_scale_multiplier=self.timestep_scale_multiplier,
|
||||
double_precision_rope=self.double_precision_rope,
|
||||
positional_embedding_theta=self.positional_embedding_theta,
|
||||
rope_type=self.rope_type,
|
||||
)
|
||||
|
||||
def _init_transformer_blocks(
|
||||
self,
|
||||
num_layers: int,
|
||||
attention_head_dim: int,
|
||||
cross_attention_dim: int,
|
||||
audio_attention_head_dim: int,
|
||||
audio_cross_attention_dim: int,
|
||||
norm_eps: float,
|
||||
attention_type: AttentionFunction | AttentionCallable,
|
||||
) -> None:
|
||||
"""Initialize transformer blocks for LTX."""
|
||||
video_config = (
|
||||
TransformerConfig(
|
||||
dim=self.inner_dim,
|
||||
heads=self.num_attention_heads,
|
||||
d_head=attention_head_dim,
|
||||
context_dim=cross_attention_dim,
|
||||
)
|
||||
if self.model_type.is_video_enabled()
|
||||
else None
|
||||
)
|
||||
audio_config = (
|
||||
TransformerConfig(
|
||||
dim=self.audio_inner_dim,
|
||||
heads=self.audio_num_attention_heads,
|
||||
d_head=audio_attention_head_dim,
|
||||
context_dim=audio_cross_attention_dim,
|
||||
)
|
||||
if self.model_type.is_audio_enabled()
|
||||
else None
|
||||
)
|
||||
self.transformer_blocks = torch.nn.ModuleList(
|
||||
[
|
||||
BasicAVTransformerBlock(
|
||||
idx=idx,
|
||||
video=video_config,
|
||||
audio=audio_config,
|
||||
rope_type=self.rope_type,
|
||||
norm_eps=norm_eps,
|
||||
attention_function=attention_type,
|
||||
)
|
||||
for idx in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
def set_gradient_checkpointing(self, enable: bool) -> None:
|
||||
"""Enable or disable gradient checkpointing for transformer blocks.
|
||||
Gradient checkpointing trades compute for memory by recomputing activations
|
||||
during the backward pass instead of storing them. This can significantly
|
||||
reduce memory usage at the cost of ~20-30% slower training.
|
||||
Args:
|
||||
enable: Whether to enable gradient checkpointing
|
||||
"""
|
||||
self._enable_gradient_checkpointing = enable
|
||||
|
||||
def _process_transformer_blocks(
|
||||
self,
|
||||
video: TransformerArgs | None,
|
||||
audio: TransformerArgs | None,
|
||||
perturbations: BatchedPerturbationConfig,
|
||||
) -> tuple[TransformerArgs, TransformerArgs]:
|
||||
"""Process transformer blocks for LTXAV."""
|
||||
|
||||
# Process transformer blocks
|
||||
for block in self.transformer_blocks:
|
||||
if self._enable_gradient_checkpointing and self.training:
|
||||
# Use gradient checkpointing to save memory during training.
|
||||
# With use_reentrant=False, we can pass dataclasses directly -
|
||||
# PyTorch will track all tensor leaves in the computation graph.
|
||||
video, audio = torch.utils.checkpoint.checkpoint(
|
||||
block,
|
||||
video,
|
||||
audio,
|
||||
perturbations,
|
||||
use_reentrant=False,
|
||||
)
|
||||
else:
|
||||
video, audio = block(
|
||||
video=video,
|
||||
audio=audio,
|
||||
perturbations=perturbations,
|
||||
)
|
||||
|
||||
return video, audio
|
||||
|
||||
def _process_output(
|
||||
self,
|
||||
scale_shift_table: torch.Tensor,
|
||||
norm_out: torch.nn.LayerNorm,
|
||||
proj_out: torch.nn.Linear,
|
||||
x: torch.Tensor,
|
||||
embedded_timestep: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Process output for LTXV."""
|
||||
# Apply scale-shift modulation
|
||||
scale_shift_values = (
|
||||
scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) + embedded_timestep[:, :, None]
|
||||
)
|
||||
shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
|
||||
|
||||
x = norm_out(x)
|
||||
x = x * (1 + scale) + shift
|
||||
x = proj_out(x)
|
||||
return x
|
||||
|
||||
def forward(
|
||||
self, video: Modality | None, audio: Modality | None, perturbations: BatchedPerturbationConfig
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Forward pass for LTX models.
|
||||
Returns:
|
||||
Processed output tensors
|
||||
"""
|
||||
if not self.model_type.is_video_enabled() and video is not None:
|
||||
raise ValueError("Video is not enabled for this model")
|
||||
if not self.model_type.is_audio_enabled() and audio is not None:
|
||||
raise ValueError("Audio is not enabled for this model")
|
||||
|
||||
video_args = self.video_args_preprocessor.prepare(video) if video is not None else None
|
||||
audio_args = self.audio_args_preprocessor.prepare(audio) if audio is not None else None
|
||||
# Process transformer blocks
|
||||
video_out, audio_out = self._process_transformer_blocks(
|
||||
video=video_args,
|
||||
audio=audio_args,
|
||||
perturbations=perturbations,
|
||||
)
|
||||
|
||||
# Process output
|
||||
vx = (
|
||||
self._process_output(
|
||||
self.scale_shift_table, self.norm_out, self.proj_out, video_out.x, video_out.embedded_timestep
|
||||
)
|
||||
if video_out is not None
|
||||
else None
|
||||
)
|
||||
ax = (
|
||||
self._process_output(
|
||||
self.audio_scale_shift_table,
|
||||
self.audio_norm_out,
|
||||
self.audio_proj_out,
|
||||
audio_out.x,
|
||||
audio_out.embedded_timestep,
|
||||
)
|
||||
if audio_out is not None
|
||||
else None
|
||||
)
|
||||
return vx, ax
|
||||
|
||||
|
||||
class LegacyX0Model(torch.nn.Module):
|
||||
"""
|
||||
Legacy X0 model implementation.
|
||||
Returns fully denoised output based on the velocities produced by the base model.
|
||||
"""
|
||||
|
||||
def __init__(self, velocity_model: LTXModel):
|
||||
super().__init__()
|
||||
self.velocity_model = velocity_model
|
||||
|
||||
def forward(
|
||||
self,
|
||||
video: Modality | None,
|
||||
audio: Modality | None,
|
||||
perturbations: BatchedPerturbationConfig,
|
||||
sigma: float,
|
||||
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
|
||||
"""
|
||||
Denoise the video and audio according to the sigma.
|
||||
Returns:
|
||||
Denoised video and audio
|
||||
"""
|
||||
vx, ax = self.velocity_model(video, audio, perturbations)
|
||||
denoised_video = to_denoised(video.latent, vx, sigma) if vx is not None else None
|
||||
denoised_audio = to_denoised(audio.latent, ax, sigma) if ax is not None else None
|
||||
return denoised_video, denoised_audio
|
||||
|
||||
|
||||
class X0Model(torch.nn.Module):
|
||||
"""
|
||||
X0 model implementation.
|
||||
Returns fully denoised outputs based on the velocities produced by the base model.
|
||||
Applies scaled denoising to the video and audio according to the timesteps = sigma * denoising_mask.
|
||||
"""
|
||||
|
||||
def __init__(self, velocity_model: LTXModel):
|
||||
super().__init__()
|
||||
self.velocity_model = velocity_model
|
||||
|
||||
def forward(
|
||||
self,
|
||||
video: Modality | None,
|
||||
audio: Modality | None,
|
||||
perturbations: BatchedPerturbationConfig,
|
||||
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
|
||||
"""
|
||||
Denoise the video and audio according to the sigma.
|
||||
Returns:
|
||||
Denoised video and audio
|
||||
"""
|
||||
vx, ax = self.velocity_model(video, audio, perturbations)
|
||||
denoised_video = to_denoised(video.latent, vx, video.timesteps) if vx is not None else None
|
||||
denoised_audio = to_denoised(audio.latent, ax, audio.timesteps) if ax is not None else None
|
||||
return denoised_video, denoised_audio
|
||||
@@ -0,0 +1,237 @@
|
||||
import torch
|
||||
|
||||
from ltx_core.loader.fuse_loras import fused_add_round_launch
|
||||
from ltx_core.loader.module_ops import ModuleOps
|
||||
from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps
|
||||
from ltx_core.model.model_protocol import ModelConfigurator
|
||||
from ltx_core.model.transformer.attention import AttentionFunction
|
||||
from ltx_core.model.transformer.model import LTXModel, LTXModelType
|
||||
from ltx_core.model.transformer.rope import LTXRopeType
|
||||
from ltx_core.utils import check_config_value
|
||||
|
||||
|
||||
class LTXModelConfigurator(ModelConfigurator[LTXModel]):
|
||||
"""
|
||||
Configurator for LTX model.
|
||||
Used to create an LTX model from a configuration dictionary.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def from_config(cls: type[LTXModel], config: dict) -> LTXModel:
|
||||
config = config.get("transformer", {})
|
||||
|
||||
check_config_value(config, "dropout", 0.0)
|
||||
check_config_value(config, "attention_bias", True)
|
||||
check_config_value(config, "num_vector_embeds", None)
|
||||
check_config_value(config, "activation_fn", "gelu-approximate")
|
||||
check_config_value(config, "num_embeds_ada_norm", 1000)
|
||||
check_config_value(config, "use_linear_projection", False)
|
||||
check_config_value(config, "only_cross_attention", False)
|
||||
check_config_value(config, "cross_attention_norm", True)
|
||||
check_config_value(config, "double_self_attention", False)
|
||||
check_config_value(config, "upcast_attention", False)
|
||||
check_config_value(config, "standardization_norm", "rms_norm")
|
||||
check_config_value(config, "norm_elementwise_affine", False)
|
||||
check_config_value(config, "qk_norm", "rms_norm")
|
||||
check_config_value(config, "positional_embedding_type", "rope")
|
||||
check_config_value(config, "use_audio_video_cross_attention", True)
|
||||
check_config_value(config, "share_ff", False)
|
||||
check_config_value(config, "av_cross_ada_norm", True)
|
||||
check_config_value(config, "use_middle_indices_grid", True)
|
||||
|
||||
return LTXModel(
|
||||
model_type=LTXModelType.AudioVideo,
|
||||
num_attention_heads=config.get("num_attention_heads", 32),
|
||||
attention_head_dim=config.get("attention_head_dim", 128),
|
||||
in_channels=config.get("in_channels", 128),
|
||||
out_channels=config.get("out_channels", 128),
|
||||
num_layers=config.get("num_layers", 48),
|
||||
cross_attention_dim=config.get("cross_attention_dim", 4096),
|
||||
norm_eps=config.get("norm_eps", 1e-06),
|
||||
attention_type=AttentionFunction(config.get("attention_type", "default")),
|
||||
caption_channels=config.get("caption_channels", 3840),
|
||||
positional_embedding_theta=config.get("positional_embedding_theta", 10000.0),
|
||||
positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]),
|
||||
timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000),
|
||||
use_middle_indices_grid=config.get("use_middle_indices_grid", True),
|
||||
audio_num_attention_heads=config.get("audio_num_attention_heads", 32),
|
||||
audio_attention_head_dim=config.get("audio_attention_head_dim", 64),
|
||||
audio_in_channels=config.get("audio_in_channels", 128),
|
||||
audio_out_channels=config.get("audio_out_channels", 128),
|
||||
audio_cross_attention_dim=config.get("audio_cross_attention_dim", 2048),
|
||||
audio_positional_embedding_max_pos=config.get("audio_positional_embedding_max_pos", [20]),
|
||||
av_ca_timestep_scale_multiplier=config.get("av_ca_timestep_scale_multiplier", 1),
|
||||
rope_type=LTXRopeType(config.get("rope_type", "interleaved")),
|
||||
double_precision_rope=config.get("frequencies_precision", False) == "float64",
|
||||
)
|
||||
|
||||
|
||||
class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
|
||||
"""
|
||||
Configurator for LTX video only model.
|
||||
Used to create an LTX video only model from a configuration dictionary.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def from_config(cls: type[LTXModel], config: dict) -> LTXModel:
|
||||
config = config.get("transformer", {})
|
||||
|
||||
check_config_value(config, "dropout", 0.0)
|
||||
check_config_value(config, "attention_bias", True)
|
||||
check_config_value(config, "num_vector_embeds", None)
|
||||
check_config_value(config, "activation_fn", "gelu-approximate")
|
||||
check_config_value(config, "num_embeds_ada_norm", 1000)
|
||||
check_config_value(config, "use_linear_projection", False)
|
||||
check_config_value(config, "only_cross_attention", False)
|
||||
check_config_value(config, "cross_attention_norm", True)
|
||||
check_config_value(config, "double_self_attention", False)
|
||||
check_config_value(config, "upcast_attention", False)
|
||||
check_config_value(config, "standardization_norm", "rms_norm")
|
||||
check_config_value(config, "norm_elementwise_affine", False)
|
||||
check_config_value(config, "qk_norm", "rms_norm")
|
||||
check_config_value(config, "positional_embedding_type", "rope")
|
||||
check_config_value(config, "use_middle_indices_grid", True)
|
||||
|
||||
return LTXModel(
|
||||
model_type=LTXModelType.VideoOnly,
|
||||
num_attention_heads=config.get("num_attention_heads", 32),
|
||||
attention_head_dim=config.get("attention_head_dim", 128),
|
||||
in_channels=config.get("in_channels", 128),
|
||||
out_channels=config.get("out_channels", 128),
|
||||
num_layers=config.get("num_layers", 48),
|
||||
cross_attention_dim=config.get("cross_attention_dim", 4096),
|
||||
norm_eps=config.get("norm_eps", 1e-06),
|
||||
attention_type=AttentionFunction(config.get("attention_type", "default")),
|
||||
caption_channels=config.get("caption_channels", 3840),
|
||||
positional_embedding_theta=config.get("positional_embedding_theta", 10000.0),
|
||||
positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]),
|
||||
timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000),
|
||||
use_middle_indices_grid=config.get("use_middle_indices_grid", True),
|
||||
rope_type=LTXRopeType(config.get("rope_type", "interleaved")),
|
||||
double_precision_rope=config.get("frequencies_precision", False) == "float64",
|
||||
)
|
||||
|
||||
|
||||
def _naive_weight_or_bias_downcast(key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
|
||||
"""
|
||||
Downcast the weight or bias to the float8_e4m3fn dtype.
|
||||
"""
|
||||
return [KeyValueOperationResult(key, value.to(dtype=torch.float8_e4m3fn))]
|
||||
|
||||
|
||||
def _upcast_and_round(
|
||||
weight: torch.Tensor, dtype: torch.dtype, with_stochastic_rounding: bool = False, seed: int = 0
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Upcast the weight to the given dtype and optionally apply stochastic rounding.
|
||||
Input weight needs to have float8_e4m3fn or float8_e5m2 dtype.
|
||||
"""
|
||||
if not with_stochastic_rounding:
|
||||
return weight.to(dtype)
|
||||
return fused_add_round_launch(torch.zeros_like(weight, dtype=dtype), weight, seed)
|
||||
|
||||
|
||||
def replace_fwd_with_upcast(layer: torch.nn.Linear, with_stochastic_rounding: bool = False, seed: int = 0) -> None:
|
||||
"""
|
||||
Replace linear.forward and rms_norm.forward with a version that:
|
||||
- upcasts weight and bias to input's dtype
|
||||
- returns F.linear or F.rms_norm calculated in that dtype
|
||||
"""
|
||||
|
||||
layer.original_forward = layer.forward
|
||||
|
||||
def new_linear_forward(*args, **_kwargs) -> torch.Tensor:
|
||||
# assume first arg is the input tensor
|
||||
x = args[0]
|
||||
w_up = _upcast_and_round(layer.weight, x.dtype, with_stochastic_rounding, seed)
|
||||
b_up = None
|
||||
|
||||
if layer.bias is not None:
|
||||
b_up = _upcast_and_round(layer.bias, x.dtype, with_stochastic_rounding, seed)
|
||||
|
||||
return torch.nn.functional.linear(x, w_up, b_up)
|
||||
|
||||
layer.forward = new_linear_forward
|
||||
|
||||
|
||||
def amend_forward_with_upcast(
|
||||
model: torch.nn.Module, with_stochastic_rounding: bool = False, seed: int = 0
|
||||
) -> torch.nn.Module:
|
||||
"""
|
||||
Replace the forward method of the model's Linear and RMSNorm layers to forward
|
||||
with upcast and optional stochastic rounding.
|
||||
"""
|
||||
for m in model.modules():
|
||||
if isinstance(m, (torch.nn.Linear)):
|
||||
replace_fwd_with_upcast(m, with_stochastic_rounding, seed)
|
||||
return model
|
||||
|
||||
|
||||
LTXV_MODEL_COMFY_RENAMING_MAP = (
|
||||
SDOps("LTXV_MODEL_COMFY_PREFIX_MAP")
|
||||
.with_matching(prefix="model.diffusion_model.")
|
||||
.with_replacement("model.diffusion_model.", "")
|
||||
)
|
||||
|
||||
LTXV_MODEL_COMFY_RENAMING_WITH_TRANSFORMER_LINEAR_DOWNCAST_MAP = (
|
||||
SDOps("LTXV_MODEL_COMFY_PREFIX_MAP")
|
||||
.with_matching(prefix="model.diffusion_model.")
|
||||
.with_replacement("model.diffusion_model.", "")
|
||||
.with_kv_operation(
|
||||
key_prefix="transformer_blocks.", key_suffix=".to_q.weight", operation=_naive_weight_or_bias_downcast
|
||||
)
|
||||
.with_kv_operation(
|
||||
key_prefix="transformer_blocks.", key_suffix=".to_q.bias", operation=_naive_weight_or_bias_downcast
|
||||
)
|
||||
.with_kv_operation(
|
||||
key_prefix="transformer_blocks.", key_suffix=".to_k.weight", operation=_naive_weight_or_bias_downcast
|
||||
)
|
||||
.with_kv_operation(
|
||||
key_prefix="transformer_blocks.", key_suffix=".to_k.bias", operation=_naive_weight_or_bias_downcast
|
||||
)
|
||||
.with_kv_operation(
|
||||
key_prefix="transformer_blocks.", key_suffix=".to_v.weight", operation=_naive_weight_or_bias_downcast
|
||||
)
|
||||
.with_kv_operation(
|
||||
key_prefix="transformer_blocks.", key_suffix=".to_v.bias", operation=_naive_weight_or_bias_downcast
|
||||
)
|
||||
.with_kv_operation(
|
||||
key_prefix="transformer_blocks.", key_suffix=".to_out.0.weight", operation=_naive_weight_or_bias_downcast
|
||||
)
|
||||
.with_kv_operation(
|
||||
key_prefix="transformer_blocks.", key_suffix=".to_out.0.bias", operation=_naive_weight_or_bias_downcast
|
||||
)
|
||||
.with_kv_operation(
|
||||
key_prefix="transformer_blocks.", key_suffix=".ff.net.0.proj.weight", operation=_naive_weight_or_bias_downcast
|
||||
)
|
||||
.with_kv_operation(
|
||||
key_prefix="transformer_blocks.", key_suffix=".ff.net.0.proj.bias", operation=_naive_weight_or_bias_downcast
|
||||
)
|
||||
.with_kv_operation(
|
||||
key_prefix="transformer_blocks.", key_suffix=".ff.net.2.weight", operation=_naive_weight_or_bias_downcast
|
||||
)
|
||||
.with_kv_operation(
|
||||
key_prefix="transformer_blocks.", key_suffix=".ff.net.2.bias", operation=_naive_weight_or_bias_downcast
|
||||
)
|
||||
)
|
||||
|
||||
UPCAST_DURING_INFERENCE = ModuleOps(
|
||||
name="upcast_fp8_during_linear_forward",
|
||||
matcher=lambda model: isinstance(model, LTXModel),
|
||||
mutator=lambda model: amend_forward_with_upcast(model, False),
|
||||
)
|
||||
|
||||
|
||||
class UpcastWithStochasticRounding(ModuleOps):
|
||||
"""
|
||||
ModuleOps for upcasting the model's float8_e4m3fn weights and biases to the bfloat16 dtype
|
||||
and applying stochastic rounding during linear forward.
|
||||
"""
|
||||
|
||||
def __new__(cls, seed: int = 0):
|
||||
return super().__new__(
|
||||
cls,
|
||||
name="upcast_fp8_during_linear_forward_with_stochastic_rounding",
|
||||
matcher=lambda model: isinstance(model, LTXModel),
|
||||
mutator=lambda model: amend_forward_with_upcast(model, True, seed),
|
||||
)
|
||||
@@ -0,0 +1,204 @@
|
||||
import functools
|
||||
import math
|
||||
from enum import Enum
|
||||
from typing import Callable, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from einops import rearrange
|
||||
|
||||
|
||||
class LTXRopeType(Enum):
|
||||
INTERLEAVED = "interleaved"
|
||||
SPLIT = "split"
|
||||
|
||||
|
||||
def apply_rotary_emb(
|
||||
input_tensor: torch.Tensor,
|
||||
freqs_cis: Tuple[torch.Tensor, torch.Tensor],
|
||||
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
|
||||
) -> torch.Tensor:
|
||||
if rope_type == LTXRopeType.INTERLEAVED:
|
||||
return apply_interleaved_rotary_emb(input_tensor, *freqs_cis)
|
||||
elif rope_type == LTXRopeType.SPLIT:
|
||||
return apply_split_rotary_emb(input_tensor, *freqs_cis)
|
||||
else:
|
||||
raise ValueError(f"Invalid rope type: {rope_type}")
|
||||
|
||||
|
||||
def apply_interleaved_rotary_emb(
|
||||
input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
t_dup = rearrange(input_tensor, "... (d r) -> ... d r", r=2)
|
||||
t1, t2 = t_dup.unbind(dim=-1)
|
||||
t_dup = torch.stack((-t2, t1), dim=-1)
|
||||
input_tensor_rot = rearrange(t_dup, "... d r -> ... (d r)")
|
||||
|
||||
out = input_tensor * cos_freqs + input_tensor_rot * sin_freqs
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def apply_split_rotary_emb(
|
||||
input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
needs_reshape = False
|
||||
if input_tensor.ndim != 4 and cos_freqs.ndim == 4:
|
||||
b, h, t, _ = cos_freqs.shape
|
||||
input_tensor = input_tensor.reshape(b, t, h, -1).swapaxes(1, 2)
|
||||
needs_reshape = True
|
||||
|
||||
split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2)
|
||||
first_half_input = split_input[..., :1, :]
|
||||
second_half_input = split_input[..., 1:, :]
|
||||
|
||||
output = split_input * cos_freqs.unsqueeze(-2)
|
||||
first_half_output = output[..., :1, :]
|
||||
second_half_output = output[..., 1:, :]
|
||||
|
||||
first_half_output.addcmul_(-sin_freqs.unsqueeze(-2), second_half_input)
|
||||
second_half_output.addcmul_(sin_freqs.unsqueeze(-2), first_half_input)
|
||||
|
||||
output = rearrange(output, "... d r -> ... (d r)")
|
||||
if needs_reshape:
|
||||
output = output.swapaxes(1, 2).reshape(b, t, -1)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=5)
|
||||
def generate_freq_grid_np(
|
||||
positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int
|
||||
) -> torch.Tensor:
|
||||
theta = positional_embedding_theta
|
||||
start = 1
|
||||
end = theta
|
||||
|
||||
n_elem = 2 * positional_embedding_max_pos_count
|
||||
pow_indices = np.power(
|
||||
theta,
|
||||
np.linspace(
|
||||
np.log(start) / np.log(theta),
|
||||
np.log(end) / np.log(theta),
|
||||
inner_dim // n_elem,
|
||||
dtype=np.float64,
|
||||
),
|
||||
)
|
||||
return torch.tensor(pow_indices * math.pi / 2, dtype=torch.float32)
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=5)
|
||||
def generate_freq_grid_pytorch(
|
||||
positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int
|
||||
) -> torch.Tensor:
|
||||
theta = positional_embedding_theta
|
||||
start = 1
|
||||
end = theta
|
||||
n_elem = 2 * positional_embedding_max_pos_count
|
||||
|
||||
indices = theta ** (
|
||||
torch.linspace(
|
||||
math.log(start, theta),
|
||||
math.log(end, theta),
|
||||
inner_dim // n_elem,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
)
|
||||
indices = indices.to(dtype=torch.float32)
|
||||
|
||||
indices = indices * math.pi / 2
|
||||
|
||||
return indices
|
||||
|
||||
|
||||
def get_fractional_positions(indices_grid: torch.Tensor, max_pos: list[int]) -> torch.Tensor:
|
||||
n_pos_dims = indices_grid.shape[1]
|
||||
assert n_pos_dims == len(max_pos), (
|
||||
f"Number of position dimensions ({n_pos_dims}) must match max_pos length ({len(max_pos)})"
|
||||
)
|
||||
fractional_positions = torch.stack(
|
||||
[indices_grid[:, i] / max_pos[i] for i in range(n_pos_dims)],
|
||||
dim=-1,
|
||||
)
|
||||
return fractional_positions
|
||||
|
||||
|
||||
def generate_freqs(
|
||||
indices: torch.Tensor, indices_grid: torch.Tensor, max_pos: list[int], use_middle_indices_grid: bool
|
||||
) -> torch.Tensor:
|
||||
if use_middle_indices_grid:
|
||||
assert len(indices_grid.shape) == 4
|
||||
assert indices_grid.shape[-1] == 2
|
||||
indices_grid_start, indices_grid_end = indices_grid[..., 0], indices_grid[..., 1]
|
||||
indices_grid = (indices_grid_start + indices_grid_end) / 2.0
|
||||
elif len(indices_grid.shape) == 4:
|
||||
indices_grid = indices_grid[..., 0]
|
||||
|
||||
fractional_positions = get_fractional_positions(indices_grid, max_pos)
|
||||
indices = indices.to(device=fractional_positions.device)
|
||||
|
||||
freqs = (indices * (fractional_positions.unsqueeze(-1) * 2 - 1)).transpose(-1, -2).flatten(2)
|
||||
return freqs
|
||||
|
||||
|
||||
def split_freqs_cis(freqs: torch.Tensor, pad_size: int, num_attention_heads: int) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
cos_freq = freqs.cos()
|
||||
sin_freq = freqs.sin()
|
||||
|
||||
if pad_size != 0:
|
||||
cos_padding = torch.ones_like(cos_freq[:, :, :pad_size])
|
||||
sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size])
|
||||
|
||||
cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1)
|
||||
sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1)
|
||||
|
||||
# Reshape freqs to be compatible with multi-head attention
|
||||
b = cos_freq.shape[0]
|
||||
t = cos_freq.shape[1]
|
||||
|
||||
cos_freq = cos_freq.reshape(b, t, num_attention_heads, -1)
|
||||
sin_freq = sin_freq.reshape(b, t, num_attention_heads, -1)
|
||||
|
||||
cos_freq = torch.swapaxes(cos_freq, 1, 2) # (B,H,T,D//2)
|
||||
sin_freq = torch.swapaxes(sin_freq, 1, 2) # (B,H,T,D//2)
|
||||
return cos_freq, sin_freq
|
||||
|
||||
|
||||
def interleaved_freqs_cis(freqs: torch.Tensor, pad_size: int) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
cos_freq = freqs.cos().repeat_interleave(2, dim=-1)
|
||||
sin_freq = freqs.sin().repeat_interleave(2, dim=-1)
|
||||
if pad_size != 0:
|
||||
cos_padding = torch.ones_like(cos_freq[:, :, :pad_size])
|
||||
sin_padding = torch.zeros_like(cos_freq[:, :, :pad_size])
|
||||
cos_freq = torch.cat([cos_padding, cos_freq], dim=-1)
|
||||
sin_freq = torch.cat([sin_padding, sin_freq], dim=-1)
|
||||
return cos_freq, sin_freq
|
||||
|
||||
|
||||
def precompute_freqs_cis(
|
||||
indices_grid: torch.Tensor,
|
||||
dim: int,
|
||||
out_dtype: torch.dtype,
|
||||
theta: float = 10000.0,
|
||||
max_pos: list[int] | None = None,
|
||||
use_middle_indices_grid: bool = False,
|
||||
num_attention_heads: int = 32,
|
||||
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
|
||||
freq_grid_generator: Callable[[float, int, int, torch.device], torch.Tensor] = generate_freq_grid_pytorch,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if max_pos is None:
|
||||
max_pos = [20, 2048, 2048]
|
||||
|
||||
indices = freq_grid_generator(theta, indices_grid.shape[1], dim)
|
||||
freqs = generate_freqs(indices, indices_grid, max_pos, use_middle_indices_grid)
|
||||
|
||||
if rope_type == LTXRopeType.SPLIT:
|
||||
expected_freqs = dim // 2
|
||||
current_freqs = freqs.shape[-1]
|
||||
pad_size = expected_freqs - current_freqs
|
||||
cos_freq, sin_freq = split_freqs_cis(freqs, pad_size, num_attention_heads)
|
||||
else:
|
||||
# 2 because of cos and sin by 3 for (t, x, y), 1 for temporal only
|
||||
n_elem = 2 * indices_grid.shape[1]
|
||||
cos_freq, sin_freq = interleaved_freqs_cis(freqs, dim % n_elem)
|
||||
return cos_freq.to(out_dtype), sin_freq.to(out_dtype)
|
||||
@@ -0,0 +1,27 @@
|
||||
import torch
|
||||
|
||||
|
||||
class PixArtAlphaTextProjection(torch.nn.Module):
|
||||
"""
|
||||
Projects caption embeddings. Also handles dropout for classifier-free guidance.
|
||||
Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py
|
||||
"""
|
||||
|
||||
def __init__(self, in_features: int, hidden_size: int, out_features: int | None = None, act_fn: str = "gelu_tanh"):
|
||||
super().__init__()
|
||||
if out_features is None:
|
||||
out_features = hidden_size
|
||||
self.linear_1 = torch.nn.Linear(in_features=in_features, out_features=hidden_size, bias=True)
|
||||
if act_fn == "gelu_tanh":
|
||||
self.act_1 = torch.nn.GELU(approximate="tanh")
|
||||
elif act_fn == "silu":
|
||||
self.act_1 = torch.nn.SiLU()
|
||||
else:
|
||||
raise ValueError(f"Unknown activation function: {act_fn}")
|
||||
self.linear_2 = torch.nn.Linear(in_features=hidden_size, out_features=out_features, bias=True)
|
||||
|
||||
def forward(self, caption: torch.Tensor) -> torch.Tensor:
|
||||
hidden_states = self.linear_1(caption)
|
||||
hidden_states = self.act_1(hidden_states)
|
||||
hidden_states = self.linear_2(hidden_states)
|
||||
return hidden_states
|
||||
@@ -0,0 +1,143 @@
|
||||
import math
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def get_timestep_embedding(
|
||||
timesteps: torch.Tensor,
|
||||
embedding_dim: int,
|
||||
flip_sin_to_cos: bool = False,
|
||||
downscale_freq_shift: float = 1,
|
||||
scale: float = 1,
|
||||
max_period: int = 10000,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
|
||||
Args
|
||||
timesteps (torch.Tensor):
|
||||
a 1-D Tensor of N indices, one per batch element. These may be fractional.
|
||||
embedding_dim (int):
|
||||
the dimension of the output.
|
||||
flip_sin_to_cos (bool):
|
||||
Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False)
|
||||
downscale_freq_shift (float):
|
||||
Controls the delta between frequencies between dimensions
|
||||
scale (float):
|
||||
Scaling factor applied to the embeddings.
|
||||
max_period (int):
|
||||
Controls the maximum frequency of the embeddings
|
||||
Returns
|
||||
torch.Tensor: an [N x dim] Tensor of positional embeddings.
|
||||
"""
|
||||
assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
|
||||
|
||||
half_dim = embedding_dim // 2
|
||||
exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device)
|
||||
exponent = exponent / (half_dim - downscale_freq_shift)
|
||||
|
||||
emb = torch.exp(exponent)
|
||||
emb = timesteps[:, None].float() * emb[None, :]
|
||||
|
||||
# scale embeddings
|
||||
emb = scale * emb
|
||||
|
||||
# concat sine and cosine embeddings
|
||||
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
|
||||
|
||||
# flip sine and cosine embeddings
|
||||
if flip_sin_to_cos:
|
||||
emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
|
||||
|
||||
# zero pad
|
||||
if embedding_dim % 2 == 1:
|
||||
emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
|
||||
return emb
|
||||
|
||||
|
||||
class TimestepEmbedding(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
time_embed_dim: int,
|
||||
out_dim: int | None = None,
|
||||
post_act_fn: str | None = None,
|
||||
cond_proj_dim: int | None = None,
|
||||
sample_proj_bias: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.linear_1 = torch.nn.Linear(in_channels, time_embed_dim, sample_proj_bias)
|
||||
|
||||
if cond_proj_dim is not None:
|
||||
self.cond_proj = torch.nn.Linear(cond_proj_dim, in_channels, bias=False)
|
||||
else:
|
||||
self.cond_proj = None
|
||||
|
||||
self.act = torch.nn.SiLU()
|
||||
time_embed_dim_out = out_dim if out_dim is not None else time_embed_dim
|
||||
|
||||
self.linear_2 = torch.nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias)
|
||||
|
||||
if post_act_fn is None:
|
||||
self.post_act = None
|
||||
|
||||
def forward(self, sample: torch.Tensor, condition: torch.Tensor | None = None) -> torch.Tensor:
|
||||
if condition is not None:
|
||||
sample = sample + self.cond_proj(condition)
|
||||
sample = self.linear_1(sample)
|
||||
|
||||
if self.act is not None:
|
||||
sample = self.act(sample)
|
||||
|
||||
sample = self.linear_2(sample)
|
||||
|
||||
if self.post_act is not None:
|
||||
sample = self.post_act(sample)
|
||||
return sample
|
||||
|
||||
|
||||
class Timesteps(torch.nn.Module):
|
||||
def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1):
|
||||
super().__init__()
|
||||
self.num_channels = num_channels
|
||||
self.flip_sin_to_cos = flip_sin_to_cos
|
||||
self.downscale_freq_shift = downscale_freq_shift
|
||||
self.scale = scale
|
||||
|
||||
def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
|
||||
t_emb = get_timestep_embedding(
|
||||
timesteps,
|
||||
self.num_channels,
|
||||
flip_sin_to_cos=self.flip_sin_to_cos,
|
||||
downscale_freq_shift=self.downscale_freq_shift,
|
||||
scale=self.scale,
|
||||
)
|
||||
return t_emb
|
||||
|
||||
|
||||
class PixArtAlphaCombinedTimestepSizeEmbeddings(torch.nn.Module):
|
||||
"""
|
||||
For PixArt-Alpha.
|
||||
Reference:
|
||||
https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L164C9-L168C29
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
embedding_dim: int,
|
||||
size_emb_dim: int,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.outdim = size_emb_dim
|
||||
self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
|
||||
self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
timestep: torch.Tensor,
|
||||
hidden_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
timesteps_proj = self.time_proj(timestep)
|
||||
timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D)
|
||||
return timesteps_emb
|
||||
@@ -0,0 +1,274 @@
|
||||
from dataclasses import dataclass, replace
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType
|
||||
from ltx_core.model.transformer.attention import Attention, AttentionCallable, AttentionFunction
|
||||
from ltx_core.model.transformer.feed_forward import FeedForward
|
||||
from ltx_core.model.transformer.rope import LTXRopeType
|
||||
from ltx_core.model.transformer.transformer_args import TransformerArgs
|
||||
from ltx_core.utils import rms_norm
|
||||
|
||||
|
||||
@dataclass
|
||||
class TransformerConfig:
|
||||
dim: int
|
||||
heads: int
|
||||
d_head: int
|
||||
context_dim: int
|
||||
|
||||
|
||||
class BasicAVTransformerBlock(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
idx: int,
|
||||
video: TransformerConfig | None = None,
|
||||
audio: TransformerConfig | None = None,
|
||||
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
|
||||
norm_eps: float = 1e-6,
|
||||
attention_function: AttentionFunction | AttentionCallable = AttentionFunction.DEFAULT,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.idx = idx
|
||||
if video is not None:
|
||||
self.attn1 = Attention(
|
||||
query_dim=video.dim,
|
||||
heads=video.heads,
|
||||
dim_head=video.d_head,
|
||||
context_dim=None,
|
||||
rope_type=rope_type,
|
||||
norm_eps=norm_eps,
|
||||
attention_function=attention_function,
|
||||
)
|
||||
self.attn2 = Attention(
|
||||
query_dim=video.dim,
|
||||
context_dim=video.context_dim,
|
||||
heads=video.heads,
|
||||
dim_head=video.d_head,
|
||||
rope_type=rope_type,
|
||||
norm_eps=norm_eps,
|
||||
attention_function=attention_function,
|
||||
)
|
||||
self.ff = FeedForward(video.dim, dim_out=video.dim)
|
||||
self.scale_shift_table = torch.nn.Parameter(torch.empty(6, video.dim))
|
||||
|
||||
if audio is not None:
|
||||
self.audio_attn1 = Attention(
|
||||
query_dim=audio.dim,
|
||||
heads=audio.heads,
|
||||
dim_head=audio.d_head,
|
||||
context_dim=None,
|
||||
rope_type=rope_type,
|
||||
norm_eps=norm_eps,
|
||||
attention_function=attention_function,
|
||||
)
|
||||
self.audio_attn2 = Attention(
|
||||
query_dim=audio.dim,
|
||||
context_dim=audio.context_dim,
|
||||
heads=audio.heads,
|
||||
dim_head=audio.d_head,
|
||||
rope_type=rope_type,
|
||||
norm_eps=norm_eps,
|
||||
attention_function=attention_function,
|
||||
)
|
||||
self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim)
|
||||
self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(6, audio.dim))
|
||||
|
||||
if audio is not None and video is not None:
|
||||
# Q: Video, K,V: Audio
|
||||
self.audio_to_video_attn = Attention(
|
||||
query_dim=video.dim,
|
||||
context_dim=audio.dim,
|
||||
heads=audio.heads,
|
||||
dim_head=audio.d_head,
|
||||
rope_type=rope_type,
|
||||
norm_eps=norm_eps,
|
||||
attention_function=attention_function,
|
||||
)
|
||||
|
||||
# Q: Audio, K,V: Video
|
||||
self.video_to_audio_attn = Attention(
|
||||
query_dim=audio.dim,
|
||||
context_dim=video.dim,
|
||||
heads=audio.heads,
|
||||
dim_head=audio.d_head,
|
||||
rope_type=rope_type,
|
||||
norm_eps=norm_eps,
|
||||
attention_function=attention_function,
|
||||
)
|
||||
|
||||
self.scale_shift_table_a2v_ca_audio = torch.nn.Parameter(torch.empty(5, audio.dim))
|
||||
self.scale_shift_table_a2v_ca_video = torch.nn.Parameter(torch.empty(5, video.dim))
|
||||
|
||||
self.norm_eps = norm_eps
|
||||
|
||||
def get_ada_values(
|
||||
self, scale_shift_table: torch.Tensor, batch_size: int, timestep: torch.Tensor, indices: slice
|
||||
) -> tuple[torch.Tensor, ...]:
|
||||
num_ada_params = scale_shift_table.shape[0]
|
||||
|
||||
ada_values = (
|
||||
scale_shift_table[indices].unsqueeze(0).unsqueeze(0).to(device=timestep.device, dtype=timestep.dtype)
|
||||
+ timestep.reshape(batch_size, timestep.shape[1], num_ada_params, -1)[:, :, indices, :]
|
||||
).unbind(dim=2)
|
||||
return ada_values
|
||||
|
||||
def get_av_ca_ada_values(
|
||||
self,
|
||||
scale_shift_table: torch.Tensor,
|
||||
batch_size: int,
|
||||
scale_shift_timestep: torch.Tensor,
|
||||
gate_timestep: torch.Tensor,
|
||||
num_scale_shift_values: int = 4,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
scale_shift_ada_values = self.get_ada_values(
|
||||
scale_shift_table[:num_scale_shift_values, :], batch_size, scale_shift_timestep, slice(None, None)
|
||||
)
|
||||
gate_ada_values = self.get_ada_values(
|
||||
scale_shift_table[num_scale_shift_values:, :], batch_size, gate_timestep, slice(None, None)
|
||||
)
|
||||
|
||||
scale_shift_chunks = [t.squeeze(2) for t in scale_shift_ada_values]
|
||||
gate_ada_values = [t.squeeze(2) for t in gate_ada_values]
|
||||
|
||||
return (*scale_shift_chunks, *gate_ada_values)
|
||||
|
||||
def forward( # noqa: PLR0915
|
||||
self,
|
||||
video: TransformerArgs | None,
|
||||
audio: TransformerArgs | None,
|
||||
perturbations: BatchedPerturbationConfig | None = None,
|
||||
) -> tuple[TransformerArgs | None, TransformerArgs | None]:
|
||||
batch_size = video.x.shape[0]
|
||||
if perturbations is None:
|
||||
perturbations = BatchedPerturbationConfig.empty(batch_size)
|
||||
|
||||
vx = video.x if video is not None else None
|
||||
ax = audio.x if audio is not None else None
|
||||
|
||||
run_vx = video is not None and video.enabled and vx.numel() > 0
|
||||
run_ax = audio is not None and audio.enabled and ax.numel() > 0
|
||||
|
||||
run_a2v = run_vx and (audio is not None and audio.enabled and ax.numel() > 0)
|
||||
run_v2a = run_ax and (video is not None and video.enabled and vx.numel() > 0)
|
||||
|
||||
if run_vx:
|
||||
vshift_msa, vscale_msa, vgate_msa = self.get_ada_values(
|
||||
self.scale_shift_table, vx.shape[0], video.timesteps, slice(0, 3)
|
||||
)
|
||||
if not perturbations.all_in_batch(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx):
|
||||
norm_vx = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa
|
||||
v_mask = perturbations.mask_like(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx, vx)
|
||||
vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings) * vgate_msa * v_mask
|
||||
|
||||
vx = vx + self.attn2(rms_norm(vx, eps=self.norm_eps), context=video.context, mask=video.context_mask)
|
||||
|
||||
del vshift_msa, vscale_msa, vgate_msa
|
||||
|
||||
if run_ax:
|
||||
ashift_msa, ascale_msa, agate_msa = self.get_ada_values(
|
||||
self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(0, 3)
|
||||
)
|
||||
|
||||
if not perturbations.all_in_batch(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx):
|
||||
norm_ax = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa
|
||||
a_mask = perturbations.mask_like(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx, ax)
|
||||
ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings) * agate_msa * a_mask
|
||||
|
||||
ax = ax + self.audio_attn2(rms_norm(ax, eps=self.norm_eps), context=audio.context, mask=audio.context_mask)
|
||||
|
||||
del ashift_msa, ascale_msa, agate_msa
|
||||
|
||||
# Audio - Video cross attention.
|
||||
if run_a2v or run_v2a:
|
||||
vx_norm3 = rms_norm(vx, eps=self.norm_eps)
|
||||
ax_norm3 = rms_norm(ax, eps=self.norm_eps)
|
||||
|
||||
(
|
||||
scale_ca_audio_hidden_states_a2v,
|
||||
shift_ca_audio_hidden_states_a2v,
|
||||
scale_ca_audio_hidden_states_v2a,
|
||||
shift_ca_audio_hidden_states_v2a,
|
||||
gate_out_v2a,
|
||||
) = self.get_av_ca_ada_values(
|
||||
self.scale_shift_table_a2v_ca_audio,
|
||||
ax.shape[0],
|
||||
audio.cross_scale_shift_timestep,
|
||||
audio.cross_gate_timestep,
|
||||
)
|
||||
|
||||
(
|
||||
scale_ca_video_hidden_states_a2v,
|
||||
shift_ca_video_hidden_states_a2v,
|
||||
scale_ca_video_hidden_states_v2a,
|
||||
shift_ca_video_hidden_states_v2a,
|
||||
gate_out_a2v,
|
||||
) = self.get_av_ca_ada_values(
|
||||
self.scale_shift_table_a2v_ca_video,
|
||||
vx.shape[0],
|
||||
video.cross_scale_shift_timestep,
|
||||
video.cross_gate_timestep,
|
||||
)
|
||||
|
||||
if run_a2v:
|
||||
vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_a2v) + shift_ca_video_hidden_states_a2v
|
||||
ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_a2v) + shift_ca_audio_hidden_states_a2v
|
||||
a2v_mask = perturbations.mask_like(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx, vx)
|
||||
vx = vx + (
|
||||
self.audio_to_video_attn(
|
||||
vx_scaled,
|
||||
context=ax_scaled,
|
||||
pe=video.cross_positional_embeddings,
|
||||
k_pe=audio.cross_positional_embeddings,
|
||||
)
|
||||
* gate_out_a2v
|
||||
* a2v_mask
|
||||
)
|
||||
|
||||
if run_v2a:
|
||||
ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_v2a) + shift_ca_audio_hidden_states_v2a
|
||||
vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_v2a) + shift_ca_video_hidden_states_v2a
|
||||
v2a_mask = perturbations.mask_like(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx, ax)
|
||||
ax = ax + (
|
||||
self.video_to_audio_attn(
|
||||
ax_scaled,
|
||||
context=vx_scaled,
|
||||
pe=audio.cross_positional_embeddings,
|
||||
k_pe=video.cross_positional_embeddings,
|
||||
)
|
||||
* gate_out_v2a
|
||||
* v2a_mask
|
||||
)
|
||||
|
||||
del gate_out_a2v, gate_out_v2a
|
||||
del (
|
||||
scale_ca_video_hidden_states_a2v,
|
||||
shift_ca_video_hidden_states_a2v,
|
||||
scale_ca_audio_hidden_states_a2v,
|
||||
shift_ca_audio_hidden_states_a2v,
|
||||
scale_ca_video_hidden_states_v2a,
|
||||
shift_ca_video_hidden_states_v2a,
|
||||
scale_ca_audio_hidden_states_v2a,
|
||||
shift_ca_audio_hidden_states_v2a,
|
||||
)
|
||||
|
||||
if run_vx:
|
||||
vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values(
|
||||
self.scale_shift_table, vx.shape[0], video.timesteps, slice(3, None)
|
||||
)
|
||||
vx_scaled = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_mlp) + vshift_mlp
|
||||
vx = vx + self.ff(vx_scaled) * vgate_mlp
|
||||
|
||||
del vshift_mlp, vscale_mlp, vgate_mlp
|
||||
|
||||
if run_ax:
|
||||
ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values(
|
||||
self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(3, None)
|
||||
)
|
||||
ax_scaled = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_mlp) + ashift_mlp
|
||||
ax = ax + self.audio_ff(ax_scaled) * agate_mlp
|
||||
|
||||
del ashift_mlp, ascale_mlp, agate_mlp
|
||||
|
||||
return replace(video, x=vx) if video is not None else None, replace(audio, x=ax) if audio is not None else None
|
||||
@@ -0,0 +1,239 @@
|
||||
from dataclasses import dataclass, replace
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.model.transformer.adaln import AdaLayerNormSingle
|
||||
from ltx_core.model.transformer.modality import Modality
|
||||
from ltx_core.model.transformer.rope import (
|
||||
LTXRopeType,
|
||||
generate_freq_grid_np,
|
||||
generate_freq_grid_pytorch,
|
||||
precompute_freqs_cis,
|
||||
)
|
||||
from ltx_core.model.transformer.text_projection import PixArtAlphaTextProjection
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TransformerArgs:
|
||||
x: torch.Tensor
|
||||
context: torch.Tensor
|
||||
context_mask: torch.Tensor
|
||||
timesteps: torch.Tensor
|
||||
embedded_timestep: torch.Tensor
|
||||
positional_embeddings: torch.Tensor
|
||||
cross_positional_embeddings: torch.Tensor | None
|
||||
cross_scale_shift_timestep: torch.Tensor | None
|
||||
cross_gate_timestep: torch.Tensor | None
|
||||
enabled: bool
|
||||
|
||||
|
||||
class TransformerArgsPreprocessor:
|
||||
def __init__( # noqa: PLR0913
|
||||
self,
|
||||
patchify_proj: torch.nn.Linear,
|
||||
adaln: AdaLayerNormSingle,
|
||||
caption_projection: PixArtAlphaTextProjection,
|
||||
inner_dim: int,
|
||||
max_pos: list[int],
|
||||
num_attention_heads: int,
|
||||
use_middle_indices_grid: bool,
|
||||
timestep_scale_multiplier: int,
|
||||
double_precision_rope: bool,
|
||||
positional_embedding_theta: float,
|
||||
rope_type: LTXRopeType,
|
||||
) -> None:
|
||||
self.patchify_proj = patchify_proj
|
||||
self.adaln = adaln
|
||||
self.caption_projection = caption_projection
|
||||
self.inner_dim = inner_dim
|
||||
self.max_pos = max_pos
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.use_middle_indices_grid = use_middle_indices_grid
|
||||
self.timestep_scale_multiplier = timestep_scale_multiplier
|
||||
self.double_precision_rope = double_precision_rope
|
||||
self.positional_embedding_theta = positional_embedding_theta
|
||||
self.rope_type = rope_type
|
||||
|
||||
def _prepare_timestep(
|
||||
self, timestep: torch.Tensor, batch_size: int, hidden_dtype: torch.dtype
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Prepare timestep embeddings."""
|
||||
|
||||
timestep = timestep * self.timestep_scale_multiplier
|
||||
timestep, embedded_timestep = self.adaln(
|
||||
timestep.flatten(),
|
||||
hidden_dtype=hidden_dtype,
|
||||
)
|
||||
|
||||
# Second dimension is 1 or number of tokens (if timestep_per_token)
|
||||
timestep = timestep.view(batch_size, -1, timestep.shape[-1])
|
||||
embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.shape[-1])
|
||||
return timestep, embedded_timestep
|
||||
|
||||
def _prepare_context(
|
||||
self,
|
||||
context: torch.Tensor,
|
||||
x: torch.Tensor,
|
||||
attention_mask: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
"""Prepare context for transformer blocks."""
|
||||
batch_size = x.shape[0]
|
||||
context = self.caption_projection(context)
|
||||
context = context.view(batch_size, -1, x.shape[-1])
|
||||
|
||||
return context, attention_mask
|
||||
|
||||
def _prepare_attention_mask(self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype) -> torch.Tensor | None:
|
||||
"""Prepare attention mask."""
|
||||
if attention_mask is None or torch.is_floating_point(attention_mask):
|
||||
return attention_mask
|
||||
|
||||
return (attention_mask - 1).to(x_dtype).reshape(
|
||||
(attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
|
||||
) * torch.finfo(x_dtype).max
|
||||
|
||||
def _prepare_positional_embeddings(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
inner_dim: int,
|
||||
max_pos: list[int],
|
||||
use_middle_indices_grid: bool,
|
||||
num_attention_heads: int,
|
||||
x_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""Prepare positional embeddings."""
|
||||
freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch
|
||||
pe = precompute_freqs_cis(
|
||||
positions,
|
||||
dim=inner_dim,
|
||||
out_dtype=x_dtype,
|
||||
theta=self.positional_embedding_theta,
|
||||
max_pos=max_pos,
|
||||
use_middle_indices_grid=use_middle_indices_grid,
|
||||
num_attention_heads=num_attention_heads,
|
||||
rope_type=self.rope_type,
|
||||
freq_grid_generator=freq_grid_generator,
|
||||
)
|
||||
return pe
|
||||
|
||||
def prepare(
|
||||
self,
|
||||
modality: Modality,
|
||||
) -> TransformerArgs:
|
||||
x = self.patchify_proj(modality.latent)
|
||||
timestep, embedded_timestep = self._prepare_timestep(modality.timesteps, x.shape[0], modality.latent.dtype)
|
||||
context, attention_mask = self._prepare_context(modality.context, x, modality.context_mask)
|
||||
attention_mask = self._prepare_attention_mask(attention_mask, modality.latent.dtype)
|
||||
pe = self._prepare_positional_embeddings(
|
||||
positions=modality.positions,
|
||||
inner_dim=self.inner_dim,
|
||||
max_pos=self.max_pos,
|
||||
use_middle_indices_grid=self.use_middle_indices_grid,
|
||||
num_attention_heads=self.num_attention_heads,
|
||||
x_dtype=modality.latent.dtype,
|
||||
)
|
||||
return TransformerArgs(
|
||||
x=x,
|
||||
context=context,
|
||||
context_mask=attention_mask,
|
||||
timesteps=timestep,
|
||||
embedded_timestep=embedded_timestep,
|
||||
positional_embeddings=pe,
|
||||
cross_positional_embeddings=None,
|
||||
cross_scale_shift_timestep=None,
|
||||
cross_gate_timestep=None,
|
||||
enabled=modality.enabled,
|
||||
)
|
||||
|
||||
|
||||
class MultiModalTransformerArgsPreprocessor:
|
||||
def __init__( # noqa: PLR0913
|
||||
self,
|
||||
patchify_proj: torch.nn.Linear,
|
||||
adaln: AdaLayerNormSingle,
|
||||
caption_projection: PixArtAlphaTextProjection,
|
||||
cross_scale_shift_adaln: AdaLayerNormSingle,
|
||||
cross_gate_adaln: AdaLayerNormSingle,
|
||||
inner_dim: int,
|
||||
max_pos: list[int],
|
||||
num_attention_heads: int,
|
||||
cross_pe_max_pos: int,
|
||||
use_middle_indices_grid: bool,
|
||||
audio_cross_attention_dim: int,
|
||||
timestep_scale_multiplier: int,
|
||||
double_precision_rope: bool,
|
||||
positional_embedding_theta: float,
|
||||
rope_type: LTXRopeType,
|
||||
av_ca_timestep_scale_multiplier: int,
|
||||
) -> None:
|
||||
self.simple_preprocessor = TransformerArgsPreprocessor(
|
||||
patchify_proj=patchify_proj,
|
||||
adaln=adaln,
|
||||
caption_projection=caption_projection,
|
||||
inner_dim=inner_dim,
|
||||
max_pos=max_pos,
|
||||
num_attention_heads=num_attention_heads,
|
||||
use_middle_indices_grid=use_middle_indices_grid,
|
||||
timestep_scale_multiplier=timestep_scale_multiplier,
|
||||
double_precision_rope=double_precision_rope,
|
||||
positional_embedding_theta=positional_embedding_theta,
|
||||
rope_type=rope_type,
|
||||
)
|
||||
self.cross_scale_shift_adaln = cross_scale_shift_adaln
|
||||
self.cross_gate_adaln = cross_gate_adaln
|
||||
self.cross_pe_max_pos = cross_pe_max_pos
|
||||
self.audio_cross_attention_dim = audio_cross_attention_dim
|
||||
self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier
|
||||
|
||||
def prepare(
|
||||
self,
|
||||
modality: Modality,
|
||||
) -> TransformerArgs:
|
||||
transformer_args = self.simple_preprocessor.prepare(modality)
|
||||
cross_pe = self.simple_preprocessor._prepare_positional_embeddings(
|
||||
positions=modality.positions[:, 0:1, :],
|
||||
inner_dim=self.audio_cross_attention_dim,
|
||||
max_pos=[self.cross_pe_max_pos],
|
||||
use_middle_indices_grid=True,
|
||||
num_attention_heads=self.simple_preprocessor.num_attention_heads,
|
||||
x_dtype=modality.latent.dtype,
|
||||
)
|
||||
|
||||
cross_scale_shift_timestep, cross_gate_timestep = self._prepare_cross_attention_timestep(
|
||||
timestep=modality.timesteps,
|
||||
timestep_scale_multiplier=self.simple_preprocessor.timestep_scale_multiplier,
|
||||
batch_size=transformer_args.x.shape[0],
|
||||
hidden_dtype=modality.latent.dtype,
|
||||
)
|
||||
|
||||
return replace(
|
||||
transformer_args,
|
||||
cross_positional_embeddings=cross_pe,
|
||||
cross_scale_shift_timestep=cross_scale_shift_timestep,
|
||||
cross_gate_timestep=cross_gate_timestep,
|
||||
)
|
||||
|
||||
def _prepare_cross_attention_timestep(
|
||||
self,
|
||||
timestep: torch.Tensor,
|
||||
timestep_scale_multiplier: int,
|
||||
batch_size: int,
|
||||
hidden_dtype: torch.dtype,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Prepare cross attention timestep embeddings."""
|
||||
timestep = timestep * timestep_scale_multiplier
|
||||
|
||||
av_ca_factor = self.av_ca_timestep_scale_multiplier / timestep_scale_multiplier
|
||||
|
||||
scale_shift_timestep, _ = self.cross_scale_shift_adaln(
|
||||
timestep.flatten(),
|
||||
hidden_dtype=hidden_dtype,
|
||||
)
|
||||
scale_shift_timestep = scale_shift_timestep.view(batch_size, -1, scale_shift_timestep.shape[-1])
|
||||
gate_noise_timestep, _ = self.cross_gate_adaln(
|
||||
timestep.flatten() * av_ca_factor,
|
||||
hidden_dtype=hidden_dtype,
|
||||
)
|
||||
gate_noise_timestep = gate_noise_timestep.view(batch_size, -1, gate_noise_timestep.shape[-1])
|
||||
|
||||
return scale_shift_timestep, gate_noise_timestep
|
||||
@@ -0,0 +1,10 @@
|
||||
"""Latent upsampler model components."""
|
||||
|
||||
from ltx_core.model.upsampler.model import LatentUpsampler, upsample_video
|
||||
from ltx_core.model.upsampler.model_configurator import LatentUpsamplerConfigurator
|
||||
|
||||
__all__ = [
|
||||
"LatentUpsampler",
|
||||
"LatentUpsamplerConfigurator",
|
||||
"upsample_video",
|
||||
]
|
||||
@@ -0,0 +1,53 @@
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange
|
||||
|
||||
|
||||
class BlurDownsample(torch.nn.Module):
|
||||
"""
|
||||
Anti-aliased spatial downsampling by integer stride using a fixed separable binomial kernel.
|
||||
Applies only on H,W. Works for dims=2 or dims=3 (per-frame).
|
||||
"""
|
||||
|
||||
def __init__(self, dims: int, stride: int, kernel_size: int = 5) -> None:
|
||||
super().__init__()
|
||||
assert dims in (2, 3)
|
||||
assert isinstance(stride, int)
|
||||
assert stride >= 1
|
||||
assert kernel_size >= 3
|
||||
assert kernel_size % 2 == 1
|
||||
self.dims = dims
|
||||
self.stride = stride
|
||||
self.kernel_size = kernel_size
|
||||
|
||||
# 5x5 separable binomial kernel using binomial coefficients [1, 4, 6, 4, 1] from
|
||||
# the 4th row of Pascal's triangle. This kernel is used for anti-aliasing and
|
||||
# provides a smooth approximation of a Gaussian filter (often called a "binomial filter").
|
||||
# The 2D kernel is constructed as the outer product and normalized.
|
||||
k = torch.tensor([math.comb(kernel_size - 1, k) for k in range(kernel_size)])
|
||||
k2d = k[:, None] @ k[None, :]
|
||||
k2d = (k2d / k2d.sum()).float() # shape (kernel_size, kernel_size)
|
||||
self.register_buffer("kernel", k2d[None, None, :, :]) # (1, 1, kernel_size, kernel_size)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if self.stride == 1:
|
||||
return x
|
||||
|
||||
if self.dims == 2:
|
||||
return self._apply_2d(x)
|
||||
else:
|
||||
# dims == 3: apply per-frame on H,W
|
||||
b, _, f, _, _ = x.shape
|
||||
x = rearrange(x, "b c f h w -> (b f) c h w")
|
||||
x = self._apply_2d(x)
|
||||
h2, w2 = x.shape[-2:]
|
||||
x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f, h=h2, w=w2)
|
||||
return x
|
||||
|
||||
def _apply_2d(self, x2d: torch.Tensor) -> torch.Tensor:
|
||||
c = x2d.shape[1]
|
||||
weight = self.kernel.expand(c, 1, self.kernel_size, self.kernel_size) # depthwise
|
||||
x2d = F.conv2d(x2d, weight=weight, bias=None, stride=self.stride, padding=self.kernel_size // 2, groups=c)
|
||||
return x2d
|
||||
@@ -0,0 +1,142 @@
|
||||
import torch
|
||||
from einops import rearrange
|
||||
|
||||
from ltx_core.model.upsampler.pixel_shuffle import PixelShuffleND
|
||||
from ltx_core.model.upsampler.res_block import ResBlock
|
||||
from ltx_core.model.upsampler.spatial_rational_resampler import SpatialRationalResampler
|
||||
from ltx_core.model.video_vae import VideoEncoder
|
||||
|
||||
|
||||
class LatentUpsampler(torch.nn.Module):
|
||||
"""
|
||||
Model to upsample VAE latents spatially and/or temporally.
|
||||
Args:
|
||||
in_channels (`int`): Number of channels in the input latent
|
||||
mid_channels (`int`): Number of channels in the middle layers
|
||||
num_blocks_per_stage (`int`): Number of ResBlocks to use in each stage (pre/post upsampling)
|
||||
dims (`int`): Number of dimensions for convolutions (2 or 3)
|
||||
spatial_upsample (`bool`): Whether to spatially upsample the latent
|
||||
temporal_upsample (`bool`): Whether to temporally upsample the latent
|
||||
spatial_scale (`float`): Scale factor for spatial upsampling
|
||||
rational_resampler (`bool`): Whether to use a rational resampler for spatial upsampling
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int = 128,
|
||||
mid_channels: int = 512,
|
||||
num_blocks_per_stage: int = 4,
|
||||
dims: int = 3,
|
||||
spatial_upsample: bool = True,
|
||||
temporal_upsample: bool = False,
|
||||
spatial_scale: float = 2.0,
|
||||
rational_resampler: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.mid_channels = mid_channels
|
||||
self.num_blocks_per_stage = num_blocks_per_stage
|
||||
self.dims = dims
|
||||
self.spatial_upsample = spatial_upsample
|
||||
self.temporal_upsample = temporal_upsample
|
||||
self.spatial_scale = float(spatial_scale)
|
||||
self.rational_resampler = rational_resampler
|
||||
|
||||
conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d
|
||||
|
||||
self.initial_conv = conv(in_channels, mid_channels, kernel_size=3, padding=1)
|
||||
self.initial_norm = torch.nn.GroupNorm(32, mid_channels)
|
||||
self.initial_activation = torch.nn.SiLU()
|
||||
|
||||
self.res_blocks = torch.nn.ModuleList([ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)])
|
||||
|
||||
if spatial_upsample and temporal_upsample:
|
||||
self.upsampler = torch.nn.Sequential(
|
||||
torch.nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1),
|
||||
PixelShuffleND(3),
|
||||
)
|
||||
elif spatial_upsample:
|
||||
if rational_resampler:
|
||||
self.upsampler = SpatialRationalResampler(mid_channels=mid_channels, scale=self.spatial_scale)
|
||||
else:
|
||||
self.upsampler = torch.nn.Sequential(
|
||||
torch.nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1),
|
||||
PixelShuffleND(2),
|
||||
)
|
||||
elif temporal_upsample:
|
||||
self.upsampler = torch.nn.Sequential(
|
||||
torch.nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1),
|
||||
PixelShuffleND(1),
|
||||
)
|
||||
else:
|
||||
raise ValueError("Either spatial_upsample or temporal_upsample must be True")
|
||||
|
||||
self.post_upsample_res_blocks = torch.nn.ModuleList(
|
||||
[ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]
|
||||
)
|
||||
|
||||
self.final_conv = conv(mid_channels, in_channels, kernel_size=3, padding=1)
|
||||
|
||||
def forward(self, latent: torch.Tensor) -> torch.Tensor:
|
||||
b, _, f, _, _ = latent.shape
|
||||
|
||||
if self.dims == 2:
|
||||
x = rearrange(latent, "b c f h w -> (b f) c h w")
|
||||
x = self.initial_conv(x)
|
||||
x = self.initial_norm(x)
|
||||
x = self.initial_activation(x)
|
||||
|
||||
for block in self.res_blocks:
|
||||
x = block(x)
|
||||
|
||||
x = self.upsampler(x)
|
||||
|
||||
for block in self.post_upsample_res_blocks:
|
||||
x = block(x)
|
||||
|
||||
x = self.final_conv(x)
|
||||
x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
|
||||
else:
|
||||
x = self.initial_conv(latent)
|
||||
x = self.initial_norm(x)
|
||||
x = self.initial_activation(x)
|
||||
|
||||
for block in self.res_blocks:
|
||||
x = block(x)
|
||||
|
||||
if self.temporal_upsample:
|
||||
x = self.upsampler(x)
|
||||
# remove the first frame after upsampling.
|
||||
# This is done because the first frame encodes one pixel frame.
|
||||
x = x[:, :, 1:, :, :]
|
||||
elif isinstance(self.upsampler, SpatialRationalResampler):
|
||||
x = self.upsampler(x)
|
||||
else:
|
||||
x = rearrange(x, "b c f h w -> (b f) c h w")
|
||||
x = self.upsampler(x)
|
||||
x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
|
||||
|
||||
for block in self.post_upsample_res_blocks:
|
||||
x = block(x)
|
||||
|
||||
x = self.final_conv(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def upsample_video(latent: torch.Tensor, video_encoder: VideoEncoder, upsampler: "LatentUpsampler") -> torch.Tensor:
|
||||
"""
|
||||
Apply upsampling to the latent representation using the provided upsampler,
|
||||
with normalization and un-normalization based on the video encoder's per-channel statistics.
|
||||
Args:
|
||||
latent: Input latent tensor of shape [B, C, F, H, W].
|
||||
video_encoder: VideoEncoder with per_channel_statistics for normalization.
|
||||
upsampler: LatentUpsampler module to perform upsampling.
|
||||
Returns:
|
||||
torch.Tensor: Upsampled and re-normalized latent tensor.
|
||||
"""
|
||||
latent = video_encoder.per_channel_statistics.un_normalize(latent)
|
||||
latent = upsampler(latent)
|
||||
latent = video_encoder.per_channel_statistics.normalize(latent)
|
||||
return latent
|
||||
@@ -0,0 +1,30 @@
|
||||
from ltx_core.model.model_protocol import ModelConfigurator
|
||||
from ltx_core.model.upsampler.model import LatentUpsampler
|
||||
|
||||
|
||||
class LatentUpsamplerConfigurator(ModelConfigurator[LatentUpsampler]):
|
||||
"""
|
||||
Configurator for LatentUpsampler model.
|
||||
Used to create a LatentUpsampler model from a configuration dictionary.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def from_config(cls: type[LatentUpsampler], config: dict) -> LatentUpsampler:
|
||||
in_channels = config.get("in_channels", 128)
|
||||
mid_channels = config.get("mid_channels", 512)
|
||||
num_blocks_per_stage = config.get("num_blocks_per_stage", 4)
|
||||
dims = config.get("dims", 3)
|
||||
spatial_upsample = config.get("spatial_upsample", True)
|
||||
temporal_upsample = config.get("temporal_upsample", False)
|
||||
spatial_scale = config.get("spatial_scale", 2.0)
|
||||
rational_resampler = config.get("rational_resampler", False)
|
||||
return LatentUpsampler(
|
||||
in_channels=in_channels,
|
||||
mid_channels=mid_channels,
|
||||
num_blocks_per_stage=num_blocks_per_stage,
|
||||
dims=dims,
|
||||
spatial_upsample=spatial_upsample,
|
||||
temporal_upsample=temporal_upsample,
|
||||
spatial_scale=spatial_scale,
|
||||
rational_resampler=rational_resampler,
|
||||
)
|
||||
@@ -0,0 +1,54 @@
|
||||
import torch
|
||||
from einops import rearrange
|
||||
|
||||
|
||||
class PixelShuffleND(torch.nn.Module):
|
||||
"""
|
||||
N-dimensional pixel shuffle operation for upsampling tensors.
|
||||
Args:
|
||||
dims (int): Number of dimensions to apply pixel shuffle to.
|
||||
- 1: Temporal (e.g., frames)
|
||||
- 2: Spatial (e.g., height and width)
|
||||
- 3: Spatiotemporal (e.g., depth, height, width)
|
||||
upscale_factors (tuple[int, int, int], optional): Upscaling factors for each dimension.
|
||||
For dims=1, only the first value is used.
|
||||
For dims=2, the first two values are used.
|
||||
For dims=3, all three values are used.
|
||||
The input tensor is rearranged so that the channel dimension is split into
|
||||
smaller channels and upscaling factors, and the upscaling factors are moved
|
||||
into the corresponding spatial/temporal dimensions.
|
||||
Note:
|
||||
This operation is equivalent to the patchifier operation in for the models. Consider
|
||||
using this class instead.
|
||||
"""
|
||||
|
||||
def __init__(self, dims: int, upscale_factors: tuple[int, int, int] = (2, 2, 2)):
|
||||
super().__init__()
|
||||
assert dims in [1, 2, 3], "dims must be 1, 2, or 3"
|
||||
self.dims = dims
|
||||
self.upscale_factors = upscale_factors
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if self.dims == 3:
|
||||
return rearrange(
|
||||
x,
|
||||
"b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
|
||||
p1=self.upscale_factors[0],
|
||||
p2=self.upscale_factors[1],
|
||||
p3=self.upscale_factors[2],
|
||||
)
|
||||
elif self.dims == 2:
|
||||
return rearrange(
|
||||
x,
|
||||
"b (c p1 p2) h w -> b c (h p1) (w p2)",
|
||||
p1=self.upscale_factors[0],
|
||||
p2=self.upscale_factors[1],
|
||||
)
|
||||
elif self.dims == 1:
|
||||
return rearrange(
|
||||
x,
|
||||
"b (c p1) f h w -> b c (f p1) h w",
|
||||
p1=self.upscale_factors[0],
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported dims: {self.dims}")
|
||||
@@ -0,0 +1,37 @@
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class ResBlock(torch.nn.Module):
|
||||
"""
|
||||
Residual block with two convolutional layers, group normalization, and SiLU activation.
|
||||
Args:
|
||||
channels (int): Number of input and output channels.
|
||||
mid_channels (Optional[int]): Number of channels in the intermediate convolution layer. Defaults to `channels`
|
||||
if not specified.
|
||||
dims (int): Dimensionality of the convolution (2 for Conv2d, 3 for Conv3d). Defaults to 3.
|
||||
"""
|
||||
|
||||
def __init__(self, channels: int, mid_channels: Optional[int] = None, dims: int = 3):
|
||||
super().__init__()
|
||||
if mid_channels is None:
|
||||
mid_channels = channels
|
||||
|
||||
conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d
|
||||
|
||||
self.conv1 = conv(channels, mid_channels, kernel_size=3, padding=1)
|
||||
self.norm1 = torch.nn.GroupNorm(32, mid_channels)
|
||||
self.conv2 = conv(mid_channels, channels, kernel_size=3, padding=1)
|
||||
self.norm2 = torch.nn.GroupNorm(32, channels)
|
||||
self.activation = torch.nn.SiLU()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
residual = x
|
||||
x = self.conv1(x)
|
||||
x = self.norm1(x)
|
||||
x = self.activation(x)
|
||||
x = self.conv2(x)
|
||||
x = self.norm2(x)
|
||||
x = self.activation(x + residual)
|
||||
return x
|
||||
@@ -0,0 +1,47 @@
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
|
||||
from ltx_core.model.upsampler.blur_downsample import BlurDownsample
|
||||
from ltx_core.model.upsampler.pixel_shuffle import PixelShuffleND
|
||||
|
||||
|
||||
def _rational_for_scale(scale: float) -> Tuple[int, int]:
|
||||
mapping = {0.75: (3, 4), 1.5: (3, 2), 2.0: (2, 1), 4.0: (4, 1)}
|
||||
if float(scale) not in mapping:
|
||||
raise ValueError(f"Unsupported scale {scale}. Choose from {list(mapping.keys())}")
|
||||
return mapping[float(scale)]
|
||||
|
||||
|
||||
class SpatialRationalResampler(torch.nn.Module):
|
||||
"""
|
||||
Fully-learned rational spatial scaling: up by 'num' via PixelShuffle, then anti-aliased
|
||||
downsample by 'den' using fixed blur + stride. Operates on H,W only.
|
||||
For dims==3, work per-frame for spatial scaling (temporal axis untouched).
|
||||
Args:
|
||||
mid_channels (`int`): Number of intermediate channels for the convolution layer
|
||||
scale (`float`): Spatial scaling factor. Supported values are:
|
||||
- 0.75: Downsample by 3/4 (reduce spatial size)
|
||||
- 1.5: Upsample by 3/2 (increase spatial size)
|
||||
- 2.0: Upsample by 2x (double spatial size)
|
||||
- 4.0: Upsample by 4x (quadruple spatial size)
|
||||
Any other value will raise a ValueError.
|
||||
"""
|
||||
|
||||
def __init__(self, mid_channels: int, scale: float):
|
||||
super().__init__()
|
||||
self.scale = float(scale)
|
||||
self.num, self.den = _rational_for_scale(self.scale)
|
||||
self.conv = torch.nn.Conv2d(mid_channels, (self.num**2) * mid_channels, kernel_size=3, padding=1)
|
||||
self.pixel_shuffle = PixelShuffleND(2, upscale_factors=(self.num, self.num))
|
||||
self.blur_down = BlurDownsample(dims=2, stride=self.den)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
b, _, f, _, _ = x.shape
|
||||
x = rearrange(x, "b c f h w -> (b f) c h w")
|
||||
x = self.conv(x)
|
||||
x = self.pixel_shuffle(x)
|
||||
x = self.blur_down(x)
|
||||
x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
|
||||
return x
|
||||
@@ -0,0 +1,24 @@
|
||||
"""Video VAE package."""
|
||||
|
||||
from ltx_core.model.video_vae.model_configurator import (
|
||||
VAE_DECODER_COMFY_KEYS_FILTER,
|
||||
VAE_ENCODER_COMFY_KEYS_FILTER,
|
||||
VideoDecoderConfigurator,
|
||||
VideoEncoderConfigurator,
|
||||
)
|
||||
from ltx_core.model.video_vae.tiling import SpatialTilingConfig, TemporalTilingConfig, TilingConfig
|
||||
from ltx_core.model.video_vae.video_vae import VideoDecoder, VideoEncoder, decode_video, get_video_chunks_number
|
||||
|
||||
__all__ = [
|
||||
"VAE_DECODER_COMFY_KEYS_FILTER",
|
||||
"VAE_ENCODER_COMFY_KEYS_FILTER",
|
||||
"SpatialTilingConfig",
|
||||
"TemporalTilingConfig",
|
||||
"TilingConfig",
|
||||
"VideoDecoder",
|
||||
"VideoDecoderConfigurator",
|
||||
"VideoEncoder",
|
||||
"VideoEncoderConfigurator",
|
||||
"decode_video",
|
||||
"get_video_chunks_number",
|
||||
]
|
||||
@@ -0,0 +1,317 @@
|
||||
from typing import Tuple, Union
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from ltx_core.model.video_vae.enums import PaddingModeType
|
||||
|
||||
|
||||
def make_conv_nd( # noqa: PLR0913
|
||||
dims: Union[int, Tuple[int, int]],
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: int,
|
||||
stride: int = 1,
|
||||
padding: int = 0,
|
||||
dilation: int = 1,
|
||||
groups: int = 1,
|
||||
bias: bool = True,
|
||||
causal: bool = False,
|
||||
spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
|
||||
temporal_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
|
||||
) -> nn.Module:
|
||||
if not (spatial_padding_mode == temporal_padding_mode or causal):
|
||||
raise NotImplementedError("spatial and temporal padding modes must be equal")
|
||||
if dims == 2:
|
||||
return nn.Conv2d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
padding_mode=spatial_padding_mode.value,
|
||||
)
|
||||
elif dims == 3:
|
||||
if causal:
|
||||
return CausalConv3d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
return nn.Conv3d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
padding_mode=spatial_padding_mode.value,
|
||||
)
|
||||
elif dims == (2, 1):
|
||||
return DualConv3d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
bias=bias,
|
||||
padding_mode=spatial_padding_mode.value,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"unsupported dimensions: {dims}")
|
||||
|
||||
|
||||
def make_linear_nd(
|
||||
dims: int,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
bias: bool = True,
|
||||
) -> nn.Module:
|
||||
if dims == 2:
|
||||
return nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias)
|
||||
elif dims in (3, (2, 1)):
|
||||
return nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias)
|
||||
else:
|
||||
raise ValueError(f"unsupported dimensions: {dims}")
|
||||
|
||||
|
||||
class DualConv3d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: int,
|
||||
stride: Union[int, Tuple[int, int, int]] = 1,
|
||||
padding: Union[int, Tuple[int, int, int]] = 0,
|
||||
dilation: Union[int, Tuple[int, int, int]] = 1,
|
||||
groups: int = 1,
|
||||
bias: bool = True,
|
||||
padding_mode: str = "zeros",
|
||||
) -> None:
|
||||
super(DualConv3d, self).__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.padding_mode = padding_mode
|
||||
# Ensure kernel_size, stride, padding, and dilation are tuples of length 3
|
||||
if isinstance(kernel_size, int):
|
||||
kernel_size = (kernel_size, kernel_size, kernel_size)
|
||||
if kernel_size == (1, 1, 1):
|
||||
raise ValueError("kernel_size must be greater than 1. Use make_linear_nd instead.")
|
||||
if isinstance(stride, int):
|
||||
stride = (stride, stride, stride)
|
||||
if isinstance(padding, int):
|
||||
padding = (padding, padding, padding)
|
||||
if isinstance(dilation, int):
|
||||
dilation = (dilation, dilation, dilation)
|
||||
|
||||
# Set parameters for convolutions
|
||||
self.groups = groups
|
||||
self.bias = bias
|
||||
|
||||
# Define the size of the channels after the first convolution
|
||||
intermediate_channels = out_channels if in_channels < out_channels else in_channels
|
||||
|
||||
# Define parameters for the first convolution
|
||||
self.weight1 = nn.Parameter(
|
||||
torch.Tensor(
|
||||
intermediate_channels,
|
||||
in_channels // groups,
|
||||
1,
|
||||
kernel_size[1],
|
||||
kernel_size[2],
|
||||
)
|
||||
)
|
||||
self.stride1 = (1, stride[1], stride[2])
|
||||
self.padding1 = (0, padding[1], padding[2])
|
||||
self.dilation1 = (1, dilation[1], dilation[2])
|
||||
if bias:
|
||||
self.bias1 = nn.Parameter(torch.Tensor(intermediate_channels))
|
||||
else:
|
||||
self.register_parameter("bias1", None)
|
||||
|
||||
# Define parameters for the second convolution
|
||||
self.weight2 = nn.Parameter(torch.Tensor(out_channels, intermediate_channels // groups, kernel_size[0], 1, 1))
|
||||
self.stride2 = (stride[0], 1, 1)
|
||||
self.padding2 = (padding[0], 0, 0)
|
||||
self.dilation2 = (dilation[0], 1, 1)
|
||||
if bias:
|
||||
self.bias2 = nn.Parameter(torch.Tensor(out_channels))
|
||||
else:
|
||||
self.register_parameter("bias2", None)
|
||||
|
||||
# Initialize weights and biases
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
nn.init.kaiming_uniform_(self.weight1, a=torch.sqrt(5))
|
||||
nn.init.kaiming_uniform_(self.weight2, a=torch.sqrt(5))
|
||||
if self.bias:
|
||||
fan_in1, _ = nn.init._calculate_fan_in_and_fan_out(self.weight1)
|
||||
bound1 = 1 / torch.sqrt(fan_in1)
|
||||
nn.init.uniform_(self.bias1, -bound1, bound1)
|
||||
fan_in2, _ = nn.init._calculate_fan_in_and_fan_out(self.weight2)
|
||||
bound2 = 1 / torch.sqrt(fan_in2)
|
||||
nn.init.uniform_(self.bias2, -bound2, bound2)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
use_conv3d: bool = False,
|
||||
skip_time_conv: bool = False,
|
||||
) -> torch.Tensor:
|
||||
if use_conv3d:
|
||||
return self.forward_with_3d(x=x, skip_time_conv=skip_time_conv)
|
||||
else:
|
||||
return self.forward_with_2d(x=x, skip_time_conv=skip_time_conv)
|
||||
|
||||
def forward_with_3d(self, x: torch.Tensor, skip_time_conv: bool = False) -> torch.Tensor:
|
||||
# First convolution
|
||||
x = F.conv3d(
|
||||
x,
|
||||
self.weight1,
|
||||
self.bias1,
|
||||
self.stride1,
|
||||
self.padding1,
|
||||
self.dilation1,
|
||||
self.groups,
|
||||
padding_mode=self.padding_mode,
|
||||
)
|
||||
|
||||
if skip_time_conv:
|
||||
return x
|
||||
|
||||
# Second convolution
|
||||
x = F.conv3d(
|
||||
x,
|
||||
self.weight2,
|
||||
self.bias2,
|
||||
self.stride2,
|
||||
self.padding2,
|
||||
self.dilation2,
|
||||
self.groups,
|
||||
padding_mode=self.padding_mode,
|
||||
)
|
||||
|
||||
return x
|
||||
|
||||
def forward_with_2d(self, x: torch.Tensor, skip_time_conv: bool = False) -> torch.Tensor:
|
||||
b, _, _, h, w = x.shape
|
||||
|
||||
# First 2D convolution
|
||||
x = rearrange(x, "b c d h w -> (b d) c h w")
|
||||
# Squeeze the depth dimension out of weight1 since it's 1
|
||||
weight1 = self.weight1.squeeze(2)
|
||||
# Select stride, padding, and dilation for the 2D convolution
|
||||
stride1 = (self.stride1[1], self.stride1[2])
|
||||
padding1 = (self.padding1[1], self.padding1[2])
|
||||
dilation1 = (self.dilation1[1], self.dilation1[2])
|
||||
x = F.conv2d(
|
||||
x,
|
||||
weight1,
|
||||
self.bias1,
|
||||
stride1,
|
||||
padding1,
|
||||
dilation1,
|
||||
self.groups,
|
||||
padding_mode=self.padding_mode,
|
||||
)
|
||||
|
||||
_, _, h, w = x.shape
|
||||
|
||||
if skip_time_conv:
|
||||
x = rearrange(x, "(b d) c h w -> b c d h w", b=b)
|
||||
return x
|
||||
|
||||
# Second convolution which is essentially treated as a 1D convolution across the 'd' dimension
|
||||
x = rearrange(x, "(b d) c h w -> (b h w) c d", b=b)
|
||||
|
||||
# Reshape weight2 to match the expected dimensions for conv1d
|
||||
weight2 = self.weight2.squeeze(-1).squeeze(-1)
|
||||
# Use only the relevant dimension for stride, padding, and dilation for the 1D convolution
|
||||
stride2 = self.stride2[0]
|
||||
padding2 = self.padding2[0]
|
||||
dilation2 = self.dilation2[0]
|
||||
x = F.conv1d(
|
||||
x,
|
||||
weight2,
|
||||
self.bias2,
|
||||
stride2,
|
||||
padding2,
|
||||
dilation2,
|
||||
self.groups,
|
||||
padding_mode=self.padding_mode,
|
||||
)
|
||||
x = rearrange(x, "(b h w) c d -> b c d h w", b=b, h=h, w=w)
|
||||
|
||||
return x
|
||||
|
||||
@property
|
||||
def weight(self) -> torch.Tensor:
|
||||
return self.weight2
|
||||
|
||||
|
||||
class CausalConv3d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: int = 3,
|
||||
stride: Union[int, Tuple[int]] = 1,
|
||||
dilation: int = 1,
|
||||
groups: int = 1,
|
||||
bias: bool = True,
|
||||
spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
|
||||
kernel_size = (kernel_size, kernel_size, kernel_size)
|
||||
self.time_kernel_size = kernel_size[0]
|
||||
|
||||
dilation = (dilation, 1, 1)
|
||||
|
||||
height_pad = kernel_size[1] // 2
|
||||
width_pad = kernel_size[2] // 2
|
||||
padding = (0, height_pad, width_pad)
|
||||
|
||||
self.conv = nn.Conv3d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
dilation=dilation,
|
||||
padding=padding,
|
||||
padding_mode=spatial_padding_mode.value,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor, causal: bool = True) -> torch.Tensor:
|
||||
if causal:
|
||||
first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, self.time_kernel_size - 1, 1, 1))
|
||||
x = torch.concatenate((first_frame_pad, x), dim=2)
|
||||
else:
|
||||
first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1))
|
||||
last_frame_pad = x[:, :, -1:, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1))
|
||||
x = torch.concatenate((first_frame_pad, x, last_frame_pad), dim=2)
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
@property
|
||||
def weight(self) -> torch.Tensor:
|
||||
return self.conv.weight
|
||||
@@ -0,0 +1,20 @@
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class NormLayerType(Enum):
|
||||
GROUP_NORM = "group_norm"
|
||||
PIXEL_NORM = "pixel_norm"
|
||||
|
||||
|
||||
class LogVarianceType(Enum):
|
||||
PER_CHANNEL = "per_channel"
|
||||
UNIFORM = "uniform"
|
||||
CONSTANT = "constant"
|
||||
NONE = "none"
|
||||
|
||||
|
||||
class PaddingModeType(Enum):
|
||||
ZEROS = "zeros"
|
||||
REFLECT = "reflect"
|
||||
REPLICATE = "replicate"
|
||||
CIRCULAR = "circular"
|
||||
@@ -0,0 +1,77 @@
|
||||
from ltx_core.loader.sd_ops import SDOps
|
||||
from ltx_core.model.model_protocol import ModelConfigurator
|
||||
from ltx_core.model.video_vae.enums import LogVarianceType, NormLayerType, PaddingModeType
|
||||
from ltx_core.model.video_vae.video_vae import VideoDecoder, VideoEncoder
|
||||
|
||||
|
||||
class VideoEncoderConfigurator(ModelConfigurator[VideoEncoder]):
|
||||
"""Configurator for creating a video VAE Encoder from a configuration dictionary."""
|
||||
|
||||
@classmethod
|
||||
def from_config(cls: type[VideoEncoder], config: dict) -> VideoEncoder:
|
||||
config = config.get("vae", {})
|
||||
convolution_dimensions = config.get("dims", 3)
|
||||
in_channels = config.get("in_channels", 3)
|
||||
latent_channels = config.get("latent_channels", 128)
|
||||
encoder_spatial_padding_mode = PaddingModeType(config.get("encoder_spatial_padding_mode", "zeros"))
|
||||
encoder_blocks = config.get("encoder_blocks", [])
|
||||
patch_size = config.get("patch_size", 4)
|
||||
norm_layer_str = config.get("norm_layer", "pixel_norm")
|
||||
latent_log_var_str = config.get("latent_log_var", "uniform")
|
||||
|
||||
return VideoEncoder(
|
||||
convolution_dimensions=convolution_dimensions,
|
||||
in_channels=in_channels,
|
||||
out_channels=latent_channels,
|
||||
encoder_blocks=encoder_blocks,
|
||||
patch_size=patch_size,
|
||||
norm_layer=NormLayerType(norm_layer_str),
|
||||
latent_log_var=LogVarianceType(latent_log_var_str),
|
||||
encoder_spatial_padding_mode=encoder_spatial_padding_mode,
|
||||
)
|
||||
|
||||
|
||||
class VideoDecoderConfigurator(ModelConfigurator[VideoDecoder]):
|
||||
"""Configurator for creating a video VAE Decoder from a configuration dictionary."""
|
||||
|
||||
@classmethod
|
||||
def from_config(cls: type[VideoDecoder], config: dict) -> VideoDecoder:
|
||||
config = config.get("vae", {})
|
||||
convolution_dimensions = config.get("dims", 3)
|
||||
latent_channels = config.get("latent_channels", 128)
|
||||
decoder_spatial_padding_mode = PaddingModeType(config.get("decoder_spatial_padding_mode", "reflect"))
|
||||
out_channels = config.get("out_channels", 3)
|
||||
decoder_blocks = config.get("decoder_blocks", [])
|
||||
patch_size = config.get("patch_size", 4)
|
||||
norm_layer_str = config.get("norm_layer", "pixel_norm")
|
||||
causal = config.get("causal_decoder", False)
|
||||
timestep_conditioning = config.get("timestep_conditioning", True)
|
||||
|
||||
return VideoDecoder(
|
||||
convolution_dimensions=convolution_dimensions,
|
||||
in_channels=latent_channels,
|
||||
out_channels=out_channels,
|
||||
decoder_blocks=decoder_blocks,
|
||||
patch_size=patch_size,
|
||||
norm_layer=NormLayerType(norm_layer_str),
|
||||
causal=causal,
|
||||
timestep_conditioning=timestep_conditioning,
|
||||
decoder_spatial_padding_mode=decoder_spatial_padding_mode,
|
||||
)
|
||||
|
||||
|
||||
VAE_DECODER_COMFY_KEYS_FILTER = (
|
||||
SDOps("VAE_DECODER_COMFY_KEYS_FILTER")
|
||||
.with_matching(prefix="vae.decoder.")
|
||||
.with_matching(prefix="vae.per_channel_statistics.")
|
||||
.with_replacement("vae.decoder.", "")
|
||||
.with_replacement("vae.per_channel_statistics.", "per_channel_statistics.")
|
||||
)
|
||||
|
||||
VAE_ENCODER_COMFY_KEYS_FILTER = (
|
||||
SDOps("VAE_ENCODER_COMFY_KEYS_FILTER")
|
||||
.with_matching(prefix="vae.encoder.")
|
||||
.with_matching(prefix="vae.per_channel_statistics.")
|
||||
.with_replacement("vae.encoder.", "")
|
||||
.with_replacement("vae.per_channel_statistics.", "per_channel_statistics.")
|
||||
)
|
||||
@@ -0,0 +1,3 @@
|
||||
from ltx_core.model.common.normalization import PixelNorm, build_normalization_layer
|
||||
|
||||
__all__ = ["PixelNorm", "build_normalization_layer"]
|
||||
@@ -0,0 +1,85 @@
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import nn
|
||||
|
||||
|
||||
def patchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor:
|
||||
"""
|
||||
Rearrange spatial dimensions into channels. Divides image into patch_size x patch_size blocks
|
||||
and moves pixels from each block into separate channels (space-to-depth).
|
||||
Args:
|
||||
x: Input tensor (4D or 5D)
|
||||
patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, divides HxW into 4x4 blocks.
|
||||
patch_size_t: Temporal patch size for frames. Default=1 (no temporal patching).
|
||||
For 5D: (B, C, F, H, W) -> (B, Cx(patch_size_hw^2)x(patch_size_t), F/patch_size_t, H/patch_size_hw, W/patch_size_hw)
|
||||
Example: (B, 3, 33, 512, 512) with patch_size_hw=4, patch_size_t=1 -> (B, 48, 33, 128, 128)
|
||||
"""
|
||||
if patch_size_hw == 1 and patch_size_t == 1:
|
||||
return x
|
||||
if x.dim() == 4:
|
||||
x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size_hw, r=patch_size_hw)
|
||||
elif x.dim() == 5:
|
||||
x = rearrange(
|
||||
x,
|
||||
"b c (f p) (h q) (w r) -> b (c p r q) f h w",
|
||||
p=patch_size_t,
|
||||
q=patch_size_hw,
|
||||
r=patch_size_hw,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid input shape: {x.shape}")
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def unpatchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor:
|
||||
"""
|
||||
Rearrange channels back into spatial dimensions. Inverse of patchify - moves pixels from
|
||||
channels back into patch_size x patch_size blocks (depth-to-space).
|
||||
Args:
|
||||
x: Input tensor (4D or 5D)
|
||||
patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, expands HxW by 4x.
|
||||
patch_size_t: Temporal patch size for frames. Default=1 (no temporal expansion).
|
||||
For 5D: (B, Cx(patch_size_hw^2)x(patch_size_t), F, H, W) -> (B, C, Fxpatch_size_t, Hxpatch_size_hw, Wxpatch_size_hw)
|
||||
Example: (B, 48, 33, 128, 128) with patch_size_hw=4, patch_size_t=1 -> (B, 3, 33, 512, 512)
|
||||
"""
|
||||
if patch_size_hw == 1 and patch_size_t == 1:
|
||||
return x
|
||||
|
||||
if x.dim() == 4:
|
||||
x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size_hw, r=patch_size_hw)
|
||||
elif x.dim() == 5:
|
||||
x = rearrange(
|
||||
x,
|
||||
"b (c p r q) f h w -> b c (f p) (h q) (w r)",
|
||||
p=patch_size_t,
|
||||
q=patch_size_hw,
|
||||
r=patch_size_hw,
|
||||
)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class PerChannelStatistics(nn.Module):
|
||||
"""
|
||||
Per-channel statistics for normalizing and denormalizing the latent representation.
|
||||
This statics is computed over the entire dataset and stored in model's checkpoint under VAE state_dict.
|
||||
"""
|
||||
|
||||
def __init__(self, latent_channels: int = 128):
|
||||
super().__init__()
|
||||
self.register_buffer("std-of-means", torch.empty(latent_channels))
|
||||
self.register_buffer("mean-of-means", torch.empty(latent_channels))
|
||||
self.register_buffer("mean-of-stds", torch.empty(latent_channels))
|
||||
self.register_buffer("mean-of-stds_over_std-of-means", torch.empty(latent_channels))
|
||||
self.register_buffer("channel", torch.empty(latent_channels))
|
||||
|
||||
def un_normalize(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return (x * self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)) + self.get_buffer("mean-of-means").view(
|
||||
1, -1, 1, 1, 1
|
||||
).to(x)
|
||||
|
||||
def normalize(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return (x - self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)) / self.get_buffer("std-of-means").view(
|
||||
1, -1, 1, 1, 1
|
||||
).to(x)
|
||||
@@ -0,0 +1,277 @@
|
||||
from typing import Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from ltx_core.model.common.normalization import PixelNorm
|
||||
from ltx_core.model.transformer.timestep_embedding import PixArtAlphaCombinedTimestepSizeEmbeddings
|
||||
from ltx_core.model.video_vae.convolution import make_conv_nd, make_linear_nd
|
||||
from ltx_core.model.video_vae.enums import NormLayerType, PaddingModeType
|
||||
|
||||
|
||||
class ResnetBlock3D(nn.Module):
|
||||
r"""
|
||||
A Resnet block.
|
||||
Parameters:
|
||||
in_channels (`int`): The number of channels in the input.
|
||||
out_channels (`int`, *optional*, default to be `None`):
|
||||
The number of output channels for the first conv layer. If None, same as `in_channels`.
|
||||
dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use.
|
||||
groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer.
|
||||
eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dims: Union[int, Tuple[int, int]],
|
||||
in_channels: int,
|
||||
out_channels: Optional[int] = None,
|
||||
dropout: float = 0.0,
|
||||
groups: int = 32,
|
||||
eps: float = 1e-6,
|
||||
norm_layer: NormLayerType = NormLayerType.PIXEL_NORM,
|
||||
inject_noise: bool = False,
|
||||
timestep_conditioning: bool = False,
|
||||
spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
|
||||
):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
out_channels = in_channels if out_channels is None else out_channels
|
||||
self.out_channels = out_channels
|
||||
self.inject_noise = inject_noise
|
||||
|
||||
if norm_layer == NormLayerType.GROUP_NORM:
|
||||
self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
|
||||
elif norm_layer == NormLayerType.PIXEL_NORM:
|
||||
self.norm1 = PixelNorm()
|
||||
|
||||
self.non_linearity = nn.SiLU()
|
||||
|
||||
self.conv1 = make_conv_nd(
|
||||
dims,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
causal=True,
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
|
||||
if inject_noise:
|
||||
self.per_channel_scale1 = nn.Parameter(torch.zeros((in_channels, 1, 1)))
|
||||
|
||||
if norm_layer == NormLayerType.GROUP_NORM:
|
||||
self.norm2 = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps, affine=True)
|
||||
elif norm_layer == NormLayerType.PIXEL_NORM:
|
||||
self.norm2 = PixelNorm()
|
||||
|
||||
self.dropout = torch.nn.Dropout(dropout)
|
||||
|
||||
self.conv2 = make_conv_nd(
|
||||
dims,
|
||||
out_channels,
|
||||
out_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
causal=True,
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
|
||||
if inject_noise:
|
||||
self.per_channel_scale2 = nn.Parameter(torch.zeros((in_channels, 1, 1)))
|
||||
|
||||
self.conv_shortcut = (
|
||||
make_linear_nd(dims=dims, in_channels=in_channels, out_channels=out_channels)
|
||||
if in_channels != out_channels
|
||||
else nn.Identity()
|
||||
)
|
||||
|
||||
# Using GroupNorm with 1 group is equivalent to LayerNorm but works with (B, C, ...) layout
|
||||
# avoiding the need for dimension rearrangement used in standard nn.LayerNorm
|
||||
self.norm3 = (
|
||||
nn.GroupNorm(num_groups=1, num_channels=in_channels, eps=eps, affine=True)
|
||||
if in_channels != out_channels
|
||||
else nn.Identity()
|
||||
)
|
||||
|
||||
self.timestep_conditioning = timestep_conditioning
|
||||
|
||||
if timestep_conditioning:
|
||||
self.scale_shift_table = nn.Parameter(torch.randn(4, in_channels) / in_channels**0.5)
|
||||
|
||||
def _feed_spatial_noise(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
per_channel_scale: torch.Tensor,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
) -> torch.Tensor:
|
||||
spatial_shape = hidden_states.shape[-2:]
|
||||
device = hidden_states.device
|
||||
dtype = hidden_states.dtype
|
||||
|
||||
# similar to the "explicit noise inputs" method in style-gan
|
||||
spatial_noise = torch.randn(spatial_shape, device=device, dtype=dtype, generator=generator)[None]
|
||||
scaled_noise = (spatial_noise * per_channel_scale)[None, :, None, ...]
|
||||
hidden_states = hidden_states + scaled_noise
|
||||
|
||||
return hidden_states
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_tensor: torch.Tensor,
|
||||
causal: bool = True,
|
||||
timestep: Optional[torch.Tensor] = None,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
) -> torch.Tensor:
|
||||
hidden_states = input_tensor
|
||||
batch_size = hidden_states.shape[0]
|
||||
|
||||
hidden_states = self.norm1(hidden_states)
|
||||
if self.timestep_conditioning:
|
||||
if timestep is None:
|
||||
raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True")
|
||||
ada_values = self.scale_shift_table[None, ..., None, None, None].to(
|
||||
device=hidden_states.device, dtype=hidden_states.dtype
|
||||
) + timestep.reshape(
|
||||
batch_size,
|
||||
4,
|
||||
-1,
|
||||
timestep.shape[-3],
|
||||
timestep.shape[-2],
|
||||
timestep.shape[-1],
|
||||
)
|
||||
shift1, scale1, shift2, scale2 = ada_values.unbind(dim=1)
|
||||
|
||||
hidden_states = hidden_states * (1 + scale1) + shift1
|
||||
|
||||
hidden_states = self.non_linearity(hidden_states)
|
||||
|
||||
hidden_states = self.conv1(hidden_states, causal=causal)
|
||||
|
||||
if self.inject_noise:
|
||||
hidden_states = self._feed_spatial_noise(
|
||||
hidden_states,
|
||||
self.per_channel_scale1.to(device=hidden_states.device, dtype=hidden_states.dtype),
|
||||
generator=generator,
|
||||
)
|
||||
|
||||
hidden_states = self.norm2(hidden_states)
|
||||
|
||||
if self.timestep_conditioning:
|
||||
hidden_states = hidden_states * (1 + scale2) + shift2
|
||||
|
||||
hidden_states = self.non_linearity(hidden_states)
|
||||
|
||||
hidden_states = self.dropout(hidden_states)
|
||||
|
||||
hidden_states = self.conv2(hidden_states, causal=causal)
|
||||
|
||||
if self.inject_noise:
|
||||
hidden_states = self._feed_spatial_noise(
|
||||
hidden_states,
|
||||
self.per_channel_scale2.to(device=hidden_states.device, dtype=hidden_states.dtype),
|
||||
generator=generator,
|
||||
)
|
||||
|
||||
input_tensor = self.norm3(input_tensor)
|
||||
|
||||
batch_size = input_tensor.shape[0]
|
||||
|
||||
input_tensor = self.conv_shortcut(input_tensor)
|
||||
|
||||
output_tensor = input_tensor + hidden_states
|
||||
|
||||
return output_tensor
|
||||
|
||||
|
||||
class UNetMidBlock3D(nn.Module):
|
||||
"""
|
||||
A 3D UNet mid-block [`UNetMidBlock3D`] with multiple residual blocks.
|
||||
Args:
|
||||
in_channels (`int`): The number of input channels.
|
||||
dropout (`float`, *optional*, defaults to 0.0): The dropout rate.
|
||||
num_layers (`int`, *optional*, defaults to 1): The number of residual blocks.
|
||||
resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks.
|
||||
resnet_groups (`int`, *optional*, defaults to 32):
|
||||
The number of groups to use in the group normalization layers of the resnet blocks.
|
||||
norm_layer (`str`, *optional*, defaults to `group_norm`):
|
||||
The normalization layer to use. Can be either `group_norm` or `pixel_norm`.
|
||||
inject_noise (`bool`, *optional*, defaults to `False`):
|
||||
Whether to inject noise into the hidden states.
|
||||
timestep_conditioning (`bool`, *optional*, defaults to `False`):
|
||||
Whether to condition the hidden states on the timestep.
|
||||
Returns:
|
||||
`torch.Tensor`: The output of the last residual block, which is a tensor of shape `(batch_size,
|
||||
in_channels, height, width)`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dims: Union[int, Tuple[int, int]],
|
||||
in_channels: int,
|
||||
dropout: float = 0.0,
|
||||
num_layers: int = 1,
|
||||
resnet_eps: float = 1e-6,
|
||||
resnet_groups: int = 32,
|
||||
norm_layer: NormLayerType = NormLayerType.GROUP_NORM,
|
||||
inject_noise: bool = False,
|
||||
timestep_conditioning: bool = False,
|
||||
spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
|
||||
):
|
||||
super().__init__()
|
||||
resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
|
||||
|
||||
self.timestep_conditioning = timestep_conditioning
|
||||
|
||||
if timestep_conditioning:
|
||||
self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(
|
||||
embedding_dim=in_channels * 4, size_emb_dim=0
|
||||
)
|
||||
|
||||
self.res_blocks = nn.ModuleList(
|
||||
[
|
||||
ResnetBlock3D(
|
||||
dims=dims,
|
||||
in_channels=in_channels,
|
||||
out_channels=in_channels,
|
||||
eps=resnet_eps,
|
||||
groups=resnet_groups,
|
||||
dropout=dropout,
|
||||
norm_layer=norm_layer,
|
||||
inject_noise=inject_noise,
|
||||
timestep_conditioning=timestep_conditioning,
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
for _ in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
causal: bool = True,
|
||||
timestep: Optional[torch.Tensor] = None,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
) -> torch.Tensor:
|
||||
timestep_embed = None
|
||||
if self.timestep_conditioning:
|
||||
if timestep is None:
|
||||
raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True")
|
||||
batch_size = hidden_states.shape[0]
|
||||
timestep_embed = self.time_embedder(
|
||||
timestep=timestep.flatten(),
|
||||
hidden_dtype=hidden_states.dtype,
|
||||
)
|
||||
timestep_embed = timestep_embed.view(batch_size, timestep_embed.shape[-1], 1, 1, 1)
|
||||
|
||||
for resnet in self.res_blocks:
|
||||
hidden_states = resnet(
|
||||
hidden_states,
|
||||
causal=causal,
|
||||
timestep=timestep_embed,
|
||||
generator=generator,
|
||||
)
|
||||
|
||||
return hidden_states
|
||||
@@ -0,0 +1,123 @@
|
||||
import math
|
||||
from typing import Tuple, Union
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import nn
|
||||
|
||||
from .convolution import make_conv_nd
|
||||
from .enums import PaddingModeType
|
||||
|
||||
|
||||
class SpaceToDepthDownsample(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dims: Union[int, Tuple[int, int]],
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
stride: Tuple[int, int, int],
|
||||
spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
|
||||
):
|
||||
super().__init__()
|
||||
self.stride = stride
|
||||
self.group_size = in_channels * math.prod(stride) // out_channels
|
||||
self.conv = make_conv_nd(
|
||||
dims=dims,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels // math.prod(stride),
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
causal=True,
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
causal: bool = True,
|
||||
) -> torch.Tensor:
|
||||
if self.stride[0] == 2:
|
||||
x = torch.cat([x[:, :, :1, :, :], x], dim=2) # duplicate first frames for padding
|
||||
|
||||
# skip connection
|
||||
x_in = rearrange(
|
||||
x,
|
||||
"b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w",
|
||||
p1=self.stride[0],
|
||||
p2=self.stride[1],
|
||||
p3=self.stride[2],
|
||||
)
|
||||
x_in = rearrange(x_in, "b (c g) d h w -> b c g d h w", g=self.group_size)
|
||||
x_in = x_in.mean(dim=2)
|
||||
|
||||
# conv
|
||||
x = self.conv(x, causal=causal)
|
||||
x = rearrange(
|
||||
x,
|
||||
"b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w",
|
||||
p1=self.stride[0],
|
||||
p2=self.stride[1],
|
||||
p3=self.stride[2],
|
||||
)
|
||||
|
||||
x = x + x_in
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class DepthToSpaceUpsample(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dims: int | Tuple[int, int],
|
||||
in_channels: int,
|
||||
stride: Tuple[int, int, int],
|
||||
residual: bool = False,
|
||||
out_channels_reduction_factor: int = 1,
|
||||
spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
|
||||
):
|
||||
super().__init__()
|
||||
self.stride = stride
|
||||
self.out_channels = math.prod(stride) * in_channels // out_channels_reduction_factor
|
||||
self.conv = make_conv_nd(
|
||||
dims=dims,
|
||||
in_channels=in_channels,
|
||||
out_channels=self.out_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
causal=True,
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
self.residual = residual
|
||||
self.out_channels_reduction_factor = out_channels_reduction_factor
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
causal: bool = True,
|
||||
) -> torch.Tensor:
|
||||
if self.residual:
|
||||
# Reshape and duplicate the input to match the output shape
|
||||
x_in = rearrange(
|
||||
x,
|
||||
"b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
|
||||
p1=self.stride[0],
|
||||
p2=self.stride[1],
|
||||
p3=self.stride[2],
|
||||
)
|
||||
num_repeat = math.prod(self.stride) // self.out_channels_reduction_factor
|
||||
x_in = x_in.repeat(1, num_repeat, 1, 1, 1)
|
||||
if self.stride[0] == 2:
|
||||
x_in = x_in[:, :, 1:, :, :]
|
||||
x = self.conv(x, causal=causal)
|
||||
x = rearrange(
|
||||
x,
|
||||
"b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
|
||||
p1=self.stride[0],
|
||||
p2=self.stride[1],
|
||||
p3=self.stride[2],
|
||||
)
|
||||
if self.stride[0] == 2:
|
||||
x = x[:, :, 1:, :, :]
|
||||
if self.residual:
|
||||
x = x + x_in
|
||||
return x
|
||||
@@ -0,0 +1,260 @@
|
||||
import itertools
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, List, NamedTuple, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def compute_trapezoidal_mask_1d(
|
||||
length: int,
|
||||
ramp_left: int,
|
||||
ramp_right: int,
|
||||
left_starts_from_0: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Generate a 1D trapezoidal blending mask with linear ramps.
|
||||
Args:
|
||||
length: Output length of the mask.
|
||||
ramp_left: Fade-in length on the left.
|
||||
ramp_right: Fade-out length on the right.
|
||||
left_starts_from_0: Whether the ramp starts from 0 or first non-zero value.
|
||||
Useful for temporal tiles where the first tile is causal.
|
||||
Returns:
|
||||
A 1D tensor of shape `(length,)` with values in [0, 1].
|
||||
"""
|
||||
if length <= 0:
|
||||
raise ValueError("Mask length must be positive.")
|
||||
|
||||
ramp_left = max(0, min(ramp_left, length))
|
||||
ramp_right = max(0, min(ramp_right, length))
|
||||
|
||||
mask = torch.ones(length)
|
||||
|
||||
if ramp_left > 0:
|
||||
interval_length = ramp_left + 1 if left_starts_from_0 else ramp_left + 2
|
||||
fade_in = torch.linspace(0.0, 1.0, interval_length)[:-1]
|
||||
if not left_starts_from_0:
|
||||
fade_in = fade_in[1:]
|
||||
mask[:ramp_left] *= fade_in
|
||||
|
||||
if ramp_right > 0:
|
||||
fade_out = torch.linspace(1.0, 0.0, steps=ramp_right + 2)[1:-1]
|
||||
mask[-ramp_right:] *= fade_out
|
||||
|
||||
return mask.clamp_(0, 1)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SpatialTilingConfig:
|
||||
"""Configuration for dividing each frame into spatial tiles with optional overlap.
|
||||
Args:
|
||||
tile_size_in_pixels (int): Size of each tile in pixels. Must be at least 64 and divisible by 32.
|
||||
tile_overlap_in_pixels (int, optional): Overlap between tiles in pixels. Must be divisible by 32. Defaults to 0.
|
||||
"""
|
||||
|
||||
tile_size_in_pixels: int
|
||||
tile_overlap_in_pixels: int = 0
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.tile_size_in_pixels < 64:
|
||||
raise ValueError(f"tile_size_in_pixels must be at least 64, got {self.tile_size_in_pixels}")
|
||||
if self.tile_size_in_pixels % 32 != 0:
|
||||
raise ValueError(f"tile_size_in_pixels must be divisible by 32, got {self.tile_size_in_pixels}")
|
||||
if self.tile_overlap_in_pixels % 32 != 0:
|
||||
raise ValueError(f"tile_overlap_in_pixels must be divisible by 32, got {self.tile_overlap_in_pixels}")
|
||||
if self.tile_overlap_in_pixels >= self.tile_size_in_pixels:
|
||||
raise ValueError(
|
||||
f"Overlap must be less than tile size, got {self.tile_overlap_in_pixels} and {self.tile_size_in_pixels}"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TemporalTilingConfig:
|
||||
"""Configuration for dividing a video into temporal tiles (chunks of frames) with optional overlap.
|
||||
Args:
|
||||
tile_size_in_frames (int): Number of frames in each tile. Must be at least 16 and divisible by 8.
|
||||
tile_overlap_in_frames (int, optional): Number of overlapping frames between consecutive tiles.
|
||||
Must be divisible by 8. Defaults to 0.
|
||||
"""
|
||||
|
||||
tile_size_in_frames: int
|
||||
tile_overlap_in_frames: int = 0
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.tile_size_in_frames < 16:
|
||||
raise ValueError(f"tile_size_in_frames must be at least 16, got {self.tile_size_in_frames}")
|
||||
if self.tile_size_in_frames % 8 != 0:
|
||||
raise ValueError(f"tile_size_in_frames must be divisible by 8, got {self.tile_size_in_frames}")
|
||||
if self.tile_overlap_in_frames % 8 != 0:
|
||||
raise ValueError(f"tile_overlap_in_frames must be divisible by 8, got {self.tile_overlap_in_frames}")
|
||||
if self.tile_overlap_in_frames >= self.tile_size_in_frames:
|
||||
raise ValueError(
|
||||
f"Overlap must be less than tile size, got {self.tile_overlap_in_frames} and {self.tile_size_in_frames}"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TilingConfig:
|
||||
"""Configuration for splitting video into tiles with optional overlap.
|
||||
Attributes:
|
||||
spatial_config: Configuration for splitting spatial dimensions into tiles.
|
||||
temporal_config: Configuration for splitting temporal dimension into tiles.
|
||||
"""
|
||||
|
||||
spatial_config: SpatialTilingConfig | None = None
|
||||
temporal_config: TemporalTilingConfig | None = None
|
||||
|
||||
@classmethod
|
||||
def default(cls) -> "TilingConfig":
|
||||
return cls(
|
||||
spatial_config=SpatialTilingConfig(tile_size_in_pixels=512, tile_overlap_in_pixels=64),
|
||||
temporal_config=TemporalTilingConfig(tile_size_in_frames=64, tile_overlap_in_frames=24),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DimensionIntervals:
|
||||
"""Intervals which a single dimension of the latent space is split into.
|
||||
Each interval is defined by its start, end, left ramp, and right ramp.
|
||||
The start and end are the indices of the first and last element (exclusive) in the interval.
|
||||
Ramps are regions of the interval where the value of the mask tensor is
|
||||
interpolated between 0 and 1 for blending with neighboring intervals.
|
||||
The left ramp and right ramp values are the lengths of the left and right ramps.
|
||||
"""
|
||||
|
||||
starts: List[int]
|
||||
ends: List[int]
|
||||
left_ramps: List[int]
|
||||
right_ramps: List[int]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LatentIntervals:
|
||||
"""Intervals which the latent tensor of given shape is split into.
|
||||
Each dimension of the latent space is split into intervals based on the length along said dimension.
|
||||
"""
|
||||
|
||||
original_shape: torch.Size
|
||||
dimension_intervals: Tuple[DimensionIntervals, ...]
|
||||
|
||||
|
||||
# Operation to split a single dimension of the tensor into intervals based on the length along the dimension.
|
||||
SplitOperation = Callable[[int], DimensionIntervals]
|
||||
# Operation to map the intervals in input dimension to slices and masks along a corresponding output dimension.
|
||||
MappingOperation = Callable[[DimensionIntervals], tuple[list[slice], list[torch.Tensor | None]]]
|
||||
|
||||
|
||||
def default_split_operation(length: int) -> DimensionIntervals:
|
||||
return DimensionIntervals(starts=[0], ends=[length], left_ramps=[0], right_ramps=[0])
|
||||
|
||||
|
||||
DEFAULT_SPLIT_OPERATION: SplitOperation = default_split_operation
|
||||
|
||||
|
||||
def default_mapping_operation(
|
||||
_intervals: DimensionIntervals,
|
||||
) -> tuple[list[slice], list[torch.Tensor | None]]:
|
||||
return [slice(0, None)], [None]
|
||||
|
||||
|
||||
DEFAULT_MAPPING_OPERATION: MappingOperation = default_mapping_operation
|
||||
|
||||
|
||||
class Tile(NamedTuple):
|
||||
"""
|
||||
Represents a single tile.
|
||||
Attributes:
|
||||
in_coords:
|
||||
Tuple of slices specifying where to cut the tile from the INPUT tensor.
|
||||
out_coords:
|
||||
Tuple of slices specifying where this tile's OUTPUT should be placed in the reconstructed OUTPUT tensor.
|
||||
masks_1d:
|
||||
Per-dimension masks in OUTPUT units.
|
||||
These are used to create all-dimensional blending mask.
|
||||
Methods:
|
||||
blend_mask:
|
||||
Create a single N-D mask from the per-dimension masks.
|
||||
"""
|
||||
|
||||
in_coords: Tuple[slice, ...]
|
||||
out_coords: Tuple[slice, ...]
|
||||
masks_1d: Tuple[Tuple[torch.Tensor, ...]]
|
||||
|
||||
@property
|
||||
def blend_mask(self) -> torch.Tensor:
|
||||
num_dims = len(self.out_coords)
|
||||
per_dimension_masks: List[torch.Tensor] = []
|
||||
|
||||
for dim_idx in range(num_dims):
|
||||
mask_1d = self.masks_1d[dim_idx]
|
||||
view_shape = [1] * num_dims
|
||||
if mask_1d is None:
|
||||
# Broadcast mask along this dimension (length 1).
|
||||
one = torch.ones(1)
|
||||
|
||||
view_shape[dim_idx] = 1
|
||||
per_dimension_masks.append(one.view(*view_shape))
|
||||
continue
|
||||
|
||||
# Reshape (L,) -> (1, ..., L, ..., 1) so masks across dimensions broadcast-multiply.
|
||||
view_shape[dim_idx] = mask_1d.shape[0]
|
||||
per_dimension_masks.append(mask_1d.view(*view_shape))
|
||||
|
||||
# Multiply per-dimension masks to form the full N-D mask (separable blending window).
|
||||
combined_mask = per_dimension_masks[0]
|
||||
for mask in per_dimension_masks[1:]:
|
||||
combined_mask = combined_mask * mask
|
||||
|
||||
return combined_mask
|
||||
|
||||
|
||||
def create_tiles_from_intervals_and_mappers(
|
||||
intervals: LatentIntervals,
|
||||
mappers: List[MappingOperation],
|
||||
) -> List[Tile]:
|
||||
full_dim_input_slices = []
|
||||
full_dim_output_slices = []
|
||||
full_dim_masks_1d = []
|
||||
for axis_index in range(len(intervals.original_shape)):
|
||||
dimension_intervals = intervals.dimension_intervals[axis_index]
|
||||
starts = dimension_intervals.starts
|
||||
ends = dimension_intervals.ends
|
||||
input_slices = [slice(s, e) for s, e in zip(starts, ends, strict=True)]
|
||||
output_slices, masks_1d = mappers[axis_index](dimension_intervals)
|
||||
full_dim_input_slices.append(input_slices)
|
||||
full_dim_output_slices.append(output_slices)
|
||||
full_dim_masks_1d.append(masks_1d)
|
||||
|
||||
tiles = []
|
||||
tile_in_coords = list(itertools.product(*full_dim_input_slices))
|
||||
tile_out_coords = list(itertools.product(*full_dim_output_slices))
|
||||
tile_mask_1ds = list(itertools.product(*full_dim_masks_1d))
|
||||
for in_coord, out_coord, mask_1d in zip(tile_in_coords, tile_out_coords, tile_mask_1ds, strict=True):
|
||||
tiles.append(
|
||||
Tile(
|
||||
in_coords=in_coord,
|
||||
out_coords=out_coord,
|
||||
masks_1d=mask_1d,
|
||||
)
|
||||
)
|
||||
return tiles
|
||||
|
||||
|
||||
def create_tiles(
|
||||
latent_shape: torch.Size,
|
||||
splitters: List[SplitOperation],
|
||||
mappers: List[MappingOperation],
|
||||
) -> List[Tile]:
|
||||
if len(splitters) != len(latent_shape):
|
||||
raise ValueError(
|
||||
f"Number of splitters must be equal to number of dimensions in latent shape, "
|
||||
f"got {len(splitters)} and {len(latent_shape)}"
|
||||
)
|
||||
if len(mappers) != len(latent_shape):
|
||||
raise ValueError(
|
||||
f"Number of mappers must be equal to number of dimensions in latent shape, "
|
||||
f"got {len(mappers)} and {len(latent_shape)}"
|
||||
)
|
||||
intervals = [splitter(length) for splitter, length in zip(splitters, latent_shape, strict=True)]
|
||||
latent_intervals = LatentIntervals(original_shape=latent_shape, dimension_intervals=tuple(intervals))
|
||||
return create_tiles_from_intervals_and_mappers(latent_intervals, mappers)
|
||||
@@ -0,0 +1,925 @@
|
||||
from dataclasses import replace
|
||||
from typing import Any, Callable, Iterator, List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import nn
|
||||
|
||||
from ltx_core.model.common.normalization import PixelNorm
|
||||
from ltx_core.model.transformer.timestep_embedding import PixArtAlphaCombinedTimestepSizeEmbeddings
|
||||
from ltx_core.model.video_vae.convolution import make_conv_nd
|
||||
from ltx_core.model.video_vae.enums import LogVarianceType, NormLayerType, PaddingModeType
|
||||
from ltx_core.model.video_vae.ops import PerChannelStatistics, patchify, unpatchify
|
||||
from ltx_core.model.video_vae.resnet import ResnetBlock3D, UNetMidBlock3D
|
||||
from ltx_core.model.video_vae.sampling import DepthToSpaceUpsample, SpaceToDepthDownsample
|
||||
from ltx_core.model.video_vae.tiling import (
|
||||
DEFAULT_MAPPING_OPERATION,
|
||||
DEFAULT_SPLIT_OPERATION,
|
||||
DimensionIntervals,
|
||||
MappingOperation,
|
||||
SplitOperation,
|
||||
Tile,
|
||||
TilingConfig,
|
||||
compute_trapezoidal_mask_1d,
|
||||
create_tiles,
|
||||
)
|
||||
from ltx_core.types import SpatioTemporalScaleFactors, VideoLatentShape
|
||||
|
||||
|
||||
def _make_encoder_block(
|
||||
block_name: str,
|
||||
block_config: dict[str, Any],
|
||||
in_channels: int,
|
||||
convolution_dimensions: int,
|
||||
norm_layer: NormLayerType,
|
||||
norm_num_groups: int,
|
||||
spatial_padding_mode: PaddingModeType,
|
||||
) -> Tuple[nn.Module, int]:
|
||||
out_channels = in_channels
|
||||
|
||||
if block_name == "res_x":
|
||||
block = UNetMidBlock3D(
|
||||
dims=convolution_dimensions,
|
||||
in_channels=in_channels,
|
||||
num_layers=block_config["num_layers"],
|
||||
resnet_eps=1e-6,
|
||||
resnet_groups=norm_num_groups,
|
||||
norm_layer=norm_layer,
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
elif block_name == "res_x_y":
|
||||
out_channels = in_channels * block_config.get("multiplier", 2)
|
||||
block = ResnetBlock3D(
|
||||
dims=convolution_dimensions,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
eps=1e-6,
|
||||
groups=norm_num_groups,
|
||||
norm_layer=norm_layer,
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
elif block_name == "compress_time":
|
||||
block = make_conv_nd(
|
||||
dims=convolution_dimensions,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=3,
|
||||
stride=(2, 1, 1),
|
||||
causal=True,
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
elif block_name == "compress_space":
|
||||
block = make_conv_nd(
|
||||
dims=convolution_dimensions,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=3,
|
||||
stride=(1, 2, 2),
|
||||
causal=True,
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
elif block_name == "compress_all":
|
||||
block = make_conv_nd(
|
||||
dims=convolution_dimensions,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=3,
|
||||
stride=(2, 2, 2),
|
||||
causal=True,
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
elif block_name == "compress_all_x_y":
|
||||
out_channels = in_channels * block_config.get("multiplier", 2)
|
||||
block = make_conv_nd(
|
||||
dims=convolution_dimensions,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=3,
|
||||
stride=(2, 2, 2),
|
||||
causal=True,
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
elif block_name == "compress_all_res":
|
||||
out_channels = in_channels * block_config.get("multiplier", 2)
|
||||
block = SpaceToDepthDownsample(
|
||||
dims=convolution_dimensions,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
stride=(2, 2, 2),
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
elif block_name == "compress_space_res":
|
||||
out_channels = in_channels * block_config.get("multiplier", 2)
|
||||
block = SpaceToDepthDownsample(
|
||||
dims=convolution_dimensions,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
stride=(1, 2, 2),
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
elif block_name == "compress_time_res":
|
||||
out_channels = in_channels * block_config.get("multiplier", 2)
|
||||
block = SpaceToDepthDownsample(
|
||||
dims=convolution_dimensions,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
stride=(2, 1, 1),
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"unknown block: {block_name}")
|
||||
|
||||
return block, out_channels
|
||||
|
||||
|
||||
class VideoEncoder(nn.Module):
|
||||
_DEFAULT_NORM_NUM_GROUPS = 32
|
||||
"""
|
||||
Variational Autoencoder Encoder. Encodes video frames into a latent representation.
|
||||
The encoder compresses the input video through a series of downsampling operations controlled by
|
||||
patch_size and encoder_blocks. The output is a normalized latent tensor with shape (B, 128, F', H', W').
|
||||
Compression Behavior:
|
||||
The total compression is determined by:
|
||||
1. Initial spatial compression via patchify: H -> H/4, W -> W/4 (patch_size=4)
|
||||
2. Sequential compression through encoder_blocks based on their stride patterns
|
||||
Compression blocks apply 2x compression in specified dimensions:
|
||||
- "compress_time" / "compress_time_res": temporal only
|
||||
- "compress_space" / "compress_space_res": spatial only (H and W)
|
||||
- "compress_all" / "compress_all_res": all dimensions (F, H, W)
|
||||
- "res_x" / "res_x_y": no compression
|
||||
Standard LTX Video configuration:
|
||||
- patch_size=4
|
||||
- encoder_blocks: 1x compress_space_res, 1x compress_time_res, 2x compress_all_res
|
||||
- Final dimensions: F' = 1 + (F-1)/8, H' = H/32, W' = W/32
|
||||
- Example: (B, 3, 33, 512, 512) -> (B, 128, 5, 16, 16)
|
||||
- Note: Input must have 1 + 8*k frames (e.g., 1, 9, 17, 25, 33...)
|
||||
Args:
|
||||
convolution_dimensions: The number of dimensions to use in convolutions (2D or 3D).
|
||||
in_channels: The number of input channels. For RGB images, this is 3.
|
||||
out_channels: The number of output channels (latent channels). For latent channels, this is 128.
|
||||
encoder_blocks: The list of blocks to construct the encoder. Each block is a tuple of (block_name, params)
|
||||
where params is either an int (num_layers) or a dict with configuration.
|
||||
patch_size: The patch size for initial spatial compression. Should be a power of 2.
|
||||
norm_layer: The normalization layer to use. Can be either `group_norm` or `pixel_norm`.
|
||||
latent_log_var: The log variance mode. Can be either `per_channel`, `uniform`, `constant` or `none`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
convolution_dimensions: int = 3,
|
||||
in_channels: int = 3,
|
||||
out_channels: int = 128,
|
||||
encoder_blocks: List[Tuple[str, int]] | List[Tuple[str, dict[str, Any]]] = [], # noqa: B006
|
||||
patch_size: int = 4,
|
||||
norm_layer: NormLayerType = NormLayerType.PIXEL_NORM,
|
||||
latent_log_var: LogVarianceType = LogVarianceType.UNIFORM,
|
||||
encoder_spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.patch_size = patch_size
|
||||
self.norm_layer = norm_layer
|
||||
self.latent_channels = out_channels
|
||||
self.latent_log_var = latent_log_var
|
||||
self._norm_num_groups = self._DEFAULT_NORM_NUM_GROUPS
|
||||
|
||||
# Per-channel statistics for normalizing latents
|
||||
self.per_channel_statistics = PerChannelStatistics(latent_channels=out_channels)
|
||||
|
||||
in_channels = in_channels * patch_size**2
|
||||
feature_channels = out_channels
|
||||
|
||||
self.conv_in = make_conv_nd(
|
||||
dims=convolution_dimensions,
|
||||
in_channels=in_channels,
|
||||
out_channels=feature_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
causal=True,
|
||||
spatial_padding_mode=encoder_spatial_padding_mode,
|
||||
)
|
||||
|
||||
self.down_blocks = nn.ModuleList([])
|
||||
|
||||
for block_name, block_params in encoder_blocks:
|
||||
# Convert int to dict format for uniform handling
|
||||
block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params
|
||||
|
||||
block, feature_channels = _make_encoder_block(
|
||||
block_name=block_name,
|
||||
block_config=block_config,
|
||||
in_channels=feature_channels,
|
||||
convolution_dimensions=convolution_dimensions,
|
||||
norm_layer=norm_layer,
|
||||
norm_num_groups=self._norm_num_groups,
|
||||
spatial_padding_mode=encoder_spatial_padding_mode,
|
||||
)
|
||||
|
||||
self.down_blocks.append(block)
|
||||
|
||||
# out
|
||||
if norm_layer == NormLayerType.GROUP_NORM:
|
||||
self.conv_norm_out = nn.GroupNorm(num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6)
|
||||
elif norm_layer == NormLayerType.PIXEL_NORM:
|
||||
self.conv_norm_out = PixelNorm()
|
||||
|
||||
self.conv_act = nn.SiLU()
|
||||
|
||||
conv_out_channels = out_channels
|
||||
if latent_log_var == LogVarianceType.PER_CHANNEL:
|
||||
conv_out_channels *= 2
|
||||
elif latent_log_var in {LogVarianceType.UNIFORM, LogVarianceType.CONSTANT}:
|
||||
conv_out_channels += 1
|
||||
elif latent_log_var != LogVarianceType.NONE:
|
||||
raise ValueError(f"Invalid latent_log_var: {latent_log_var}")
|
||||
|
||||
self.conv_out = make_conv_nd(
|
||||
dims=convolution_dimensions,
|
||||
in_channels=feature_channels,
|
||||
out_channels=conv_out_channels,
|
||||
kernel_size=3,
|
||||
padding=1,
|
||||
causal=True,
|
||||
spatial_padding_mode=encoder_spatial_padding_mode,
|
||||
)
|
||||
|
||||
def forward(self, sample: torch.Tensor) -> torch.Tensor:
|
||||
r"""
|
||||
Encode video frames into normalized latent representation.
|
||||
Args:
|
||||
sample: Input video (B, C, F, H, W). F must be 1 + 8*k (e.g., 1, 9, 17, 25, 33...).
|
||||
Returns:
|
||||
Normalized latent means (B, 128, F', H', W') where F' = 1+(F-1)/8, H' = H/32, W' = W/32.
|
||||
Example: (B, 3, 33, 512, 512) -> (B, 128, 5, 16, 16).
|
||||
"""
|
||||
# Validate frame count
|
||||
frames_count = sample.shape[2]
|
||||
if ((frames_count - 1) % 8) != 0:
|
||||
raise ValueError(
|
||||
"Invalid number of frames: Encode input must have 1 + 8 * x frames "
|
||||
"(e.g., 1, 9, 17, ...). Please check your input."
|
||||
)
|
||||
|
||||
# Initial spatial compression: trade spatial resolution for channel depth
|
||||
# This reduces H,W by patch_size and increases channels, making convolutions more efficient
|
||||
# Example: (B, 3, F, 512, 512) -> (B, 48, F, 128, 128) with patch_size=4
|
||||
sample = patchify(sample, patch_size_hw=self.patch_size, patch_size_t=1)
|
||||
sample = self.conv_in(sample)
|
||||
|
||||
for down_block in self.down_blocks:
|
||||
sample = down_block(sample)
|
||||
|
||||
sample = self.conv_norm_out(sample)
|
||||
sample = self.conv_act(sample)
|
||||
sample = self.conv_out(sample)
|
||||
|
||||
if self.latent_log_var == LogVarianceType.UNIFORM:
|
||||
# Uniform Variance: model outputs N means and 1 shared log-variance channel.
|
||||
# We need to expand the single logvar to match the number of means channels
|
||||
# to create a format compatible with PER_CHANNEL (means + logvar, each with N channels).
|
||||
# Sample shape: (B, N+1, ...) where N = latent_channels (e.g., 128 means + 1 logvar = 129)
|
||||
# Target shape: (B, 2*N, ...) where first N are means, last N are logvar
|
||||
|
||||
if sample.shape[1] < 2:
|
||||
raise ValueError(
|
||||
f"Invalid channel count for UNIFORM mode: expected at least 2 channels "
|
||||
f"(N means + 1 logvar), got {sample.shape[1]}"
|
||||
)
|
||||
|
||||
# Extract means (first N channels) and logvar (last 1 channel)
|
||||
means = sample[:, :-1, ...] # (B, N, ...)
|
||||
logvar = sample[:, -1:, ...] # (B, 1, ...)
|
||||
|
||||
# Repeat logvar N times to match means channels
|
||||
# Use expand/repeat pattern that works for both 4D and 5D tensors
|
||||
num_channels = means.shape[1]
|
||||
repeat_shape = [1, num_channels] + [1] * (sample.ndim - 2)
|
||||
repeated_logvar = logvar.repeat(*repeat_shape) # (B, N, ...)
|
||||
|
||||
# Concatenate to create (B, 2*N, ...) format: [means, repeated_logvar]
|
||||
sample = torch.cat([means, repeated_logvar], dim=1)
|
||||
elif self.latent_log_var == LogVarianceType.CONSTANT:
|
||||
sample = sample[:, :-1, ...]
|
||||
approx_ln_0 = -30 # this is the minimal clamp value in DiagonalGaussianDistribution objects
|
||||
sample = torch.cat(
|
||||
[sample, torch.ones_like(sample, device=sample.device) * approx_ln_0],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
# Split into means and logvar, then normalize means
|
||||
means, _ = torch.chunk(sample, 2, dim=1)
|
||||
return self.per_channel_statistics.normalize(means)
|
||||
|
||||
|
||||
def _make_decoder_block(
|
||||
block_name: str,
|
||||
block_config: dict[str, Any],
|
||||
in_channels: int,
|
||||
convolution_dimensions: int,
|
||||
norm_layer: NormLayerType,
|
||||
timestep_conditioning: bool,
|
||||
norm_num_groups: int,
|
||||
spatial_padding_mode: PaddingModeType,
|
||||
) -> Tuple[nn.Module, int]:
|
||||
out_channels = in_channels
|
||||
if block_name == "res_x":
|
||||
block = UNetMidBlock3D(
|
||||
dims=convolution_dimensions,
|
||||
in_channels=in_channels,
|
||||
num_layers=block_config["num_layers"],
|
||||
resnet_eps=1e-6,
|
||||
resnet_groups=norm_num_groups,
|
||||
norm_layer=norm_layer,
|
||||
inject_noise=block_config.get("inject_noise", False),
|
||||
timestep_conditioning=timestep_conditioning,
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
elif block_name == "attn_res_x":
|
||||
block = UNetMidBlock3D(
|
||||
dims=convolution_dimensions,
|
||||
in_channels=in_channels,
|
||||
num_layers=block_config["num_layers"],
|
||||
resnet_groups=norm_num_groups,
|
||||
norm_layer=norm_layer,
|
||||
inject_noise=block_config.get("inject_noise", False),
|
||||
timestep_conditioning=timestep_conditioning,
|
||||
attention_head_dim=block_config["attention_head_dim"],
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
elif block_name == "res_x_y":
|
||||
out_channels = in_channels // block_config.get("multiplier", 2)
|
||||
block = ResnetBlock3D(
|
||||
dims=convolution_dimensions,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
eps=1e-6,
|
||||
groups=norm_num_groups,
|
||||
norm_layer=norm_layer,
|
||||
inject_noise=block_config.get("inject_noise", False),
|
||||
timestep_conditioning=False,
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
elif block_name == "compress_time":
|
||||
block = DepthToSpaceUpsample(
|
||||
dims=convolution_dimensions,
|
||||
in_channels=in_channels,
|
||||
stride=(2, 1, 1),
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
elif block_name == "compress_space":
|
||||
block = DepthToSpaceUpsample(
|
||||
dims=convolution_dimensions,
|
||||
in_channels=in_channels,
|
||||
stride=(1, 2, 2),
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
elif block_name == "compress_all":
|
||||
out_channels = in_channels // block_config.get("multiplier", 1)
|
||||
block = DepthToSpaceUpsample(
|
||||
dims=convolution_dimensions,
|
||||
in_channels=in_channels,
|
||||
stride=(2, 2, 2),
|
||||
residual=block_config.get("residual", False),
|
||||
out_channels_reduction_factor=block_config.get("multiplier", 1),
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"unknown layer: {block_name}")
|
||||
|
||||
return block, out_channels
|
||||
|
||||
|
||||
class VideoDecoder(nn.Module):
|
||||
_DEFAULT_NORM_NUM_GROUPS = 32
|
||||
"""
|
||||
Variational Autoencoder Decoder. Decodes latent representation into video frames.
|
||||
The decoder upsamples latents through a series of upsampling operations (inverse of encoder).
|
||||
Output dimensions: F = 8x(F'-1) + 1, H = 32xH', W = 32xW' for standard LTX Video configuration.
|
||||
Upsampling blocks expand dimensions by 2x in specified dimensions:
|
||||
- "compress_time": temporal only
|
||||
- "compress_space": spatial only (H and W)
|
||||
- "compress_all": all dimensions (F, H, W)
|
||||
- "res_x" / "res_x_y" / "attn_res_x": no upsampling
|
||||
Causal Mode:
|
||||
causal=False (standard): Symmetric padding, allows future frame dependencies.
|
||||
causal=True: Causal padding, each frame depends only on past/current frames.
|
||||
First frame removed after temporal upsampling in both modes. Output shape unchanged.
|
||||
Example: (B, 128, 5, 16, 16) -> (B, 3, 33, 512, 512) for both modes.
|
||||
Args:
|
||||
convolution_dimensions: The number of dimensions to use in convolutions (2D or 3D).
|
||||
in_channels: The number of input channels (latent channels). Default is 128.
|
||||
out_channels: The number of output channels. For RGB images, this is 3.
|
||||
decoder_blocks: The list of blocks to construct the decoder. Each block is a tuple of (block_name, params)
|
||||
where params is either an int (num_layers) or a dict with configuration.
|
||||
patch_size: Final spatial expansion factor. For standard LTX Video, use 4 for 4x spatial expansion:
|
||||
H -> Hx4, W -> Wx4. Should be a power of 2.
|
||||
norm_layer: The normalization layer to use. Can be either `group_norm` or `pixel_norm`.
|
||||
causal: Whether to use causal convolutions. For standard LTX Video, use False for symmetric padding.
|
||||
When True, uses causal padding (past/current frames only).
|
||||
timestep_conditioning: Whether to condition the decoder on timestep for denoising.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
convolution_dimensions: int = 3,
|
||||
in_channels: int = 128,
|
||||
out_channels: int = 3,
|
||||
decoder_blocks: List[Tuple[str, int | dict]] = [], # noqa: B006
|
||||
patch_size: int = 4,
|
||||
norm_layer: NormLayerType = NormLayerType.PIXEL_NORM,
|
||||
causal: bool = False,
|
||||
timestep_conditioning: bool = False,
|
||||
decoder_spatial_padding_mode: PaddingModeType = PaddingModeType.REFLECT,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# Spatiotemporal downscaling between decoded video space and VAE latents.
|
||||
# According to the LTXV paper, the standard configuration downsamples
|
||||
# video inputs by a factor of 8 in the temporal dimension and 32 in
|
||||
# each spatial dimension (height and width). This parameter determines how
|
||||
# many video frames and pixels correspond to a single latent cell.
|
||||
self.video_downscale_factors = SpatioTemporalScaleFactors(
|
||||
time=8,
|
||||
width=32,
|
||||
height=32,
|
||||
)
|
||||
|
||||
self.patch_size = patch_size
|
||||
out_channels = out_channels * patch_size**2
|
||||
self.causal = causal
|
||||
self.timestep_conditioning = timestep_conditioning
|
||||
self._norm_num_groups = self._DEFAULT_NORM_NUM_GROUPS
|
||||
|
||||
# Per-channel statistics for denormalizing latents
|
||||
self.per_channel_statistics = PerChannelStatistics(latent_channels=in_channels)
|
||||
|
||||
# Noise and timestep parameters for decoder conditioning
|
||||
self.decode_noise_scale = 0.025
|
||||
self.decode_timestep = 0.05
|
||||
|
||||
# Compute initial feature_channels by going through blocks in reverse
|
||||
# This determines the channel width at the start of the decoder
|
||||
feature_channels = in_channels
|
||||
for block_name, block_params in list(reversed(decoder_blocks)):
|
||||
block_config = block_params if isinstance(block_params, dict) else {}
|
||||
if block_name == "res_x_y":
|
||||
feature_channels = feature_channels * block_config.get("multiplier", 2)
|
||||
if block_name == "compress_all":
|
||||
feature_channels = feature_channels * block_config.get("multiplier", 1)
|
||||
|
||||
self.conv_in = make_conv_nd(
|
||||
dims=convolution_dimensions,
|
||||
in_channels=in_channels,
|
||||
out_channels=feature_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
causal=True,
|
||||
spatial_padding_mode=decoder_spatial_padding_mode,
|
||||
)
|
||||
|
||||
self.up_blocks = nn.ModuleList([])
|
||||
|
||||
for block_name, block_params in list(reversed(decoder_blocks)):
|
||||
# Convert int to dict format for uniform handling
|
||||
block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params
|
||||
|
||||
block, feature_channels = _make_decoder_block(
|
||||
block_name=block_name,
|
||||
block_config=block_config,
|
||||
in_channels=feature_channels,
|
||||
convolution_dimensions=convolution_dimensions,
|
||||
norm_layer=norm_layer,
|
||||
timestep_conditioning=timestep_conditioning,
|
||||
norm_num_groups=self._norm_num_groups,
|
||||
spatial_padding_mode=decoder_spatial_padding_mode,
|
||||
)
|
||||
|
||||
self.up_blocks.append(block)
|
||||
|
||||
if norm_layer == NormLayerType.GROUP_NORM:
|
||||
self.conv_norm_out = nn.GroupNorm(num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6)
|
||||
elif norm_layer == NormLayerType.PIXEL_NORM:
|
||||
self.conv_norm_out = PixelNorm()
|
||||
|
||||
self.conv_act = nn.SiLU()
|
||||
self.conv_out = make_conv_nd(
|
||||
dims=convolution_dimensions,
|
||||
in_channels=feature_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=3,
|
||||
padding=1,
|
||||
causal=True,
|
||||
spatial_padding_mode=decoder_spatial_padding_mode,
|
||||
)
|
||||
|
||||
if timestep_conditioning:
|
||||
self.timestep_scale_multiplier = nn.Parameter(torch.tensor(1000.0))
|
||||
self.last_time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(
|
||||
embedding_dim=feature_channels * 2, size_emb_dim=0
|
||||
)
|
||||
self.last_scale_shift_table = nn.Parameter(torch.empty(2, feature_channels))
|
||||
|
||||
# def forward(self, sample: torch.Tensor, target_shape) -> torch.Tensor:
|
||||
def forward(
|
||||
self,
|
||||
sample: torch.Tensor,
|
||||
timestep: Optional[torch.Tensor] = None,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
) -> torch.Tensor:
|
||||
r"""
|
||||
Decode latent representation into video frames.
|
||||
Args:
|
||||
sample: Latent tensor (B, 128, F', H', W').
|
||||
timestep: Timestep for conditioning (if timestep_conditioning=True). Uses default 0.05 if None.
|
||||
generator: Random generator for deterministic noise injection (if inject_noise=True in blocks).
|
||||
Returns:
|
||||
Decoded video (B, 3, F, H, W) where F = 8x(F'-1) + 1, H = 32xH', W = 32xW'.
|
||||
Example: (B, 128, 5, 16, 16) -> (B, 3, 33, 512, 512).
|
||||
Note: First frame is removed after temporal upsampling regardless of causal mode.
|
||||
When causal=False, allows future frame dependencies in convolutions but maintains same output shape.
|
||||
"""
|
||||
batch_size = sample.shape[0]
|
||||
|
||||
# Add noise if timestep conditioning is enabled
|
||||
if self.timestep_conditioning:
|
||||
noise = (
|
||||
torch.randn(
|
||||
sample.size(),
|
||||
generator=generator,
|
||||
dtype=sample.dtype,
|
||||
device=sample.device,
|
||||
)
|
||||
* self.decode_noise_scale
|
||||
)
|
||||
|
||||
sample = noise + (1.0 - self.decode_noise_scale) * sample
|
||||
|
||||
# Denormalize latents
|
||||
sample = self.per_channel_statistics.un_normalize(sample)
|
||||
|
||||
# Use default decode_timestep if timestep not provided
|
||||
if timestep is None and self.timestep_conditioning:
|
||||
timestep = torch.full((batch_size,), self.decode_timestep, device=sample.device, dtype=sample.dtype)
|
||||
|
||||
sample = self.conv_in(sample, causal=self.causal)
|
||||
|
||||
scaled_timestep = None
|
||||
if self.timestep_conditioning:
|
||||
if timestep is None:
|
||||
raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True")
|
||||
scaled_timestep = timestep * self.timestep_scale_multiplier.to(sample)
|
||||
|
||||
for up_block in self.up_blocks:
|
||||
if isinstance(up_block, UNetMidBlock3D):
|
||||
block_kwargs = {
|
||||
"causal": self.causal,
|
||||
"timestep": scaled_timestep if self.timestep_conditioning else None,
|
||||
"generator": generator,
|
||||
}
|
||||
sample = up_block(sample, **block_kwargs)
|
||||
elif isinstance(up_block, ResnetBlock3D):
|
||||
sample = up_block(sample, causal=self.causal, generator=generator)
|
||||
else:
|
||||
sample = up_block(sample, causal=self.causal)
|
||||
|
||||
sample = self.conv_norm_out(sample)
|
||||
|
||||
if self.timestep_conditioning:
|
||||
embedded_timestep = self.last_time_embedder(
|
||||
timestep=scaled_timestep.flatten(),
|
||||
hidden_dtype=sample.dtype,
|
||||
)
|
||||
embedded_timestep = embedded_timestep.view(batch_size, embedded_timestep.shape[-1], 1, 1, 1)
|
||||
ada_values = self.last_scale_shift_table[None, ..., None, None, None].to(
|
||||
device=sample.device, dtype=sample.dtype
|
||||
) + embedded_timestep.reshape(
|
||||
batch_size,
|
||||
2,
|
||||
-1,
|
||||
embedded_timestep.shape[-3],
|
||||
embedded_timestep.shape[-2],
|
||||
embedded_timestep.shape[-1],
|
||||
)
|
||||
shift, scale = ada_values.unbind(dim=1)
|
||||
sample = sample * (1 + scale) + shift
|
||||
|
||||
sample = self.conv_act(sample)
|
||||
sample = self.conv_out(sample, causal=self.causal)
|
||||
|
||||
# Final spatial expansion: reverse the initial patchify from encoder
|
||||
# Moves pixels from channels back to spatial dimensions
|
||||
# Example: (B, 48, F, 128, 128) -> (B, 3, F, 512, 512) with patch_size=4
|
||||
sample = unpatchify(sample, patch_size_hw=self.patch_size, patch_size_t=1)
|
||||
|
||||
return sample
|
||||
|
||||
def _prepare_tiles(
|
||||
self,
|
||||
latent: torch.Tensor,
|
||||
tiling_config: TilingConfig | None = None,
|
||||
) -> List[Tile]:
|
||||
splitters = [DEFAULT_SPLIT_OPERATION] * len(latent.shape)
|
||||
mappers = [DEFAULT_MAPPING_OPERATION] * len(latent.shape)
|
||||
if tiling_config is not None and tiling_config.spatial_config is not None:
|
||||
cfg = tiling_config.spatial_config
|
||||
long_side = max(latent.shape[3], latent.shape[4])
|
||||
|
||||
def enable_on_axis(axis_idx: int, factor: int) -> None:
|
||||
size = cfg.tile_size_in_pixels // factor
|
||||
overlap = cfg.tile_overlap_in_pixels // factor
|
||||
axis_length = latent.shape[axis_idx]
|
||||
lower_threshold = max(2, overlap + 1)
|
||||
tile_size = max(lower_threshold, round(size * axis_length / long_side))
|
||||
splitters[axis_idx] = split_in_spatial(tile_size, overlap)
|
||||
mappers[axis_idx] = to_mapping_operation(map_spatial_slice, factor)
|
||||
|
||||
enable_on_axis(3, self.video_downscale_factors.height)
|
||||
enable_on_axis(4, self.video_downscale_factors.width)
|
||||
|
||||
if tiling_config is not None and tiling_config.temporal_config is not None:
|
||||
cfg = tiling_config.temporal_config
|
||||
tile_size = cfg.tile_size_in_frames // self.video_downscale_factors.time
|
||||
overlap = cfg.tile_overlap_in_frames // self.video_downscale_factors.time
|
||||
splitters[2] = split_in_temporal(tile_size, overlap)
|
||||
mappers[2] = to_mapping_operation(map_temporal_slice, self.video_downscale_factors.time)
|
||||
|
||||
return create_tiles(latent.shape, splitters, mappers)
|
||||
|
||||
def tiled_decode(
|
||||
self,
|
||||
latent: torch.Tensor,
|
||||
tiling_config: TilingConfig | None = None,
|
||||
timestep: Optional[torch.Tensor] = None,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
) -> Iterator[torch.Tensor]:
|
||||
"""
|
||||
Decode a latent tensor into video frames using tiled processing.
|
||||
Splits the latent tensor into tiles, decodes each tile individually,
|
||||
and yields video chunks as they become available.
|
||||
Args:
|
||||
latent: Input latent tensor (B, C, F', H', W').
|
||||
tiling_config: Tiling configuration for the latent tensor.
|
||||
timestep: Optional timestep for decoder conditioning.
|
||||
generator: Optional random generator for deterministic decoding.
|
||||
Yields:
|
||||
Video chunks (B, C, T, H, W) by temporal slices;
|
||||
"""
|
||||
|
||||
# Calculate full video shape from latent shape to get spatial dimensions
|
||||
full_video_shape = VideoLatentShape.from_torch_shape(latent.shape).upscale(self.video_downscale_factors)
|
||||
tiles = self._prepare_tiles(latent, tiling_config)
|
||||
|
||||
temporal_groups = self._group_tiles_by_temporal_slice(tiles)
|
||||
|
||||
# State for temporal overlap handling
|
||||
previous_chunk = None
|
||||
previous_weights = None
|
||||
previous_temporal_slice = None
|
||||
|
||||
for temporal_group_tiles in temporal_groups:
|
||||
curr_temporal_slice = temporal_group_tiles[0].out_coords[2]
|
||||
|
||||
# Calculate the shape of the temporal buffer for this group of tiles.
|
||||
# The temporal length depends on whether this is the first tile (starts at 0) or not.
|
||||
# - First tile: (frames - 1) * scale + 1
|
||||
# - Subsequent tiles: frames * scale
|
||||
# This logic is handled by TemporalAxisMapping and reflected in out_coords.
|
||||
temporal_tile_buffer_shape = full_video_shape._replace(
|
||||
frames=curr_temporal_slice.stop - curr_temporal_slice.start,
|
||||
)
|
||||
|
||||
buffer = torch.zeros(
|
||||
temporal_tile_buffer_shape.to_torch_shape(),
|
||||
device=latent.device,
|
||||
dtype=latent.dtype,
|
||||
)
|
||||
|
||||
curr_weights = self._accumulate_temporal_group_into_buffer(
|
||||
group_tiles=temporal_group_tiles,
|
||||
buffer=buffer,
|
||||
latent=latent,
|
||||
timestep=timestep,
|
||||
generator=generator,
|
||||
)
|
||||
|
||||
# Blend with previous temporal chunk if it exists
|
||||
if previous_chunk is not None:
|
||||
# Check if current temporal slice overlaps with previous temporal slice
|
||||
if previous_temporal_slice.stop > curr_temporal_slice.start:
|
||||
overlap_len = previous_temporal_slice.stop - curr_temporal_slice.start
|
||||
temporal_overlap_slice = slice(curr_temporal_slice.start - previous_temporal_slice.start, None)
|
||||
|
||||
# The overlap is already masked before it reaches this step. Each tile is accumulated into buffer
|
||||
# with its trapezoidal mask, and curr_weights accumulates the same mask. In the overlap blend we add
|
||||
# the masked values (buffer[...]) and the corresponding weights (curr_weights[...]) into the
|
||||
# previous buffers, then later normalize by weights.
|
||||
previous_chunk[:, :, temporal_overlap_slice, :, :] += buffer[:, :, slice(0, overlap_len), :, :]
|
||||
previous_weights[:, :, temporal_overlap_slice, :, :] += curr_weights[
|
||||
:, :, slice(0, overlap_len), :, :
|
||||
]
|
||||
|
||||
buffer[:, :, slice(0, overlap_len), :, :] = previous_chunk[:, :, temporal_overlap_slice, :, :]
|
||||
curr_weights[:, :, slice(0, overlap_len), :, :] = previous_weights[
|
||||
:, :, temporal_overlap_slice, :, :
|
||||
]
|
||||
|
||||
# Yield the non-overlapping part of the previous chunk
|
||||
previous_weights = previous_weights.clamp(min=1e-8)
|
||||
yield_len = curr_temporal_slice.start - previous_temporal_slice.start
|
||||
yield (previous_chunk / previous_weights)[:, :, :yield_len, :, :]
|
||||
|
||||
# Update state for next iteration
|
||||
previous_chunk = buffer
|
||||
previous_weights = curr_weights
|
||||
previous_temporal_slice = curr_temporal_slice
|
||||
|
||||
# Yield any remaining chunk
|
||||
if previous_chunk is not None:
|
||||
previous_weights = previous_weights.clamp(min=1e-8)
|
||||
yield previous_chunk / previous_weights
|
||||
|
||||
def _group_tiles_by_temporal_slice(self, tiles: List[Tile]) -> List[List[Tile]]:
|
||||
"""Group tiles by their temporal output slice."""
|
||||
if not tiles:
|
||||
return []
|
||||
|
||||
groups = []
|
||||
current_slice = tiles[0].out_coords[2]
|
||||
current_group = []
|
||||
|
||||
for tile in tiles:
|
||||
tile_slice = tile.out_coords[2]
|
||||
if tile_slice == current_slice:
|
||||
current_group.append(tile)
|
||||
else:
|
||||
groups.append(current_group)
|
||||
current_slice = tile_slice
|
||||
current_group = [tile]
|
||||
|
||||
# Add the final group
|
||||
if current_group:
|
||||
groups.append(current_group)
|
||||
|
||||
return groups
|
||||
|
||||
def _accumulate_temporal_group_into_buffer(
|
||||
self,
|
||||
group_tiles: List[Tile],
|
||||
buffer: torch.Tensor,
|
||||
latent: torch.Tensor,
|
||||
timestep: Optional[torch.Tensor],
|
||||
generator: Optional[torch.Generator],
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Decode and accumulate all tiles of a temporal group into a local buffer.
|
||||
The buffer is local to the group and always starts at time 0; temporal coordinates
|
||||
are rebased by subtracting temporal_slice.start.
|
||||
"""
|
||||
temporal_slice = group_tiles[0].out_coords[2]
|
||||
|
||||
weights = torch.zeros_like(buffer)
|
||||
|
||||
for tile in group_tiles:
|
||||
decoded_tile = self.forward(latent[tile.in_coords], timestep, generator)
|
||||
mask = tile.blend_mask.to(device=buffer.device, dtype=buffer.dtype)
|
||||
temporal_offset = tile.out_coords[2].start - temporal_slice.start
|
||||
# Use the tile's output coordinate length, not the decoded tile's length,
|
||||
# as the decoder may produce a different number of frames than expected
|
||||
expected_temporal_len = tile.out_coords[2].stop - tile.out_coords[2].start
|
||||
decoded_temporal_len = decoded_tile.shape[2]
|
||||
|
||||
# Ensure we don't exceed the buffer or decoded tile bounds
|
||||
actual_temporal_len = min(expected_temporal_len, decoded_temporal_len, buffer.shape[2] - temporal_offset)
|
||||
|
||||
chunk_coords = (
|
||||
slice(None), # batch
|
||||
slice(None), # channels
|
||||
slice(temporal_offset, temporal_offset + actual_temporal_len),
|
||||
tile.out_coords[3], # height
|
||||
tile.out_coords[4], # width
|
||||
)
|
||||
|
||||
# Slice decoded_tile and mask to match the actual length we're writing
|
||||
decoded_slice = decoded_tile[:, :, :actual_temporal_len, :, :]
|
||||
mask_slice = mask[:, :, :actual_temporal_len, :, :] if mask.shape[2] > 1 else mask
|
||||
|
||||
buffer[chunk_coords] += decoded_slice * mask_slice
|
||||
weights[chunk_coords] += mask_slice
|
||||
|
||||
return weights
|
||||
|
||||
|
||||
def decode_video(
|
||||
latent: torch.Tensor,
|
||||
video_decoder: VideoDecoder,
|
||||
tiling_config: TilingConfig | None = None,
|
||||
) -> Iterator[torch.Tensor]:
|
||||
"""
|
||||
Decode a video latent tensor with the given decoder.
|
||||
Args:
|
||||
latent: Tensor [c, f, h, w]
|
||||
video_decoder: Decoder module.
|
||||
tiling_config: Optional tiling settings.
|
||||
Yields:
|
||||
Decoded chunk [f, h, w, c], uint8 in [0, 255].
|
||||
"""
|
||||
|
||||
def convert_to_uint8(frames: torch.Tensor) -> torch.Tensor:
|
||||
frames = (((frames + 1.0) / 2.0).clamp(0.0, 1.0) * 255.0).to(torch.uint8)
|
||||
frames = rearrange(frames[0], "c f h w -> f h w c")
|
||||
return frames
|
||||
|
||||
if tiling_config is not None:
|
||||
for frames in video_decoder.tiled_decode(latent, tiling_config):
|
||||
yield convert_to_uint8(frames)
|
||||
else:
|
||||
decoded_video = video_decoder(latent)
|
||||
yield convert_to_uint8(decoded_video)
|
||||
|
||||
|
||||
def get_video_chunks_number(num_frames: int, tiling_config: TilingConfig | None = None) -> int:
|
||||
"""
|
||||
Get the number of video chunks for a given number of frames and tiling configuration.
|
||||
Args:
|
||||
num_frames: Number of frames in the video.
|
||||
tiling_config: Tiling configuration.
|
||||
Returns:
|
||||
Number of video chunks.
|
||||
"""
|
||||
if not tiling_config or not tiling_config.temporal_config:
|
||||
return 1
|
||||
cfg = tiling_config.temporal_config
|
||||
frame_stride = cfg.tile_size_in_frames - cfg.tile_overlap_in_frames
|
||||
return (num_frames - 1 + frame_stride - 1) // frame_stride
|
||||
|
||||
|
||||
def split_in_spatial(size: int, overlap: int) -> SplitOperation:
|
||||
def split(dimension_size: int) -> DimensionIntervals:
|
||||
if dimension_size <= size:
|
||||
return DEFAULT_SPLIT_OPERATION(dimension_size)
|
||||
amount = (dimension_size + size - 2 * overlap - 1) // (size - overlap)
|
||||
starts = [i * (size - overlap) for i in range(amount)]
|
||||
ends = [start + size for start in starts]
|
||||
ends[-1] = dimension_size
|
||||
left_ramps = [0] + [overlap] * (amount - 1)
|
||||
right_ramps = [overlap] * (amount - 1) + [0]
|
||||
return DimensionIntervals(starts=starts, ends=ends, left_ramps=left_ramps, right_ramps=right_ramps)
|
||||
|
||||
return split
|
||||
|
||||
|
||||
def split_in_temporal(size: int, overlap: int) -> SplitOperation:
|
||||
non_causal_split = split_in_spatial(size, overlap)
|
||||
|
||||
def split(dimension_size: int) -> DimensionIntervals:
|
||||
if dimension_size <= size:
|
||||
return DEFAULT_SPLIT_OPERATION(dimension_size)
|
||||
intervals = non_causal_split(dimension_size)
|
||||
starts = intervals.starts
|
||||
starts[1:] = [s - 1 for s in starts[1:]]
|
||||
left_ramps = intervals.left_ramps
|
||||
left_ramps[1:] = [r + 1 for r in left_ramps[1:]]
|
||||
return replace(intervals, starts=starts, left_ramps=left_ramps)
|
||||
|
||||
return split
|
||||
|
||||
|
||||
def to_mapping_operation(
|
||||
map_func: Callable[[int, int, int, int, int], Tuple[slice, torch.Tensor]],
|
||||
scale: int,
|
||||
) -> MappingOperation:
|
||||
def map_op(intervals: DimensionIntervals) -> tuple[list[slice], list[torch.Tensor | None]]:
|
||||
output_slices: list[slice] = []
|
||||
masks_1d: list[torch.Tensor | None] = []
|
||||
number_of_slices = len(intervals.starts)
|
||||
for i in range(number_of_slices):
|
||||
start = intervals.starts[i]
|
||||
end = intervals.ends[i]
|
||||
left_ramp = intervals.left_ramps[i]
|
||||
right_ramp = intervals.right_ramps[i]
|
||||
output_slice, mask_1d = map_func(start, end, left_ramp, right_ramp, scale)
|
||||
output_slices.append(output_slice)
|
||||
masks_1d.append(mask_1d)
|
||||
return output_slices, masks_1d
|
||||
|
||||
return map_op
|
||||
|
||||
|
||||
def map_temporal_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]:
|
||||
start = begin * scale
|
||||
stop = 1 + (end - 1) * scale
|
||||
left_ramp = 1 + (left_ramp - 1) * scale
|
||||
right_ramp = right_ramp * scale
|
||||
|
||||
return slice(start, stop), compute_trapezoidal_mask_1d(stop - start, left_ramp, right_ramp, True)
|
||||
|
||||
|
||||
def map_spatial_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]:
|
||||
start = begin * scale
|
||||
stop = end * scale
|
||||
left_ramp = left_ramp * scale
|
||||
right_ramp = right_ramp * scale
|
||||
|
||||
return slice(start, stop), compute_trapezoidal_mask_1d(stop - start, left_ramp, right_ramp, False)
|
||||
@@ -0,0 +1 @@
|
||||
"""CLIP/text encoder model components."""
|
||||
@@ -0,0 +1,31 @@
|
||||
"""Gemma text encoder components."""
|
||||
|
||||
from ltx_core.text_encoders.gemma.encoders.av_encoder import (
|
||||
AV_GEMMA_TEXT_ENCODER_KEY_OPS,
|
||||
AVGemmaEncoderOutput,
|
||||
AVGemmaTextEncoderModel,
|
||||
AVGemmaTextEncoderModelConfigurator,
|
||||
)
|
||||
from ltx_core.text_encoders.gemma.encoders.base_encoder import (
|
||||
GemmaTextEncoderModelBase,
|
||||
encode_text,
|
||||
module_ops_from_gemma_root,
|
||||
)
|
||||
from ltx_core.text_encoders.gemma.encoders.video_only_encoder import (
|
||||
VideoGemmaEncoderOutput,
|
||||
VideoGemmaTextEncoderModel,
|
||||
VideoGemmaTextEncoderModelConfigurator,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AV_GEMMA_TEXT_ENCODER_KEY_OPS",
|
||||
"AVGemmaEncoderOutput",
|
||||
"AVGemmaTextEncoderModel",
|
||||
"AVGemmaTextEncoderModelConfigurator",
|
||||
"GemmaTextEncoderModelBase",
|
||||
"VideoGemmaEncoderOutput",
|
||||
"VideoGemmaTextEncoderModel",
|
||||
"VideoGemmaTextEncoderModelConfigurator",
|
||||
"encode_text",
|
||||
"module_ops_from_gemma_root",
|
||||
]
|
||||
@@ -0,0 +1,210 @@
|
||||
import torch
|
||||
|
||||
from ltx_core.model.model_protocol import ModelConfigurator
|
||||
from ltx_core.model.transformer.attention import Attention
|
||||
from ltx_core.model.transformer.feed_forward import FeedForward
|
||||
from ltx_core.model.transformer.rope import (
|
||||
LTXRopeType,
|
||||
generate_freq_grid_np,
|
||||
generate_freq_grid_pytorch,
|
||||
precompute_freqs_cis,
|
||||
)
|
||||
from ltx_core.utils import rms_norm
|
||||
|
||||
|
||||
class _BasicTransformerBlock1D(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
heads: int,
|
||||
dim_head: int,
|
||||
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.attn1 = Attention(
|
||||
query_dim=dim,
|
||||
heads=heads,
|
||||
dim_head=dim_head,
|
||||
rope_type=rope_type,
|
||||
)
|
||||
|
||||
self.ff = FeedForward(
|
||||
dim,
|
||||
dim_out=dim,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
attention_mask: torch.Tensor | None = None,
|
||||
pe: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
# Notice that normalization is always applied before the real computation in the following blocks.
|
||||
|
||||
# 1. Normalization Before Self-Attention
|
||||
norm_hidden_states = rms_norm(hidden_states)
|
||||
|
||||
norm_hidden_states = norm_hidden_states.squeeze(1)
|
||||
|
||||
# 2. Self-Attention
|
||||
attn_output = self.attn1(norm_hidden_states, mask=attention_mask, pe=pe)
|
||||
|
||||
hidden_states = attn_output + hidden_states
|
||||
if hidden_states.ndim == 4:
|
||||
hidden_states = hidden_states.squeeze(1)
|
||||
|
||||
# 3. Normalization before Feed-Forward
|
||||
norm_hidden_states = rms_norm(hidden_states)
|
||||
|
||||
# 4. Feed-forward
|
||||
ff_output = self.ff(norm_hidden_states)
|
||||
|
||||
hidden_states = ff_output + hidden_states
|
||||
if hidden_states.ndim == 4:
|
||||
hidden_states = hidden_states.squeeze(1)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class Embeddings1DConnector(torch.nn.Module):
|
||||
"""
|
||||
Embeddings1DConnector applies a 1D transformer-based processing to sequential embeddings (e.g., for video, audio, or
|
||||
other modalities). It supports rotary positional encoding (rope), optional causal temporal positioning, and can
|
||||
substitute padded positions with learnable registers. The module is highly configurable for head size, number of
|
||||
layers, and register usage.
|
||||
Args:
|
||||
attention_head_dim (int): Dimension of each attention head (default=128).
|
||||
num_attention_heads (int): Number of attention heads (default=30).
|
||||
num_layers (int): Number of transformer layers (default=2).
|
||||
positional_embedding_theta (float): Scaling factor for position embedding (default=10000.0).
|
||||
positional_embedding_max_pos (list[int] | None): Max positions for positional embeddings (default=[1]).
|
||||
causal_temporal_positioning (bool): If True, uses causal attention (default=False).
|
||||
num_learnable_registers (int | None): Number of learnable registers to replace padded tokens. If None, disables
|
||||
register replacement. (default=128)
|
||||
rope_type (LTXRopeType): The RoPE variant to use (default=DEFAULT_ROPE_TYPE).
|
||||
double_precision_rope (bool): Use double precision rope calculation (default=False).
|
||||
"""
|
||||
|
||||
_supports_gradient_checkpointing = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
attention_head_dim: int = 128,
|
||||
num_attention_heads: int = 30,
|
||||
num_layers: int = 2,
|
||||
positional_embedding_theta: float = 10000.0,
|
||||
positional_embedding_max_pos: list[int] | None = None,
|
||||
causal_temporal_positioning: bool = False,
|
||||
num_learnable_registers: int | None = 128,
|
||||
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
|
||||
double_precision_rope: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.inner_dim = num_attention_heads * attention_head_dim
|
||||
self.causal_temporal_positioning = causal_temporal_positioning
|
||||
self.positional_embedding_theta = positional_embedding_theta
|
||||
self.positional_embedding_max_pos = (
|
||||
positional_embedding_max_pos if positional_embedding_max_pos is not None else [1]
|
||||
)
|
||||
self.rope_type = rope_type
|
||||
self.double_precision_rope = double_precision_rope
|
||||
self.transformer_1d_blocks = torch.nn.ModuleList(
|
||||
[
|
||||
_BasicTransformerBlock1D(
|
||||
dim=self.inner_dim,
|
||||
heads=num_attention_heads,
|
||||
dim_head=attention_head_dim,
|
||||
rope_type=rope_type,
|
||||
)
|
||||
for _ in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
self.num_learnable_registers = num_learnable_registers
|
||||
if self.num_learnable_registers:
|
||||
self.learnable_registers = torch.nn.Parameter(
|
||||
torch.rand(self.num_learnable_registers, self.inner_dim, dtype=torch.bfloat16) * 2.0 - 1.0
|
||||
)
|
||||
|
||||
def _replace_padded_with_learnable_registers(
|
||||
self, hidden_states: torch.Tensor, attention_mask: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
assert hidden_states.shape[1] % self.num_learnable_registers == 0, (
|
||||
f"Hidden states sequence length {hidden_states.shape[1]} must be divisible by num_learnable_registers "
|
||||
f"{self.num_learnable_registers}."
|
||||
)
|
||||
|
||||
num_registers_duplications = hidden_states.shape[1] // self.num_learnable_registers
|
||||
learnable_registers = torch.tile(self.learnable_registers, (num_registers_duplications, 1))
|
||||
attention_mask_binary = (attention_mask.squeeze(1).squeeze(1).unsqueeze(-1) >= -9000.0).int()
|
||||
|
||||
non_zero_hidden_states = hidden_states[:, attention_mask_binary.squeeze().bool(), :]
|
||||
non_zero_nums = non_zero_hidden_states.shape[1]
|
||||
pad_length = hidden_states.shape[1] - non_zero_nums
|
||||
adjusted_hidden_states = torch.nn.functional.pad(non_zero_hidden_states, pad=(0, 0, 0, pad_length), value=0)
|
||||
flipped_mask = torch.flip(attention_mask_binary, dims=[1])
|
||||
hidden_states = flipped_mask * adjusted_hidden_states + (1 - flipped_mask) * learnable_registers
|
||||
|
||||
attention_mask = torch.full_like(
|
||||
attention_mask,
|
||||
0.0,
|
||||
dtype=attention_mask.dtype,
|
||||
device=attention_mask.device,
|
||||
)
|
||||
|
||||
return hidden_states, attention_mask
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
attention_mask: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Forward pass of Embeddings1DConnector.
|
||||
Args:
|
||||
hidden_states (torch.Tensor): Input tensor of embeddings (shape [batch, seq_len, feature_dim]).
|
||||
attention_mask (torch.Tensor|None): Optional mask for valid tokens (shape compatible with hidden_states).
|
||||
Returns:
|
||||
tuple[torch.Tensor, torch.Tensor]: Processed features and the corresponding (possibly modified) mask.
|
||||
"""
|
||||
if self.num_learnable_registers:
|
||||
hidden_states, attention_mask = self._replace_padded_with_learnable_registers(hidden_states, attention_mask)
|
||||
|
||||
indices_grid = torch.arange(hidden_states.shape[1], dtype=torch.float32, device=hidden_states.device)
|
||||
indices_grid = indices_grid[None, None, :]
|
||||
freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch
|
||||
freqs_cis = precompute_freqs_cis(
|
||||
indices_grid=indices_grid,
|
||||
dim=self.inner_dim,
|
||||
out_dtype=hidden_states.dtype,
|
||||
theta=self.positional_embedding_theta,
|
||||
max_pos=self.positional_embedding_max_pos,
|
||||
num_attention_heads=self.num_attention_heads,
|
||||
rope_type=self.rope_type,
|
||||
freq_grid_generator=freq_grid_generator,
|
||||
)
|
||||
|
||||
for block in self.transformer_1d_blocks:
|
||||
hidden_states = block(hidden_states, attention_mask=attention_mask, pe=freqs_cis)
|
||||
|
||||
hidden_states = rms_norm(hidden_states)
|
||||
|
||||
return hidden_states, attention_mask
|
||||
|
||||
|
||||
class Embeddings1DConnectorConfigurator(ModelConfigurator[Embeddings1DConnector]):
|
||||
@classmethod
|
||||
def from_config(cls: type[Embeddings1DConnector], config: dict) -> Embeddings1DConnector:
|
||||
config = config.get("transformer", {})
|
||||
rope_type = LTXRopeType(config.get("rope_type", "interleaved"))
|
||||
double_precision_rope = config.get("frequencies_precision", False) == "float64"
|
||||
pe_max_pos = config.get("connector_positional_embedding_max_pos", [1])
|
||||
|
||||
connector = Embeddings1DConnector(
|
||||
positional_embedding_max_pos=pe_max_pos,
|
||||
rope_type=rope_type,
|
||||
double_precision_rope=double_precision_rope,
|
||||
)
|
||||
return connector
|
||||
@@ -0,0 +1,94 @@
|
||||
from typing import NamedTuple
|
||||
|
||||
import torch
|
||||
from transformers.models.gemma3 import Gemma3ForConditionalGeneration
|
||||
|
||||
from ltx_core.loader.sd_ops import SDOps
|
||||
from ltx_core.model.model_protocol import ModelConfigurator
|
||||
from ltx_core.text_encoders.gemma.embeddings_connector import (
|
||||
Embeddings1DConnector,
|
||||
Embeddings1DConnectorConfigurator,
|
||||
)
|
||||
from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoderModelBase
|
||||
from ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear
|
||||
from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
|
||||
|
||||
|
||||
class AVGemmaEncoderOutput(NamedTuple):
|
||||
video_encoding: torch.Tensor
|
||||
audio_encoding: torch.Tensor
|
||||
attention_mask: torch.Tensor
|
||||
|
||||
|
||||
class AVGemmaTextEncoderModel(GemmaTextEncoderModelBase):
|
||||
"""
|
||||
AVGemma Text Encoder Model.
|
||||
This class combines the tokenizer, Gemma model, feature extractor from base class and a
|
||||
video and audio embeddings connectors to provide a preprocessing for audio-visual pipeline.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
feature_extractor_linear: GemmaFeaturesExtractorProjLinear,
|
||||
embeddings_connector: Embeddings1DConnector,
|
||||
audio_embeddings_connector: Embeddings1DConnector,
|
||||
tokenizer: LTXVGemmaTokenizer | None = None,
|
||||
model: Gemma3ForConditionalGeneration | None = None,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
feature_extractor_linear=feature_extractor_linear,
|
||||
tokenizer=tokenizer,
|
||||
model=model,
|
||||
dtype=dtype,
|
||||
)
|
||||
self.embeddings_connector = embeddings_connector.to(dtype=dtype)
|
||||
self.audio_embeddings_connector = audio_embeddings_connector.to(dtype=dtype)
|
||||
|
||||
def _run_connectors(
|
||||
self, encoded_input: torch.Tensor, attention_mask: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
connector_attention_mask = self._convert_to_additive_mask(attention_mask, encoded_input.dtype)
|
||||
|
||||
encoded, encoded_connector_attention_mask = self.embeddings_connector(
|
||||
encoded_input,
|
||||
connector_attention_mask,
|
||||
)
|
||||
|
||||
# restore the mask values to int64
|
||||
attention_mask = (encoded_connector_attention_mask < 0.000001).to(torch.int64)
|
||||
attention_mask = attention_mask.reshape([encoded.shape[0], encoded.shape[1], 1])
|
||||
encoded = encoded * attention_mask
|
||||
|
||||
encoded_for_audio, _ = self.audio_embeddings_connector(encoded_input, connector_attention_mask)
|
||||
|
||||
return encoded, encoded_for_audio, attention_mask.squeeze(-1)
|
||||
|
||||
def forward(self, text: str, padding_side: str = "left") -> AVGemmaEncoderOutput:
|
||||
encoded_inputs, attention_mask = self._preprocess_text(text, padding_side)
|
||||
video_encoding, audio_encoding, attention_mask = self._run_connectors(encoded_inputs, attention_mask)
|
||||
return AVGemmaEncoderOutput(video_encoding, audio_encoding, attention_mask)
|
||||
|
||||
|
||||
class AVGemmaTextEncoderModelConfigurator(ModelConfigurator[AVGemmaTextEncoderModel]):
|
||||
@classmethod
|
||||
def from_config(cls: type["AVGemmaTextEncoderModel"], config: dict) -> "AVGemmaTextEncoderModel":
|
||||
feature_extractor_linear = GemmaFeaturesExtractorProjLinear.from_config(config)
|
||||
embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
|
||||
audio_embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
|
||||
return AVGemmaTextEncoderModel(
|
||||
feature_extractor_linear=feature_extractor_linear,
|
||||
embeddings_connector=embeddings_connector,
|
||||
audio_embeddings_connector=audio_embeddings_connector,
|
||||
)
|
||||
|
||||
|
||||
AV_GEMMA_TEXT_ENCODER_KEY_OPS = (
|
||||
SDOps("AV_GEMMA_TEXT_ENCODER_KEY_OPS")
|
||||
.with_matching(prefix="text_embedding_projection.")
|
||||
.with_matching(prefix="model.diffusion_model.audio_embeddings_connector.")
|
||||
.with_matching(prefix="model.diffusion_model.video_embeddings_connector.")
|
||||
.with_replacement("text_embedding_projection.", "feature_extractor_linear.")
|
||||
.with_replacement("model.diffusion_model.video_embeddings_connector.", "embeddings_connector.")
|
||||
.with_replacement("model.diffusion_model.audio_embeddings_connector.", "audio_embeddings_connector.")
|
||||
)
|
||||
@@ -0,0 +1,329 @@
|
||||
import functools
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from transformers import AutoImageProcessor, Gemma3ForConditionalGeneration, Gemma3Processor
|
||||
|
||||
from ltx_core.loader.module_ops import ModuleOps
|
||||
from ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear
|
||||
from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
|
||||
|
||||
|
||||
class GemmaTextEncoderModelBase(torch.nn.Module):
|
||||
"""
|
||||
Gemma Text Encoder Model.
|
||||
This base class combines the tokenizer, Gemma model and feature extractor to provide a preprocessing
|
||||
for implementation classes for multimodal pipelines. It processes input text through tokenization,
|
||||
obtains hidden states from the base language model, applies a linear feature extractor.
|
||||
Args:
|
||||
tokenizer (LTXVGemmaTokenizer): The tokenizer used for text preprocessing.
|
||||
model (Gemma3ForConditionalGeneration): The base Gemma LLM.
|
||||
feature_extractor_linear (GemmaFeaturesExtractorProjLinear): Linear projection for hidden state aggregation.
|
||||
dtype (torch.dtype, optional): The data type for model parameters (default: torch.bfloat16).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
feature_extractor_linear: GemmaFeaturesExtractorProjLinear,
|
||||
tokenizer: LTXVGemmaTokenizer | None = None,
|
||||
model: Gemma3ForConditionalGeneration | None = None,
|
||||
img_processor: Gemma3Processor | None = None,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self._gemma_root = None
|
||||
self.tokenizer = tokenizer
|
||||
self.model = model
|
||||
self.processor = img_processor
|
||||
self.feature_extractor_linear = feature_extractor_linear.to(dtype=dtype)
|
||||
|
||||
def _run_feature_extractor(
|
||||
self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, padding_side: str = "right"
|
||||
) -> torch.Tensor:
|
||||
encoded_text_features = torch.stack(hidden_states, dim=-1)
|
||||
encoded_text_features_dtype = encoded_text_features.dtype
|
||||
|
||||
sequence_lengths = attention_mask.sum(dim=-1)
|
||||
normed_concated_encoded_text_features = _norm_and_concat_padded_batch(
|
||||
encoded_text_features, sequence_lengths, padding_side=padding_side
|
||||
)
|
||||
|
||||
return self.feature_extractor_linear(normed_concated_encoded_text_features.to(encoded_text_features_dtype))
|
||||
|
||||
def _convert_to_additive_mask(self, attention_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
|
||||
return (attention_mask - 1).to(dtype).reshape(
|
||||
(attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
|
||||
) * torch.finfo(dtype).max
|
||||
|
||||
def _preprocess_text(self, text: str, padding_side: str = "left") -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
|
||||
"""
|
||||
Encode a given string into feature tensors suitable for downstream tasks.
|
||||
Args:
|
||||
text (str): Input string to encode.
|
||||
Returns:
|
||||
tuple[torch.Tensor, dict[str, torch.Tensor]]: Encoded features and a dictionary with attention mask.
|
||||
"""
|
||||
token_pairs = self.tokenizer.tokenize_with_weights(text)["gemma"]
|
||||
input_ids = torch.tensor([[t[0] for t in token_pairs]], device=self.model.device)
|
||||
attention_mask = torch.tensor([[w[1] for w in token_pairs]], device=self.model.device)
|
||||
outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)
|
||||
projected = self._run_feature_extractor(
|
||||
hidden_states=outputs.hidden_states, attention_mask=attention_mask, padding_side=padding_side
|
||||
)
|
||||
return projected, attention_mask
|
||||
|
||||
def _init_image_processor(self) -> None:
|
||||
img_processor = AutoImageProcessor.from_pretrained(self._gemma_root, local_files_only=True)
|
||||
if not self.tokenizer:
|
||||
raise ValueError("Tokenizer is not loaded, cannot load image processor")
|
||||
self.processor = Gemma3Processor(image_processor=img_processor, tokenizer=self.tokenizer.tokenizer)
|
||||
|
||||
def _enhance(
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
image: torch.Tensor | None = None,
|
||||
max_new_tokens: int = 512,
|
||||
seed: int = 42,
|
||||
) -> str:
|
||||
if self.processor is None:
|
||||
self._init_image_processor()
|
||||
text = self.processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
||||
|
||||
model_inputs = self.processor(
|
||||
text=text,
|
||||
images=image,
|
||||
return_tensors="pt",
|
||||
).to(self.model.device)
|
||||
pad_token_id = self.processor.tokenizer.pad_token_id if self.processor.tokenizer.pad_token_id is not None else 0
|
||||
model_inputs = _pad_inputs_for_attention_alignment(model_inputs, pad_token_id=pad_token_id)
|
||||
|
||||
with torch.inference_mode(), torch.random.fork_rng(devices=[self.model.device]):
|
||||
torch.manual_seed(seed)
|
||||
outputs = self.model.generate(
|
||||
**model_inputs,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=True,
|
||||
temperature=0.7,
|
||||
)
|
||||
generated_ids = outputs[0][len(model_inputs.input_ids[0]) :]
|
||||
enhanced_prompt = self.processor.tokenizer.decode(generated_ids, skip_special_tokens=True)
|
||||
|
||||
return enhanced_prompt
|
||||
|
||||
def enhance_t2v(
|
||||
self,
|
||||
prompt: str,
|
||||
max_new_tokens: int = 512,
|
||||
system_prompt: str | None = None,
|
||||
seed: int = 42,
|
||||
) -> str:
|
||||
"""Enhance a text prompt for T2V generation."""
|
||||
|
||||
system_prompt = system_prompt or self.default_gemma_t2v_system_prompt
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": f"user prompt: {prompt}"},
|
||||
]
|
||||
|
||||
return self._enhance(messages, max_new_tokens=max_new_tokens, seed=seed)
|
||||
|
||||
def enhance_i2v(
|
||||
self,
|
||||
prompt: str,
|
||||
image: torch.Tensor,
|
||||
max_new_tokens: int = 512,
|
||||
system_prompt: str | None = None,
|
||||
seed: int = 42,
|
||||
) -> str:
|
||||
"""Enhance a text prompt for I2V generation using a reference image."""
|
||||
system_prompt = system_prompt or self.default_gemma_i2v_system_prompt
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image"},
|
||||
{"type": "text", "text": f"User Raw Input Prompt: {prompt}."},
|
||||
],
|
||||
},
|
||||
]
|
||||
return self._enhance(messages, image=image, max_new_tokens=max_new_tokens, seed=seed)
|
||||
|
||||
@functools.cached_property
|
||||
def default_gemma_i2v_system_prompt(self) -> str:
|
||||
return _load_system_prompt("gemma_i2v_system_prompt.txt")
|
||||
|
||||
@functools.cached_property
|
||||
def default_gemma_t2v_system_prompt(self) -> str:
|
||||
return _load_system_prompt("gemma_t2v_system_prompt.txt")
|
||||
|
||||
def forward(self, text: str, padding_side: str = "left") -> tuple[torch.Tensor, torch.Tensor]:
|
||||
raise NotImplementedError("This method is not implemented for the base class")
|
||||
|
||||
|
||||
def _norm_and_concat_padded_batch(
|
||||
encoded_text: torch.Tensor,
|
||||
sequence_lengths: torch.Tensor,
|
||||
padding_side: str = "right",
|
||||
) -> torch.Tensor:
|
||||
"""Normalize and flatten multi-layer hidden states, respecting padding.
|
||||
Performs per-batch, per-layer normalization using masked mean and range,
|
||||
then concatenates across the layer dimension.
|
||||
Args:
|
||||
encoded_text: Hidden states of shape [batch, seq_len, hidden_dim, num_layers].
|
||||
sequence_lengths: Number of valid (non-padded) tokens per batch item.
|
||||
padding_side: Whether padding is on "left" or "right".
|
||||
Returns:
|
||||
Normalized tensor of shape [batch, seq_len, hidden_dim * num_layers],
|
||||
with padded positions zeroed out.
|
||||
"""
|
||||
b, t, d, l = encoded_text.shape # noqa: E741
|
||||
device = encoded_text.device
|
||||
|
||||
# Build mask: [B, T, 1, 1]
|
||||
token_indices = torch.arange(t, device=device)[None, :] # [1, T]
|
||||
|
||||
if padding_side == "right":
|
||||
# For right padding, valid tokens are from 0 to sequence_length-1
|
||||
mask = token_indices < sequence_lengths[:, None] # [B, T]
|
||||
elif padding_side == "left":
|
||||
# For left padding, valid tokens are from (T - sequence_length) to T-1
|
||||
start_indices = t - sequence_lengths[:, None] # [B, 1]
|
||||
mask = token_indices >= start_indices # [B, T]
|
||||
else:
|
||||
raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}")
|
||||
|
||||
mask = rearrange(mask, "b t -> b t 1 1")
|
||||
|
||||
eps = 1e-6
|
||||
|
||||
# Compute masked mean: [B, 1, 1, L]
|
||||
masked = encoded_text.masked_fill(~mask, 0.0)
|
||||
denom = (sequence_lengths * d).view(b, 1, 1, 1)
|
||||
mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + eps)
|
||||
|
||||
# Compute masked min/max: [B, 1, 1, L]
|
||||
x_min = encoded_text.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True)
|
||||
x_max = encoded_text.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True)
|
||||
range_ = x_max - x_min
|
||||
|
||||
# Normalize only the valid tokens
|
||||
normed = 8 * (encoded_text - mean) / (range_ + eps)
|
||||
|
||||
# concat to be [Batch, T, D * L] - this preserves the original structure
|
||||
normed = normed.reshape(b, t, -1) # [B, T, D * L]
|
||||
|
||||
# Apply mask to preserve original padding (set padded positions to 0)
|
||||
mask_flattened = rearrange(mask, "b t 1 1 -> b t 1").expand(-1, -1, d * l)
|
||||
normed = normed.masked_fill(~mask_flattened, 0.0)
|
||||
|
||||
return normed
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=2)
|
||||
def _load_system_prompt(prompt_name: str) -> str:
|
||||
with open(Path(__file__).parent / "prompts" / f"{prompt_name}", "r") as f:
|
||||
return f.read()
|
||||
|
||||
|
||||
def _find_matching_dir(root_path: str, pattern: str) -> str:
|
||||
"""
|
||||
Recursively search for files matching a glob pattern and return the parent directory of the first match.
|
||||
"""
|
||||
|
||||
matches = list(Path(root_path).rglob(pattern))
|
||||
if not matches:
|
||||
raise FileNotFoundError(f"No files matching pattern '{pattern}' found under {root_path}")
|
||||
return str(matches[0].parent)
|
||||
|
||||
|
||||
def module_ops_from_gemma_root(gemma_root: str) -> tuple[ModuleOps, ...]:
|
||||
gemma_path = _find_matching_dir(gemma_root, "model*.safetensors")
|
||||
tokenizer_path = _find_matching_dir(gemma_root, "tokenizer.model")
|
||||
|
||||
def load_gemma(module: GemmaTextEncoderModelBase) -> GemmaTextEncoderModelBase:
|
||||
module.model = Gemma3ForConditionalGeneration.from_pretrained(
|
||||
gemma_path, local_files_only=True, torch_dtype=torch.bfloat16
|
||||
)
|
||||
module._gemma_root = module._gemma_root or gemma_root
|
||||
return module
|
||||
|
||||
def load_tokenizer(module: GemmaTextEncoderModelBase) -> GemmaTextEncoderModelBase:
|
||||
module.tokenizer = LTXVGemmaTokenizer(tokenizer_path, 1024)
|
||||
module._gemma_root = module._gemma_root or gemma_root
|
||||
return module
|
||||
|
||||
gemma_load_ops = ModuleOps(
|
||||
"GemmaLoad",
|
||||
matcher=lambda module: isinstance(module, GemmaTextEncoderModelBase) and module.model is None,
|
||||
mutator=load_gemma,
|
||||
)
|
||||
tokenizer_load_ops = ModuleOps(
|
||||
"TokenizerLoad",
|
||||
matcher=lambda module: isinstance(module, GemmaTextEncoderModelBase) and module.tokenizer is None,
|
||||
mutator=load_tokenizer,
|
||||
)
|
||||
return (gemma_load_ops, tokenizer_load_ops)
|
||||
|
||||
|
||||
def encode_text(text_encoder: GemmaTextEncoderModelBase, prompts: list[str]) -> list[tuple[torch.Tensor, torch.Tensor]]:
|
||||
"""
|
||||
Encode a list of prompts using the provided Gemma text encoder.
|
||||
Args:
|
||||
text_encoder: The Gemma text encoder instance.
|
||||
prompts: List of prompt strings to encode.
|
||||
Returns:
|
||||
List of tuples, each containing (v_context, a_context) tensors for each prompt.
|
||||
"""
|
||||
result = []
|
||||
for prompt in prompts:
|
||||
v_context, a_context, _ = text_encoder(prompt)
|
||||
result.append((v_context, a_context))
|
||||
return result
|
||||
|
||||
|
||||
def _cat_with_padding(
|
||||
tensor: torch.Tensor,
|
||||
padding_length: int,
|
||||
value: int | float,
|
||||
) -> torch.Tensor:
|
||||
"""Concatenate a tensor with a padding tensor of the given value."""
|
||||
return torch.cat(
|
||||
[
|
||||
tensor,
|
||||
torch.full(
|
||||
(1, padding_length),
|
||||
value,
|
||||
dtype=tensor.dtype,
|
||||
device=tensor.device,
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
|
||||
def _pad_inputs_for_attention_alignment(
|
||||
model_inputs: dict[str, torch.Tensor],
|
||||
pad_token_id: int = 0,
|
||||
alignment: int = 8,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""Pad sequence length to multiple of alignment for Flash Attention compatibility.
|
||||
Flash Attention within SDPA requires sequence lengths aligned to 8 bytes.
|
||||
This pads input_ids, attention_mask, and token_type_ids (if present) to prevent
|
||||
'p.attn_bias_ptr is not correctly aligned' errors.
|
||||
"""
|
||||
seq_len = model_inputs.input_ids.shape[1]
|
||||
padded_len = ((seq_len + alignment - 1) // alignment) * alignment
|
||||
padding_length = padded_len - seq_len
|
||||
|
||||
if padding_length > 0:
|
||||
model_inputs["input_ids"] = _cat_with_padding(model_inputs.input_ids, padding_length, pad_token_id)
|
||||
|
||||
model_inputs["attention_mask"] = _cat_with_padding(model_inputs.attention_mask, padding_length, 0)
|
||||
|
||||
if "token_type_ids" in model_inputs and model_inputs["token_type_ids"] is not None:
|
||||
model_inputs["token_type_ids"] = _cat_with_padding(model_inputs["token_type_ids"], padding_length, 0)
|
||||
|
||||
return model_inputs
|
||||
+30
@@ -0,0 +1,30 @@
|
||||
You are a Creative Assistant writing concise, action-focused image-to-video prompts. Given an image (first frame) and user Raw Input Prompt, generate a prompt to guide video generation from that image.
|
||||
|
||||
#### Guidelines:
|
||||
- Analyze the Image: Identify Subject, Setting, Elements, Style and Mood.
|
||||
- Follow user Raw Input Prompt: Include all requested motion, actions, camera movements, audio, and details. If in conflict with the image, prioritize user request while maintaining visual consistency (describe transition from image to user's scene).
|
||||
- Describe only changes from the image: Don't reiterate established visual details. Inaccurate descriptions may cause scene cuts.
|
||||
- Active language: Use present-progressive verbs ("is walking," "speaking"). If no action specified, describe natural movements.
|
||||
- Chronological flow: Use temporal connectors ("as," "then," "while").
|
||||
- Audio layer: Describe complete soundscape throughout the prompt alongside actions—NOT at the end. Align audio intensity with action tempo. Include natural background audio, ambient sounds, effects, speech or music (when requested). Be specific (e.g., "soft footsteps on tile") not vague (e.g., "ambient sound").
|
||||
- Speech (only when requested): Provide exact words in quotes with character's visual/voice characteristics (e.g., "The tall man speaks in a low, gravelly voice"), language if not English and accent if relevant. If general conversation mentioned without text, generate contextual quoted dialogue. (i.e., "The man is talking" input -> the output should include exact spoken words, like: "The man is talking in an excited voice saying: 'You won't believe what I just saw!' His hands gesture expressively as he speaks, eyebrows raised with enthusiasm. The ambient sound of a quiet room underscores his animated speech.")
|
||||
- Style: Include visual style at beginning: "Style: <style>, <rest of prompt>." If unclear, omit to avoid conflicts.
|
||||
- Visual and audio only: Describe only what is seen and heard. NO smell, taste, or tactile sensations.
|
||||
- Restrained language: Avoid dramatic terms. Use mild, natural, understated phrasing.
|
||||
|
||||
#### Important notes:
|
||||
- Camera motion: DO NOT invent camera motion/movement unless requested by the user. Make sure to include camera motion only if specified in the input.
|
||||
- Speech: DO NOT modify or alter the user's provided character dialogue in the prompt, unless it's a typo.
|
||||
- No timestamps or cuts: DO NOT use timestamps or describe scene cuts unless explicitly requested.
|
||||
- Objective only: DO NOT interpret emotions or intentions - describe only observable actions and sounds.
|
||||
- Format: DO NOT use phrases like "The scene opens with..." / "The video starts...". Start directly with Style (optional) and chronological scene description.
|
||||
- Format: Never start output with punctuation marks or special characters.
|
||||
- DO NOT invent dialogue unless the user mentions speech/talking/singing/conversation.
|
||||
- Your performance is CRITICAL. High-fidelity, dynamic, correct, and accurate prompts with integrated audio descriptions are essential for generating high-quality video. Your goal is flawless execution of these rules.
|
||||
|
||||
#### Output Format (Strict):
|
||||
- Single concise paragraph in natural English. NO titles, headings, prefaces, sections, code fences, or Markdown.
|
||||
- If unsafe/invalid, return original user prompt. Never ask questions or clarifications.
|
||||
|
||||
#### Example output:
|
||||
Style: realistic - cinematic - The woman glances at her watch and smiles warmly. She speaks in a cheerful, friendly voice, "I think we're right on time!" In the background, a café barista prepares drinks at the counter. The barista calls out in a clear, upbeat tone, "Two cappuccinos ready!" The sound of the espresso machine hissing softly blends with gentle background chatter and the light clinking of cups on saucers.
|
||||
+40
@@ -0,0 +1,40 @@
|
||||
You are a Creative Assistant. Given a user's raw input prompt describing a scene or concept, expand it into a detailed video generation prompt with specific visuals and integrated audio to guide a text-to-video model.
|
||||
|
||||
#### Guidelines
|
||||
- Strictly follow all aspects of the user's raw input: include every element requested (style, visuals, motions, actions, camera movement, audio).
|
||||
- If the input is vague, invent concrete details: lighting, textures, materials, scene settings, etc.
|
||||
- For characters: describe gender, clothing, hair, expressions. DO NOT invent unrequested characters.
|
||||
- Use active language: present-progressive verbs ("is walking," "speaking"). If no action specified, describe natural movements.
|
||||
- Maintain chronological flow: use temporal connectors ("as," "then," "while").
|
||||
- Audio layer: Describe complete soundscape (background audio, ambient sounds, SFX, speech/music when requested). Integrate sounds chronologically alongside actions. Be specific (e.g., "soft footsteps on tile"), not vague (e.g., "ambient sound is present").
|
||||
- Speech (only when requested):
|
||||
- For ANY speech-related input (talking, conversation, singing, etc.), ALWAYS include exact words in quotes with voice characteristics (e.g., "The man says in an excited voice: 'You won't believe what I just saw!'").
|
||||
- Specify language if not English and accent if relevant.
|
||||
- Style: Include visual style at the beginning: "Style: <style>, <rest of prompt>." Default to cinematic-realistic if unspecified. Omit if unclear.
|
||||
- Visual and audio only: NO non-visual/auditory senses (smell, taste, touch).
|
||||
- Restrained language: Avoid dramatic/exaggerated terms. Use mild, natural phrasing.
|
||||
- Colors: Use plain terms ("red dress"), not intensified ("vibrant blue," "bright red").
|
||||
- Lighting: Use neutral descriptions ("soft overhead light"), not harsh ("blinding light").
|
||||
- Facial features: Use delicate modifiers for subtle features (i.e., "subtle freckles").
|
||||
|
||||
#### Important notes:
|
||||
- Analyze the user's raw input carefully. In cases of FPV or POV, exclude the description of the subject whose POV is requested.
|
||||
- Camera motion: DO NOT invent camera motion unless requested by the user.
|
||||
- Speech: DO NOT modify user-provided character dialogue unless it's a typo.
|
||||
- No timestamps or cuts: DO NOT use timestamps or describe scene cuts unless explicitly requested.
|
||||
- Format: DO NOT use phrases like "The scene opens with...". Start directly with Style (optional) and chronological scene description.
|
||||
- Format: DO NOT start your response with special characters.
|
||||
- DO NOT invent dialogue unless the user mentions speech/talking/singing/conversation.
|
||||
- If the user's raw input prompt is highly detailed, chronological and in the requested format: DO NOT make major edits or introduce new elements. Add/enhance audio descriptions if missing.
|
||||
|
||||
#### Output Format (Strict):
|
||||
- Single continuous paragraph in natural language (English).
|
||||
- NO titles, headings, prefaces, code fences, or Markdown.
|
||||
- If unsafe/invalid, return original user prompt. Never ask questions or clarifications.
|
||||
|
||||
Your output quality is CRITICAL. Generate visually rich, dynamic prompts with integrated audio for high-quality video generation.
|
||||
|
||||
#### Example
|
||||
Input: "A woman at a coffee shop talking on the phone"
|
||||
Output:
|
||||
Style: realistic with cinematic lighting. In a medium close-up, a woman in her early 30s with shoulder-length brown hair sits at a small wooden table by the window. She wears a cream-colored turtleneck sweater, holding a white ceramic coffee cup in one hand and a smartphone to her ear with the other. Ambient cafe sounds fill the space—espresso machine hiss, quiet conversations, gentle clinking of cups. The woman listens intently, nodding slightly, then takes a sip of her coffee and sets it down with a soft clink. Her face brightens into a warm smile as she speaks in a clear, friendly voice, 'That sounds perfect! I'd love to meet up this weekend. How about Saturday afternoon?' She laughs softly—a genuine chuckle—and shifts in her chair. Behind her, other patrons move subtly in and out of focus. 'Great, I'll see you then,' she concludes cheerfully, lowering the phone.
|
||||
@@ -0,0 +1,85 @@
|
||||
from typing import NamedTuple
|
||||
|
||||
import torch
|
||||
from transformers import Gemma3ForConditionalGeneration
|
||||
|
||||
from ltx_core.loader.sd_ops import SDOps
|
||||
from ltx_core.model.model_protocol import ModelConfigurator
|
||||
from ltx_core.text_encoders.gemma.embeddings_connector import (
|
||||
Embeddings1DConnector,
|
||||
Embeddings1DConnectorConfigurator,
|
||||
)
|
||||
from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoderModelBase
|
||||
from ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear
|
||||
from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
|
||||
|
||||
|
||||
class VideoGemmaEncoderOutput(NamedTuple):
|
||||
video_encoding: torch.Tensor
|
||||
attention_mask: torch.Tensor
|
||||
|
||||
|
||||
class VideoGemmaTextEncoderModel(GemmaTextEncoderModelBase):
|
||||
"""
|
||||
Video Gemma Text Encoder Model.
|
||||
This class combines the tokenizer, Gemma model, feature extractor from base class and a
|
||||
video embeddings connector to provide a preprocessing for video only pipeline.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
feature_extractor_linear: GemmaFeaturesExtractorProjLinear,
|
||||
embeddings_connector: Embeddings1DConnector,
|
||||
tokenizer: LTXVGemmaTokenizer | None = None,
|
||||
model: Gemma3ForConditionalGeneration | None = None,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
feature_extractor_linear=feature_extractor_linear,
|
||||
tokenizer=tokenizer,
|
||||
model=model,
|
||||
dtype=dtype,
|
||||
)
|
||||
self.embeddings_connector = embeddings_connector.to(dtype=dtype)
|
||||
|
||||
def _run_connector(
|
||||
self, encoded_input: torch.Tensor, attention_mask: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
connector_attention_mask = self._convert_to_additive_mask(attention_mask, encoded_input.dtype)
|
||||
|
||||
encoded, encoded_connector_attention_mask = self.embeddings_connector(
|
||||
encoded_input,
|
||||
connector_attention_mask,
|
||||
)
|
||||
|
||||
# restore the mask values to int64
|
||||
attention_mask = (encoded_connector_attention_mask < 0.000001).to(torch.int64)
|
||||
attention_mask = attention_mask.reshape([encoded.shape[0], encoded.shape[1], 1])
|
||||
encoded = encoded * attention_mask
|
||||
|
||||
return encoded, attention_mask.squeeze(-1)
|
||||
|
||||
def forward(self, text: str, padding_side: str = "left") -> VideoGemmaEncoderOutput:
|
||||
encoded_inputs, attention_mask = self._preprocess_text(text, padding_side)
|
||||
video_encoding, attention_mask = self._run_connector(encoded_inputs, attention_mask)
|
||||
return VideoGemmaEncoderOutput(video_encoding, attention_mask)
|
||||
|
||||
|
||||
class VideoGemmaTextEncoderModelConfigurator(ModelConfigurator[VideoGemmaTextEncoderModel]):
|
||||
@classmethod
|
||||
def from_config(cls: type["VideoGemmaTextEncoderModel"], config: dict) -> "VideoGemmaTextEncoderModel":
|
||||
feature_extractor_linear = GemmaFeaturesExtractorProjLinear.from_config(config)
|
||||
embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
|
||||
return VideoGemmaTextEncoderModel(
|
||||
feature_extractor_linear=feature_extractor_linear,
|
||||
embeddings_connector=embeddings_connector,
|
||||
)
|
||||
|
||||
|
||||
VIDEO_ONLY_GEMMA_TEXT_ENCODER_KEY_OPS = (
|
||||
SDOps("VIDEO_ONLY_GEMMA_TEXT_ENCODER_KEY_OPS")
|
||||
.with_matching(prefix="text_embedding_projection.")
|
||||
.with_matching(prefix="model.diffusion_model.embeddings_connector.")
|
||||
.with_replacement("text_embedding_projection.", "feature_extractor_linear.")
|
||||
.with_replacement("model.diffusion_model.embeddings_connector.", "embeddings_connector.")
|
||||
)
|
||||
@@ -0,0 +1,36 @@
|
||||
import torch
|
||||
|
||||
from ltx_core.model.model_protocol import ModelConfigurator
|
||||
|
||||
|
||||
class GemmaFeaturesExtractorProjLinear(torch.nn.Module, ModelConfigurator["GemmaFeaturesExtractorProjLinear"]):
|
||||
"""
|
||||
Feature extractor module for Gemma models.
|
||||
This module applies a single linear projection to the input tensor.
|
||||
It expects a flattened feature tensor of shape (batch_size, 3840*49).
|
||||
The linear layer maps this to a (batch_size, 3840) embedding.
|
||||
Attributes:
|
||||
aggregate_embed (torch.nn.Linear): Linear projection layer.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""
|
||||
Initialize the GemmaFeaturesExtractorProjLinear module.
|
||||
The input dimension is expected to be 3840 * 49, and the output is 3840.
|
||||
"""
|
||||
super().__init__()
|
||||
self.aggregate_embed = torch.nn.Linear(3840 * 49, 3840, bias=False)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Forward pass for the feature extractor.
|
||||
Args:
|
||||
x (torch.Tensor): Input tensor of shape (batch_size, 3840 * 49).
|
||||
Returns:
|
||||
torch.Tensor: Output tensor of shape (batch_size, 3840).
|
||||
"""
|
||||
return self.aggregate_embed(x)
|
||||
|
||||
@classmethod
|
||||
def from_config(cls: type["GemmaFeaturesExtractorProjLinear"], _config: dict) -> "GemmaFeaturesExtractorProjLinear":
|
||||
return cls()
|
||||
@@ -0,0 +1,64 @@
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
|
||||
class LTXVGemmaTokenizer:
|
||||
"""
|
||||
Tokenizer wrapper for Gemma models compatible with LTXV processes.
|
||||
This class wraps HuggingFace's `AutoTokenizer` for use with Gemma text encoders,
|
||||
ensuring correct settings and output formatting for downstream consumption.
|
||||
"""
|
||||
|
||||
def __init__(self, tokenizer_path: str, max_length: int = 256):
|
||||
"""
|
||||
Initialize the tokenizer.
|
||||
Args:
|
||||
tokenizer_path (str): Path to the pretrained tokenizer files or model directory.
|
||||
max_length (int, optional): Max sequence length for encoding. Defaults to 256.
|
||||
"""
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(
|
||||
tokenizer_path, local_files_only=True, model_max_length=max_length
|
||||
)
|
||||
# Gemma expects left padding for chat-style prompts; for plain text it doesn't matter much.
|
||||
self.tokenizer.padding_side = "left"
|
||||
if self.tokenizer.pad_token is None:
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
|
||||
self.max_length = max_length
|
||||
|
||||
def tokenize_with_weights(self, text: str, return_word_ids: bool = False) -> dict[str, list[tuple[int, int]]]:
|
||||
"""
|
||||
Tokenize the given text and return token IDs and attention weights.
|
||||
Args:
|
||||
text (str): The input string to tokenize.
|
||||
return_word_ids (bool, optional): If True, includes the token's position (index) in the output tuples.
|
||||
If False (default), omits the indices.
|
||||
Returns:
|
||||
dict[str, list[tuple[int, int]]] OR dict[str, list[tuple[int, int, int]]]:
|
||||
A dictionary with a "gemma" key mapping to:
|
||||
- a list of (token_id, attention_mask) tuples if return_word_ids is False;
|
||||
- a list of (token_id, attention_mask, index) tuples if return_word_ids is True.
|
||||
Example:
|
||||
>>> tokenizer = LTXVGemmaTokenizer("path/to/tokenizer", max_length=8)
|
||||
>>> tokenizer.tokenize_with_weights("hello world")
|
||||
{'gemma': [(1234, 1), (5678, 1), (2, 0), ...]}
|
||||
"""
|
||||
text = text.strip()
|
||||
encoded = self.tokenizer(
|
||||
text,
|
||||
padding="max_length",
|
||||
max_length=self.max_length,
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
input_ids = encoded.input_ids
|
||||
attention_mask = encoded.attention_mask
|
||||
tuples = [
|
||||
(token_id, attn, i) for i, (token_id, attn) in enumerate(zip(input_ids[0], attention_mask[0], strict=True))
|
||||
]
|
||||
out = {"gemma": tuples}
|
||||
|
||||
if not return_word_ids:
|
||||
# Return only (token_id, attention_mask) pairs, omitting token position
|
||||
out = {k: [(t, w) for t, w, _ in v] for k, v in out.items()}
|
||||
|
||||
return out
|
||||
@@ -0,0 +1,184 @@
|
||||
from dataclasses import dataclass, replace
|
||||
from typing import Protocol
|
||||
|
||||
import torch
|
||||
from torch._prims_common import DeviceLikeType
|
||||
|
||||
from ltx_core.components.patchifiers import (
|
||||
AudioLatentShape,
|
||||
AudioPatchifier,
|
||||
VideoLatentPatchifier,
|
||||
VideoLatentShape,
|
||||
get_pixel_coords,
|
||||
)
|
||||
from ltx_core.components.protocols import Patchifier
|
||||
from ltx_core.types import LatentState, SpatioTemporalScaleFactors
|
||||
|
||||
DEFAULT_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
|
||||
|
||||
|
||||
class LatentTools(Protocol):
|
||||
"""
|
||||
Tools for building latent states.
|
||||
"""
|
||||
|
||||
patchifier: Patchifier
|
||||
target_shape: VideoLatentShape | AudioLatentShape
|
||||
|
||||
def create_initial_state(
|
||||
self,
|
||||
device: DeviceLikeType,
|
||||
dtype: torch.dtype,
|
||||
initial_latent: torch.Tensor | None = None,
|
||||
) -> LatentState:
|
||||
"""
|
||||
Create an initial latent state. If initial_latent is provided, it will be used to create the latent state.
|
||||
"""
|
||||
...
|
||||
|
||||
def patchify(self, latent_state: LatentState) -> LatentState:
|
||||
"""
|
||||
Patchify the latent state.
|
||||
"""
|
||||
if latent_state.latent.shape != self.target_shape.to_torch_shape():
|
||||
raise ValueError(
|
||||
f"Latent state has shape {latent_state.latent.shape}, expected shape is "
|
||||
f"{self.target_shape.to_torch_shape()}"
|
||||
)
|
||||
latent_state = latent_state.clone()
|
||||
latent = self.patchifier.patchify(latent_state.latent)
|
||||
clean_latent = self.patchifier.patchify(latent_state.clean_latent)
|
||||
denoise_mask = self.patchifier.patchify(latent_state.denoise_mask)
|
||||
return replace(latent_state, latent=latent, denoise_mask=denoise_mask, clean_latent=clean_latent)
|
||||
|
||||
def unpatchify(self, latent_state: LatentState) -> LatentState:
|
||||
"""
|
||||
Unpatchify the latent state.
|
||||
"""
|
||||
latent_state = latent_state.clone()
|
||||
latent = self.patchifier.unpatchify(latent_state.latent, output_shape=self.target_shape)
|
||||
clean_latent = self.patchifier.unpatchify(latent_state.clean_latent, output_shape=self.target_shape)
|
||||
denoise_mask = self.patchifier.unpatchify(
|
||||
latent_state.denoise_mask, output_shape=self.target_shape.mask_shape()
|
||||
)
|
||||
return replace(latent_state, latent=latent, denoise_mask=denoise_mask, clean_latent=clean_latent)
|
||||
|
||||
def clear_conditioning(self, latent_state: LatentState) -> LatentState:
|
||||
"""
|
||||
Clear the conditioning from the latent state. This method removes extra tokens from the end of the latent.
|
||||
Therefore, conditioning items should add extra tokens ONLY to the end of the latent.
|
||||
"""
|
||||
latent_state = latent_state.clone()
|
||||
|
||||
num_tokens = self.patchifier.get_token_count(self.target_shape)
|
||||
latent = latent_state.latent[:, :num_tokens]
|
||||
clean_latent = latent_state.clean_latent[:, :num_tokens]
|
||||
denoise_mask = torch.ones_like(latent_state.denoise_mask)[:, :num_tokens]
|
||||
positions = latent_state.positions[:, :, :num_tokens]
|
||||
|
||||
return LatentState(latent=latent, denoise_mask=denoise_mask, positions=positions, clean_latent=clean_latent)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class VideoLatentTools(LatentTools):
|
||||
"""
|
||||
Tools for building video latent states.
|
||||
"""
|
||||
|
||||
patchifier: VideoLatentPatchifier
|
||||
target_shape: VideoLatentShape
|
||||
fps: float
|
||||
scale_factors: SpatioTemporalScaleFactors = DEFAULT_SCALE_FACTORS
|
||||
causal_fix: bool = True
|
||||
|
||||
def create_initial_state(
|
||||
self,
|
||||
device: DeviceLikeType,
|
||||
dtype: torch.dtype,
|
||||
initial_latent: torch.Tensor | None = None,
|
||||
) -> LatentState:
|
||||
if initial_latent is not None:
|
||||
assert initial_latent.shape == self.target_shape.to_torch_shape(), (
|
||||
f"Latent shape {initial_latent.shape} does not match target shape {self.target_shape.to_torch_shape()}"
|
||||
)
|
||||
else:
|
||||
initial_latent = torch.zeros(
|
||||
*self.target_shape.to_torch_shape(),
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
clean_latent = initial_latent.clone()
|
||||
|
||||
denoise_mask = torch.ones(
|
||||
*self.target_shape.mask_shape().to_torch_shape(),
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
latent_coords = self.patchifier.get_patch_grid_bounds(
|
||||
output_shape=self.target_shape,
|
||||
device=device,
|
||||
)
|
||||
|
||||
positions = get_pixel_coords(
|
||||
latent_coords=latent_coords,
|
||||
scale_factors=self.scale_factors,
|
||||
causal_fix=self.causal_fix,
|
||||
).float()
|
||||
positions[:, 0, ...] = positions[:, 0, ...] / self.fps
|
||||
|
||||
return self.patchify(
|
||||
LatentState(
|
||||
latent=initial_latent,
|
||||
denoise_mask=denoise_mask,
|
||||
positions=positions.to(dtype),
|
||||
clean_latent=clean_latent,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AudioLatentTools(LatentTools):
|
||||
"""
|
||||
Tools for building audio latent states.
|
||||
"""
|
||||
|
||||
patchifier: AudioPatchifier
|
||||
target_shape: AudioLatentShape
|
||||
|
||||
def create_initial_state(
|
||||
self,
|
||||
device: DeviceLikeType,
|
||||
dtype: torch.dtype,
|
||||
initial_latent: torch.Tensor | None = None,
|
||||
) -> LatentState:
|
||||
if initial_latent is not None:
|
||||
assert initial_latent.shape == self.target_shape.to_torch_shape(), (
|
||||
f"Latent shape {initial_latent.shape} does not match target shape {self.target_shape.to_torch_shape()}"
|
||||
)
|
||||
else:
|
||||
initial_latent = torch.zeros(
|
||||
*self.target_shape.to_torch_shape(),
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
clean_latent = initial_latent.clone()
|
||||
|
||||
denoise_mask = torch.ones(
|
||||
*self.target_shape.mask_shape().to_torch_shape(),
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
latent_coords = self.patchifier.get_patch_grid_bounds(
|
||||
output_shape=self.target_shape,
|
||||
device=device,
|
||||
)
|
||||
|
||||
return self.patchify(
|
||||
LatentState(
|
||||
latent=initial_latent, denoise_mask=denoise_mask, positions=latent_coords, clean_latent=clean_latent
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,181 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import NamedTuple
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class VideoPixelShape(NamedTuple):
|
||||
"""
|
||||
Shape of the tensor representing the video pixel array. Assumes BGR channel format.
|
||||
"""
|
||||
|
||||
batch: int
|
||||
frames: int
|
||||
height: int
|
||||
width: int
|
||||
fps: float
|
||||
|
||||
|
||||
class SpatioTemporalScaleFactors(NamedTuple):
|
||||
"""
|
||||
Describes the spatiotemporal downscaling between decoded video space and
|
||||
the corresponding VAE latent grid.
|
||||
"""
|
||||
|
||||
time: int
|
||||
width: int
|
||||
height: int
|
||||
|
||||
@classmethod
|
||||
def default(cls) -> "SpatioTemporalScaleFactors":
|
||||
return cls(time=8, width=32, height=32)
|
||||
|
||||
|
||||
VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
|
||||
|
||||
|
||||
class VideoLatentShape(NamedTuple):
|
||||
"""
|
||||
Shape of the tensor representing video in VAE latent space.
|
||||
The latent representation is a 5D tensor with dimensions ordered as
|
||||
(batch, channels, frames, height, width). Spatial and temporal dimensions
|
||||
are downscaled relative to pixel space according to the VAE's scale factors.
|
||||
"""
|
||||
|
||||
batch: int
|
||||
channels: int
|
||||
frames: int
|
||||
height: int
|
||||
width: int
|
||||
|
||||
def to_torch_shape(self) -> torch.Size:
|
||||
return torch.Size([self.batch, self.channels, self.frames, self.height, self.width])
|
||||
|
||||
@staticmethod
|
||||
def from_torch_shape(shape: torch.Size) -> "VideoLatentShape":
|
||||
return VideoLatentShape(
|
||||
batch=shape[0],
|
||||
channels=shape[1],
|
||||
frames=shape[2],
|
||||
height=shape[3],
|
||||
width=shape[4],
|
||||
)
|
||||
|
||||
def mask_shape(self) -> "VideoLatentShape":
|
||||
return self._replace(channels=1)
|
||||
|
||||
@staticmethod
|
||||
def from_pixel_shape(
|
||||
shape: VideoPixelShape,
|
||||
latent_channels: int = 128,
|
||||
scale_factors: SpatioTemporalScaleFactors = VIDEO_SCALE_FACTORS,
|
||||
) -> "VideoLatentShape":
|
||||
frames = (shape.frames - 1) // scale_factors[0] + 1
|
||||
height = shape.height // scale_factors[1]
|
||||
width = shape.width // scale_factors[2]
|
||||
|
||||
return VideoLatentShape(
|
||||
batch=shape.batch,
|
||||
channels=latent_channels,
|
||||
frames=frames,
|
||||
height=height,
|
||||
width=width,
|
||||
)
|
||||
|
||||
def upscale(self, scale_factors: SpatioTemporalScaleFactors = VIDEO_SCALE_FACTORS) -> "VideoLatentShape":
|
||||
return self._replace(
|
||||
channels=3,
|
||||
frames=(self.frames - 1) * scale_factors.time + 1,
|
||||
height=self.height * scale_factors.height,
|
||||
width=self.width * scale_factors.width,
|
||||
)
|
||||
|
||||
|
||||
class AudioLatentShape(NamedTuple):
|
||||
"""
|
||||
Shape of audio in VAE latent space: (batch, channels, frames, mel_bins).
|
||||
mel_bins is the number of frequency bins from the mel-spectrogram encoding.
|
||||
"""
|
||||
|
||||
batch: int
|
||||
channels: int
|
||||
frames: int
|
||||
mel_bins: int
|
||||
|
||||
def to_torch_shape(self) -> torch.Size:
|
||||
return torch.Size([self.batch, self.channels, self.frames, self.mel_bins])
|
||||
|
||||
def mask_shape(self) -> "AudioLatentShape":
|
||||
return self._replace(channels=1, mel_bins=1)
|
||||
|
||||
@staticmethod
|
||||
def from_torch_shape(shape: torch.Size) -> "AudioLatentShape":
|
||||
return AudioLatentShape(
|
||||
batch=shape[0],
|
||||
channels=shape[1],
|
||||
frames=shape[2],
|
||||
mel_bins=shape[3],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def from_duration(
|
||||
batch: int,
|
||||
duration: float,
|
||||
channels: int = 8,
|
||||
mel_bins: int = 16,
|
||||
sample_rate: int = 16000,
|
||||
hop_length: int = 160,
|
||||
audio_latent_downsample_factor: int = 4,
|
||||
) -> "AudioLatentShape":
|
||||
latents_per_second = float(sample_rate) / float(hop_length) / float(audio_latent_downsample_factor)
|
||||
|
||||
return AudioLatentShape(
|
||||
batch=batch,
|
||||
channels=channels,
|
||||
frames=round(duration * latents_per_second),
|
||||
mel_bins=mel_bins,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def from_video_pixel_shape(
|
||||
shape: VideoPixelShape,
|
||||
channels: int = 8,
|
||||
mel_bins: int = 16,
|
||||
sample_rate: int = 16000,
|
||||
hop_length: int = 160,
|
||||
audio_latent_downsample_factor: int = 4,
|
||||
) -> "AudioLatentShape":
|
||||
return AudioLatentShape.from_duration(
|
||||
batch=shape.batch,
|
||||
duration=float(shape.frames) / float(shape.fps),
|
||||
channels=channels,
|
||||
mel_bins=mel_bins,
|
||||
sample_rate=sample_rate,
|
||||
hop_length=hop_length,
|
||||
audio_latent_downsample_factor=audio_latent_downsample_factor,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LatentState:
|
||||
"""
|
||||
State of latents during the diffusion denoising process.
|
||||
Attributes:
|
||||
latent: The current noisy latent tensor being denoised.
|
||||
denoise_mask: Mask encoding the denoising strength for each token (1 = full denoising, 0 = no denoising).
|
||||
positions: Positional indices for each latent element, used for positional embeddings.
|
||||
clean_latent: Initial state of the latent before denoising, may include conditioning latents.
|
||||
"""
|
||||
|
||||
latent: torch.Tensor
|
||||
denoise_mask: torch.Tensor
|
||||
positions: torch.Tensor
|
||||
clean_latent: torch.Tensor
|
||||
|
||||
def clone(self) -> "LatentState":
|
||||
return LatentState(
|
||||
latent=self.latent.clone(),
|
||||
denoise_mask=self.denoise_mask.clone(),
|
||||
positions=self.positions.clone(),
|
||||
clean_latent=self.clean_latent.clone(),
|
||||
)
|
||||
@@ -0,0 +1,51 @@
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def rms_norm(x: torch.Tensor, weight: torch.Tensor | None = None, eps: float = 1e-6) -> torch.Tensor:
|
||||
"""Root-mean-square (RMS) normalize `x` over its last dimension.
|
||||
Thin wrapper around `torch.nn.functional.rms_norm` that infers the normalized
|
||||
shape and forwards `weight` and `eps`.
|
||||
"""
|
||||
return torch.nn.functional.rms_norm(x, (x.shape[-1],), weight=weight, eps=eps)
|
||||
|
||||
|
||||
def check_config_value(config: dict, key: str, expected: Any) -> None: # noqa: ANN401
|
||||
actual = config.get(key)
|
||||
if actual != expected:
|
||||
raise ValueError(f"Config value {key} is {actual}, expected {expected}")
|
||||
|
||||
|
||||
def to_velocity(
|
||||
sample: torch.Tensor,
|
||||
sigma: float | torch.Tensor,
|
||||
denoised_sample: torch.Tensor,
|
||||
calc_dtype: torch.dtype = torch.float32,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Convert the sample and its denoised version to velocity.
|
||||
Returns:
|
||||
Velocity
|
||||
"""
|
||||
if isinstance(sigma, torch.Tensor):
|
||||
sigma = sigma.to(calc_dtype).item()
|
||||
if sigma == 0:
|
||||
raise ValueError("Sigma can't be 0.0")
|
||||
return ((sample.to(calc_dtype) - denoised_sample.to(calc_dtype)) / sigma).to(sample.dtype)
|
||||
|
||||
|
||||
def to_denoised(
|
||||
sample: torch.Tensor,
|
||||
velocity: torch.Tensor,
|
||||
sigma: float | torch.Tensor,
|
||||
calc_dtype: torch.dtype = torch.float32,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Convert the sample and its denoising velocity to denoised sample.
|
||||
Returns:
|
||||
Denoised sample
|
||||
"""
|
||||
if isinstance(sigma, torch.Tensor):
|
||||
sigma = sigma.to(calc_dtype)
|
||||
return (sample.to(calc_dtype) - velocity.to(calc_dtype) * sigma).to(sample.dtype)
|
||||
Reference in New Issue
Block a user