Automated PR - 2026-01-29
This commit is contained in:
sync-bot
@@ -200,6 +200,12 @@ validation:
|
||||
- "/path/to/reference_video_1.mp4"
|
||||
- "/path/to/reference_video_2.mp4"
|
||||
|
||||
# Downscale factor for reference videos (for efficient IC-LoRA training)
|
||||
# When > 1, reference videos are processed at 1/n resolution
|
||||
# Must match the --reference-downscale-factor used during dataset preprocessing
|
||||
# Examples: 1 = same resolution, 2 = half resolution (384x384 ref for 768x768 target)
|
||||
reference_downscale_factor: 1
|
||||
|
||||
# Negative prompt to avoid unwanted artifacts
|
||||
negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
|
||||
|
||||
|
||||
@@ -111,7 +111,8 @@ IC-LoRA enables a wide range of advanced video-to-video applications, such as:
|
||||
- **Colorization**: Convert grayscale reference videos into colorized outputs
|
||||
- **Restoration and enhancement**: Denoise, upscale, or restore old or degraded videos
|
||||
|
||||
By providing paired reference and target videos, IC-LoRA can learn complex transformations that go beyond caption-based conditioning.
|
||||
By providing paired reference and target videos, IC-LoRA can learn complex transformations that go beyond caption-based
|
||||
conditioning.
|
||||
|
||||
IC-LoRA training fundamentally differs from standard LoRA and full fine-tuning:
|
||||
|
||||
@@ -140,8 +141,10 @@ training_strategy:
|
||||
### Dataset Requirements for IC-LoRA
|
||||
|
||||
- Your dataset must contain **paired videos** where each target video has a corresponding reference video
|
||||
- Reference and target videos must have **identical resolution and length**
|
||||
- Both reference and target videos should be **preprocessed together** using the same resolution buckets
|
||||
- Reference and target videos must have the **same frame count** (length)
|
||||
- Reference videos can optionally be at **lower spatial resolution** than target videos (
|
||||
see [Scaled Reference Conditioning](#scaled-reference-conditioning) below)
|
||||
- Both reference and target videos should be **preprocessed** before training
|
||||
|
||||
**Dataset structure for IC-LoRA training:**
|
||||
|
||||
@@ -181,30 +184,82 @@ validation:
|
||||
reference_videos:
|
||||
- "/path/to/reference1.mp4"
|
||||
- "/path/to/reference2.mp4"
|
||||
reference_downscale_factor: 1 # Set to match preprocessing (e.g., 2 for half resolution)
|
||||
include_reference_in_output: true # Show reference side-by-side with output
|
||||
```
|
||||
|
||||
### Scaled Reference Conditioning
|
||||
|
||||
For more efficient training and inference, you can use **downscaled reference videos** while keeping target videos at
|
||||
full resolution. This reduces the number of conditioning tokens, leading to:
|
||||
|
||||
- **Faster training** due to shorter sequence lengths
|
||||
- **Faster inference** with reduced memory usage
|
||||
- **Same aspect ratio** maintained between reference and target
|
||||
|
||||
#### How It Works
|
||||
|
||||
When the reference video has resolution `H/n × W/n` and the target video has resolution `H × W`, the trainer
|
||||
automatically detects this scale factor `n` and adjusts the positional encodings so that the reference positions
|
||||
map to the correct locations in the target coordinate space.
|
||||
|
||||
#### Preprocessing Datasets with Scaled References
|
||||
|
||||
Use the `--reference-downscale-factor` option when running `process_dataset.py`:
|
||||
|
||||
```bash
|
||||
# Process dataset with scaled reference videos (half resolution)
|
||||
uv run python scripts/process_dataset.py dataset.json \
|
||||
--resolution-buckets 768x768x25 \
|
||||
--model-path /path/to/ltx2.safetensors \
|
||||
--text-encoder-path /path/to/gemma \
|
||||
--reference-column "reference_path" \
|
||||
--reference-downscale-factor 2
|
||||
```
|
||||
|
||||
This will:
|
||||
|
||||
- Process target videos at 768×768 resolution
|
||||
- Process reference videos at 384×384 resolution (768 / 2)
|
||||
- The trainer will automatically infer the scale factor from the dimension ratio
|
||||
|
||||
**Important**: Set `reference_downscale_factor: 2` in your validation configuration to match the preprocessing:
|
||||
|
||||
```yaml
|
||||
validation:
|
||||
reference_downscale_factor: 2 # Must match the preprocessing factor
|
||||
reference_videos:
|
||||
- "/path/to/reference1.mp4"
|
||||
- "/path/to/reference2.mp4"
|
||||
```
|
||||
|
||||
> [!NOTE]
|
||||
> The scale factor must be a positive integer, and all dimensions must be divisible by 32.
|
||||
> Common scale factors are 1 (no scaling), 2 (half resolution), or 4 (quarter resolution).
|
||||
|
||||
## 📊 Training Mode Comparison
|
||||
|
||||
| Aspect | LoRA | Audio-Video LoRA | Full Fine-tuning | IC-LoRA |
|
||||
|----------------------|------------|------------------|------------------|----------------|
|
||||
| **Memory Usage** | Low | Low-Medium | High | Medium |
|
||||
| **Training Speed** | Fast | Fast | Slow | Medium |
|
||||
| **Output Size** | 100MB-few GB (depends on rank) | 100MB-few GB (depends on rank) | Tens of GB | 100MB-few GB (depends on rank) |
|
||||
| **Flexibility** | Medium | Medium | High | Specialized |
|
||||
| **Audio Support** | Optional | Yes | Optional | No |
|
||||
| **Reference Videos** | No | No | No | Yes (required) |
|
||||
| Aspect | LoRA | Audio-Video LoRA | Full Fine-tuning | IC-LoRA |
|
||||
|----------------------|--------------------------------|--------------------------------|------------------|--------------------------------|
|
||||
| **Memory Usage** | Low | Low-Medium | High | Medium |
|
||||
| **Training Speed** | Fast | Fast | Slow | Medium |
|
||||
| **Output Size** | 100MB-few GB (depends on rank) | 100MB-few GB (depends on rank) | Tens of GB | 100MB-few GB (depends on rank) |
|
||||
| **Flexibility** | Medium | Medium | High | Specialized |
|
||||
| **Audio Support** | Optional | Yes | Optional | No |
|
||||
| **Reference Videos** | No | No | No | Yes (required) |
|
||||
|
||||
## 🎬 Using Trained Models for Inference
|
||||
|
||||
After training, use the [`ltx-pipelines`](../../ltx-pipelines/) package for production inference with your trained LoRAs:
|
||||
After training, use the [`ltx-pipelines`](../../ltx-pipelines/) package for production inference with your trained
|
||||
LoRAs:
|
||||
|
||||
| Training Mode | Recommended Pipeline |
|
||||
|---------------|---------------------|
|
||||
| Training Mode | Recommended Pipeline |
|
||||
|-------------------------|-------------------------------------------------------|
|
||||
| LoRA / Audio-Video LoRA | `TI2VidOneStagePipeline` or `TI2VidTwoStagesPipeline` |
|
||||
| IC-LoRA | `ICLoraPipeline` |
|
||||
| IC-LoRA | `ICLoraPipeline` |
|
||||
|
||||
All pipelines support loading custom LoRAs via the `loras` parameter. See the [`ltx-pipelines`](../../ltx-pipelines/) package
|
||||
All pipelines support loading custom LoRAs via the `loras` parameter. See the [`ltx-pipelines`](../../ltx-pipelines/)
|
||||
package
|
||||
documentation for detailed usage instructions.
|
||||
|
||||
## 🚀 Next Steps
|
||||
@@ -216,6 +271,7 @@ Once you've chosen your training mode:
|
||||
- Start training with the [Training Guide](training-guide.md)
|
||||
|
||||
> [!TIP]
|
||||
> Need a training mode that's not covered here? See [Implementing Custom Training Strategies](custom-training-strategies.md)
|
||||
> Need a training mode that's not covered here?
|
||||
> See [Implementing Custom Training Strategies](custom-training-strategies.md)
|
||||
> to learn how to create your own strategy for specialized use cases like video inpainting, audio-only training, or
|
||||
> custom conditioning.
|
||||
|
||||
@@ -16,13 +16,14 @@ from pathlib import Path
|
||||
import typer
|
||||
from decode_latents import LatentsDecoder
|
||||
from process_captions import compute_captions_embeddings
|
||||
from process_videos import compute_latents, parse_resolution_buckets
|
||||
from process_videos import compute_latents, compute_scaled_resolution_buckets, parse_resolution_buckets
|
||||
from rich.console import Console
|
||||
|
||||
from ltx_trainer import logger
|
||||
from ltx_trainer.gpu_utils import free_gpu_memory_context
|
||||
|
||||
console = Console()
|
||||
|
||||
app = typer.Typer(
|
||||
pretty_exceptions_enable=False,
|
||||
no_args_is_help=True,
|
||||
@@ -46,6 +47,7 @@ def preprocess_dataset( # noqa: PLR0913
|
||||
device: str,
|
||||
remove_llm_prefixes: bool = False,
|
||||
reference_column: str | None = None,
|
||||
reference_downscale_factor: int = 1,
|
||||
with_audio: bool = False,
|
||||
load_text_encoder_in_8bit: bool = False,
|
||||
) -> None:
|
||||
@@ -99,14 +101,33 @@ def preprocess_dataset( # noqa: PLR0913
|
||||
|
||||
# Process reference videos if reference_column is provided
|
||||
if reference_column:
|
||||
logger.info("Processing reference videos for IC-LoRA training...")
|
||||
# Validate: scaled references with multiple buckets can cause ambiguous bucket matching
|
||||
if reference_downscale_factor > 1 and len(resolution_buckets) > 1:
|
||||
raise ValueError(
|
||||
"When using --reference-downscale-factor > 1, only a single resolution bucket is supported. "
|
||||
"Using multiple buckets with scaled references can cause ambiguous bucket matching "
|
||||
"(e.g., a 512x256 reference could match either the scaled-down 1024x512 bucket or the 512x256 "
|
||||
"bucket). Please use a single resolution bucket or set --reference-downscale-factor to 1."
|
||||
)
|
||||
|
||||
# Calculate and validate scaled resolution buckets for reference videos
|
||||
reference_buckets = compute_scaled_resolution_buckets(resolution_buckets, reference_downscale_factor)
|
||||
|
||||
if reference_downscale_factor > 1:
|
||||
logger.info(
|
||||
f"Processing reference videos for IC-LoRA training at 1/{reference_downscale_factor} resolution..."
|
||||
)
|
||||
logger.info(f"Reference resolution buckets: {reference_buckets}")
|
||||
else:
|
||||
logger.info("Processing reference videos for IC-LoRA training...")
|
||||
|
||||
reference_latents_dir = output_base / "reference_latents"
|
||||
|
||||
compute_latents(
|
||||
dataset_file=dataset_file,
|
||||
main_media_column=video_column,
|
||||
video_column=reference_column,
|
||||
resolution_buckets=resolution_buckets,
|
||||
resolution_buckets=reference_buckets,
|
||||
output_dir=str(reference_latents_dir),
|
||||
model_path=model_path,
|
||||
batch_size=batch_size,
|
||||
@@ -226,6 +247,11 @@ def main( # noqa: PLR0913
|
||||
default=False,
|
||||
help="Load the Gemma text encoder in 8-bit precision to save GPU memory (requires bitsandbytes)",
|
||||
),
|
||||
reference_downscale_factor: int = typer.Option(
|
||||
default=1,
|
||||
help="Downscale factor for reference video resolution. When > 1, reference videos are processed at "
|
||||
"1/n resolution (e.g., 2 means half resolution). Used for efficient IC-LoRA training.",
|
||||
),
|
||||
) -> None:
|
||||
"""Preprocess a video dataset by computing and saving latents and text embeddings.
|
||||
The dataset must be a CSV, JSON, or JSONL file with columns for captions and video paths.
|
||||
@@ -242,6 +268,10 @@ def main( # noqa: PLR0913
|
||||
python scripts/process_dataset.py dataset.json --resolution-buckets 768x768x25 \\
|
||||
--model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma \\
|
||||
--reference-column "reference_path"
|
||||
# Process dataset with scaled reference videos (half resolution) for efficient IC-LoRA
|
||||
python scripts/process_dataset.py dataset.json --resolution-buckets 768x768x25 \\
|
||||
--model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma \\
|
||||
--reference-column "reference_path" --reference-downscale-factor 2
|
||||
# Process dataset with audio for audio-video training
|
||||
python scripts/process_dataset.py dataset.json --resolution-buckets 768x512x97 \\
|
||||
--model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma \\
|
||||
@@ -255,6 +285,13 @@ def main( # noqa: PLR0913
|
||||
"When training with multiple resolution buckets, you must use a batch size of 1."
|
||||
)
|
||||
|
||||
# Validate reference_downscale_factor
|
||||
if reference_downscale_factor < 1:
|
||||
raise typer.BadParameter("--reference-downscale-factor must be >= 1")
|
||||
|
||||
if reference_downscale_factor > 1 and not reference_column:
|
||||
logger.warning("--reference-downscale-factor specified but no --reference-column provided. Ignoring.")
|
||||
|
||||
preprocess_dataset(
|
||||
dataset_file=dataset_path,
|
||||
caption_column=caption_column,
|
||||
@@ -270,6 +307,7 @@ def main( # noqa: PLR0913
|
||||
device=device,
|
||||
remove_llm_prefixes=remove_llm_prefixes,
|
||||
reference_column=reference_column,
|
||||
reference_downscale_factor=reference_downscale_factor,
|
||||
with_audio=with_audio,
|
||||
load_text_encoder_in_8bit=load_text_encoder_in_8bit,
|
||||
)
|
||||
|
||||
@@ -891,6 +891,50 @@ def parse_resolution_buckets(resolution_buckets_str: str) -> list[tuple[int, int
|
||||
return resolution_buckets
|
||||
|
||||
|
||||
def compute_scaled_resolution_buckets(
|
||||
resolution_buckets: list[tuple[int, int, int]],
|
||||
scale_factor: int,
|
||||
) -> list[tuple[int, int, int]]:
|
||||
"""Compute scaled resolution buckets and validate the results."""
|
||||
if scale_factor == 1:
|
||||
return resolution_buckets
|
||||
|
||||
scaled_buckets = []
|
||||
for frames, height, width in resolution_buckets:
|
||||
# Validate that scale factor evenly divides the dimensions
|
||||
if height % scale_factor != 0:
|
||||
raise ValueError(
|
||||
f"Height {height} is not evenly divisible by scale factor {scale_factor}. "
|
||||
f"Choose a scale factor that divides {height} evenly."
|
||||
)
|
||||
if width % scale_factor != 0:
|
||||
raise ValueError(
|
||||
f"Width {width} is not evenly divisible by scale factor {scale_factor}. "
|
||||
f"Choose a scale factor that divides {width} evenly."
|
||||
)
|
||||
|
||||
scaled_height = height // scale_factor
|
||||
scaled_width = width // scale_factor
|
||||
|
||||
# Validate scaled dimensions are divisible by VAE spatial factor
|
||||
if scaled_height % VAE_SPATIAL_FACTOR != 0:
|
||||
raise ValueError(
|
||||
f"Scaled height {scaled_height} (from {height} / {scale_factor}) "
|
||||
f"is not divisible by {VAE_SPATIAL_FACTOR}. "
|
||||
f"Choose a different scale factor or adjust your resolution buckets."
|
||||
)
|
||||
if scaled_width % VAE_SPATIAL_FACTOR != 0:
|
||||
raise ValueError(
|
||||
f"Scaled width {scaled_width} (from {width} / {scale_factor}) "
|
||||
f"is not divisible by {VAE_SPATIAL_FACTOR}. "
|
||||
f"Choose a different scale factor or adjust your resolution buckets."
|
||||
)
|
||||
|
||||
scaled_buckets.append((frames, scaled_height, scaled_width))
|
||||
|
||||
return scaled_buckets
|
||||
|
||||
|
||||
@app.command()
|
||||
def main( # noqa: PLR0913
|
||||
dataset_file: str = typer.Argument(
|
||||
|
||||
@@ -207,6 +207,14 @@ class ValidationConfig(ConfigBaseModel):
|
||||
"One video path must be provided for each validation prompt",
|
||||
)
|
||||
|
||||
reference_downscale_factor: int = Field(
|
||||
default=1,
|
||||
description="Downscale factor for reference videos in IC-LoRA validation. "
|
||||
"When > 1, reference videos are processed at 1/n resolution (e.g., 2 means half resolution). "
|
||||
"Must match the factor used during dataset preprocessing.",
|
||||
ge=1,
|
||||
)
|
||||
|
||||
video_dims: tuple[int, int, int] = Field(
|
||||
default=(960, 544, 97),
|
||||
description="Dimensions of validation videos (width, height, frames). "
|
||||
@@ -334,6 +342,41 @@ class ValidationConfig(ConfigBaseModel):
|
||||
|
||||
return v
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_scaled_reference_dimensions(self) -> "ValidationConfig":
|
||||
"""Validate that scaled reference dimensions are valid when reference_downscale_factor > 1."""
|
||||
if self.reference_downscale_factor > 1:
|
||||
width, height, _frames = self.video_dims
|
||||
|
||||
# Validate that downscale factor evenly divides the target dimensions
|
||||
if width % self.reference_downscale_factor != 0:
|
||||
raise ValueError(
|
||||
f"Width {width} is not evenly divisible by reference_downscale_factor "
|
||||
f"{self.reference_downscale_factor}. Choose a downscale factor that divides {width} evenly."
|
||||
)
|
||||
if height % self.reference_downscale_factor != 0:
|
||||
raise ValueError(
|
||||
f"Height {height} is not evenly divisible by reference_downscale_factor "
|
||||
f"{self.reference_downscale_factor}. Choose a downscale factor that divides {height} evenly."
|
||||
)
|
||||
|
||||
scaled_width = width // self.reference_downscale_factor
|
||||
scaled_height = height // self.reference_downscale_factor
|
||||
|
||||
# Validate scaled dimensions are divisible by 32
|
||||
if scaled_width % 32 != 0:
|
||||
raise ValueError(
|
||||
f"Scaled reference width {scaled_width} (from {width} / {self.reference_downscale_factor}) "
|
||||
f"is not divisible by 32. Choose a different downscale factor or adjust video_dims."
|
||||
)
|
||||
if scaled_height % 32 != 0:
|
||||
raise ValueError(
|
||||
f"Scaled reference height {scaled_height} (from {height} / {self.reference_downscale_factor}) "
|
||||
f"is not divisible by 32. Choose a different downscale factor or adjust video_dims."
|
||||
)
|
||||
|
||||
return self
|
||||
|
||||
|
||||
class CheckpointsConfig(ConfigBaseModel):
|
||||
"""Configuration for model checkpointing during training"""
|
||||
|
||||
@@ -218,16 +218,22 @@ def load_text_encoder(
|
||||
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
|
||||
from ltx_core.text_encoders.gemma.encoders.av_encoder import (
|
||||
AV_GEMMA_TEXT_ENCODER_KEY_OPS,
|
||||
GEMMA_MODEL_OPS,
|
||||
AVGemmaTextEncoderModelConfigurator,
|
||||
)
|
||||
from ltx_core.text_encoders.gemma.encoders.base_encoder import module_ops_from_gemma_root
|
||||
from ltx_core.utils import find_matching_file
|
||||
|
||||
torch_device = _to_torch_device(device)
|
||||
|
||||
gemma_model_folder = find_matching_file(str(gemma_model_path), "model*.safetensors").parent
|
||||
gemma_weight_paths = [str(p) for p in gemma_model_folder.rglob("*.safetensors")]
|
||||
|
||||
text_encoder = SingleGPUModelBuilder(
|
||||
model_path=str(checkpoint_path),
|
||||
model_path=(str(checkpoint_path), *gemma_weight_paths),
|
||||
model_class_configurator=AVGemmaTextEncoderModelConfigurator,
|
||||
model_sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS,
|
||||
module_ops=module_ops_from_gemma_root(str(gemma_model_path)),
|
||||
module_ops=(GEMMA_MODEL_OPS, *module_ops_from_gemma_root(str(gemma_model_path))),
|
||||
).build(device=torch_device, dtype=dtype)
|
||||
|
||||
return text_encoder
|
||||
|
||||
@@ -795,6 +795,7 @@ class LtxvTrainer:
|
||||
seed=self._config.validation.seed,
|
||||
condition_image=condition_image,
|
||||
reference_video=reference_video,
|
||||
reference_downscale_factor=self._config.validation.reference_downscale_factor,
|
||||
generate_audio=generate_audio,
|
||||
include_reference_in_output=self._config.validation.include_reference_in_output,
|
||||
cached_embeddings=cached_embeddings,
|
||||
@@ -885,8 +886,11 @@ class LtxvTrainer:
|
||||
# Cast to configured precision
|
||||
state_dict = {k: v.to(save_dtype) if isinstance(v, Tensor) else v for k, v in state_dict.items()}
|
||||
|
||||
# Save to disk
|
||||
save_file(state_dict, saved_weights_path)
|
||||
# Build metadata for safetensors file
|
||||
metadata = self._build_checkpoint_metadata()
|
||||
|
||||
# Save to disk with metadata
|
||||
save_file(state_dict, saved_weights_path, metadata=metadata)
|
||||
else:
|
||||
# Cast to configured precision
|
||||
full_state_dict = {k: v.to(save_dtype) if isinstance(v, Tensor) else v for k, v in full_state_dict.items()}
|
||||
@@ -913,6 +917,21 @@ class LtxvTrainer:
|
||||
# Update the list to only contain kept checkpoints
|
||||
self._checkpoint_paths = self._checkpoint_paths[-self._config.checkpoints.keep_last_n :]
|
||||
|
||||
def _build_checkpoint_metadata(self) -> dict[str, str]:
|
||||
"""Build metadata dictionary for safetensors checkpoint.
|
||||
Delegates to the training strategy to get strategy-specific metadata
|
||||
that downstream inference pipelines may need.
|
||||
Returns:
|
||||
Dictionary of string key-value pairs for safetensors metadata.
|
||||
Values are converted to strings for safetensors compatibility.
|
||||
"""
|
||||
raw_metadata = self._training_strategy.get_checkpoint_metadata()
|
||||
# Convert all values to strings for safetensors compatibility
|
||||
metadata = {k: str(v) for k, v in raw_metadata.items()}
|
||||
if metadata:
|
||||
logger.info(f"Saving checkpoint metadata: {metadata}")
|
||||
return metadata
|
||||
|
||||
def _save_config(self) -> None:
|
||||
"""Save the training configuration as a YAML file in the output directory."""
|
||||
if not IS_MAIN_PROCESS:
|
||||
|
||||
@@ -128,6 +128,15 @@ class TrainingStrategy(ABC):
|
||||
Scalar loss tensor
|
||||
"""
|
||||
|
||||
def get_checkpoint_metadata(self) -> dict[str, Any]:
|
||||
"""Get strategy-specific metadata to include in checkpoint files.
|
||||
Override this method in subclasses to add custom metadata,
|
||||
e.g. any parameters that a downstream inference pipeline may need.
|
||||
Returns:
|
||||
Dictionary of metadata key-value pairs (values must be JSON-serializable)
|
||||
"""
|
||||
return {}
|
||||
|
||||
def _get_video_positions(
|
||||
self,
|
||||
num_frames: int,
|
||||
|
||||
@@ -46,9 +46,13 @@ class VideoToVideoStrategy(TrainingStrategy):
|
||||
- Reference latents (clean) are concatenated with target latents (noised)
|
||||
- Video coordinates handle both reference and target sequences
|
||||
- Loss is computed only on the target portion
|
||||
Attributes:
|
||||
reference_downscale_factor: The inferred downscale factor of reference videos.
|
||||
This is computed from the first batch and cached for metadata export.
|
||||
"""
|
||||
|
||||
config: VideoToVideoConfig
|
||||
reference_downscale_factor: int | None
|
||||
|
||||
def __init__(self, config: VideoToVideoConfig):
|
||||
"""Initialize strategy with configuration.
|
||||
@@ -56,6 +60,7 @@ class VideoToVideoStrategy(TrainingStrategy):
|
||||
config: Video-to-video configuration
|
||||
"""
|
||||
super().__init__(config)
|
||||
self.reference_downscale_factor = None # Will be inferred from first batch
|
||||
|
||||
def get_data_sources(self) -> dict[str, str]:
|
||||
"""IC-LoRA training requires latents, conditions, and reference latents."""
|
||||
@@ -65,7 +70,7 @@ class VideoToVideoStrategy(TrainingStrategy):
|
||||
self.config.reference_latents_dir: "ref_latents",
|
||||
}
|
||||
|
||||
def prepare_training_inputs(
|
||||
def prepare_training_inputs( # noqa: PLR0915
|
||||
self,
|
||||
batch: dict[str, Any],
|
||||
timestep_sampler: TimestepSampler,
|
||||
@@ -86,6 +91,26 @@ class VideoToVideoStrategy(TrainingStrategy):
|
||||
ref_height = ref_latents_info["height"][0].item()
|
||||
ref_width = ref_latents_info["width"][0].item()
|
||||
|
||||
# Infer reference downscale factor from dimension ratios
|
||||
# This allows training with downscaled reference videos for efficiency
|
||||
reference_downscale_factor = self._infer_reference_downscale_factor(
|
||||
target_height=height,
|
||||
target_width=width,
|
||||
ref_height=ref_height,
|
||||
ref_width=ref_width,
|
||||
)
|
||||
|
||||
# Cache the scale factor for metadata export (only on first batch)
|
||||
if self.reference_downscale_factor is None:
|
||||
self.reference_downscale_factor = reference_downscale_factor
|
||||
elif self.reference_downscale_factor != reference_downscale_factor:
|
||||
raise ValueError(
|
||||
f"Inconsistent reference downscale factor across batches. "
|
||||
f"First batch had factor={self.reference_downscale_factor}, "
|
||||
f"but current batch has factor={reference_downscale_factor}. "
|
||||
f"All training samples must use the same reference/target resolution ratio."
|
||||
)
|
||||
|
||||
# Patchify latents: [B, C, F, H, W] -> [B, seq_len, C]
|
||||
target_latents = self._video_patchifier.patchify(target_latents)
|
||||
ref_latents = self._video_patchifier.patchify(ref_latents)
|
||||
@@ -159,6 +184,15 @@ class VideoToVideoStrategy(TrainingStrategy):
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
# Scale reference positions to match target coordinate space
|
||||
# This maps ref positions from (0, ref_H, ref_W) to (0, target_H, target_W)
|
||||
# Position tensor shape: [B, 3, seq_len, 2] where dim 1 is (time, height, width)
|
||||
if reference_downscale_factor != 1:
|
||||
ref_positions = ref_positions.clone()
|
||||
ref_positions[:, 1, ...] *= reference_downscale_factor # height axis
|
||||
ref_positions[:, 2, ...] *= reference_downscale_factor # width axis
|
||||
# Time axis (index 0) remains unchanged
|
||||
|
||||
target_positions = self._get_video_positions(
|
||||
num_frames=num_frames,
|
||||
height=height,
|
||||
@@ -221,3 +255,48 @@ class VideoToVideoStrategy(TrainingStrategy):
|
||||
loss = loss.mul(loss_mask).div(loss_mask.mean())
|
||||
|
||||
return loss.mean()
|
||||
|
||||
def get_checkpoint_metadata(self) -> dict[str, Any]:
|
||||
"""Get metadata for checkpoint files."""
|
||||
metadata: dict[str, Any] = {}
|
||||
# Always include reference_downscale_factor for IC-LoRAs so inference
|
||||
# pipelines know the expected scale factor for reference videos.
|
||||
if self.reference_downscale_factor is not None:
|
||||
metadata["reference_downscale_factor"] = self.reference_downscale_factor
|
||||
return metadata
|
||||
|
||||
@staticmethod
|
||||
def _infer_reference_downscale_factor(
|
||||
target_height: int,
|
||||
target_width: int,
|
||||
ref_height: int,
|
||||
ref_width: int,
|
||||
) -> int:
|
||||
"""Infer the reference downscale factor from target and reference dimensions."""
|
||||
# If dimensions match, no scaling needed
|
||||
if target_height == ref_height and target_width == ref_width:
|
||||
return 1
|
||||
|
||||
# Calculate scale factors for each dimension
|
||||
if target_height % ref_height != 0 or target_width % ref_width != 0:
|
||||
raise ValueError(
|
||||
f"Target dimensions ({target_height}x{target_width}) must be exact multiples "
|
||||
f"of reference dimensions ({ref_height}x{ref_width})"
|
||||
)
|
||||
|
||||
scale_h = target_height // ref_height
|
||||
scale_w = target_width // ref_width
|
||||
|
||||
if scale_h != scale_w:
|
||||
raise ValueError(
|
||||
f"Reference scale must be uniform. Got height scale {scale_h} and width scale {scale_w}. "
|
||||
f"Target: {target_height}x{target_width}, Reference: {ref_height}x{ref_width}"
|
||||
)
|
||||
|
||||
if scale_h < 1:
|
||||
raise ValueError(
|
||||
f"Reference dimensions ({ref_height}x{ref_width}) cannot be larger than "
|
||||
f"target dimensions ({target_height}x{target_width})"
|
||||
)
|
||||
|
||||
return scale_h
|
||||
|
||||
@@ -85,6 +85,7 @@ class GenerationConfig:
|
||||
seed: int = 42 # Random seed for reproducibility
|
||||
condition_image: Tensor | None = None # Optional first frame image for image-to-video
|
||||
reference_video: Tensor | None = None # For IC-LoRA: [F, C, H, W] in [0, 1]
|
||||
reference_downscale_factor: int = 1 # For IC-LoRA: downscale factor (1 = same resolution, 2 = half resolution)
|
||||
generate_audio: bool = True # Whether to generate audio alongside video
|
||||
include_reference_in_output: bool = False # For IC-LoRA: concatenate original reference with generated output
|
||||
cached_embeddings: CachedPromptEmbeddings | None = None # Pre-computed text embeddings (avoids loading Gemma)
|
||||
@@ -251,6 +252,14 @@ class ValidationSampler:
|
||||
ref_latent, ref_positions = self._encode_video(ref_video_preprocessed, config.frame_rate, device)
|
||||
ref_seq_len = ref_latent.shape[1]
|
||||
|
||||
# Scale reference positions to match target coordinate space
|
||||
# Position tensor shape: [B, 3, seq_len, 2] where dim 1 is (time, height, width)
|
||||
if config.reference_downscale_factor != 1:
|
||||
ref_positions = ref_positions.clone()
|
||||
ref_positions[:, 1, ...] *= config.reference_downscale_factor # height axis
|
||||
ref_positions[:, 2, ...] *= config.reference_downscale_factor # width axis
|
||||
# Time axis (index 0) remains unchanged
|
||||
|
||||
# Create target video state
|
||||
video_tools = self._create_video_latent_tools(config)
|
||||
target_clean_state = video_tools.create_initial_state(device=device, dtype=torch.bfloat16)
|
||||
@@ -375,13 +384,28 @@ class ValidationSampler:
|
||||
@staticmethod
|
||||
def _preprocess_reference_video(config: GenerationConfig) -> Tensor:
|
||||
"""Preprocess reference video: resize, crop, and convert to model input format.
|
||||
When reference_downscale_factor > 1, the reference video is downscaled to a smaller
|
||||
resolution for more efficient inference. The positions will be scaled up later
|
||||
to match the target coordinate space.
|
||||
Args:
|
||||
config: Generation configuration with reference_video
|
||||
config: Generation configuration
|
||||
Returns:
|
||||
Preprocessed video tensor [B, C, F, H, W] in [-1, 1] range
|
||||
"""
|
||||
ref_video = config.reference_video # [F, C, H, W] in [0, 1]
|
||||
target_height, target_width = config.height, config.width
|
||||
scale_factor = config.reference_downscale_factor
|
||||
|
||||
# Target dimensions for reference (scaled down if scale_factor > 1)
|
||||
target_height = config.height // scale_factor
|
||||
target_width = config.width // scale_factor
|
||||
|
||||
# Validate scaled dimensions
|
||||
if target_height % 32 != 0 or target_width % 32 != 0:
|
||||
raise ValueError(
|
||||
f"Scaled reference dimensions ({target_height}x{target_width}) must be divisible by 32. "
|
||||
f"Original: {config.height}x{config.width}, scale_factor: {scale_factor}"
|
||||
)
|
||||
|
||||
current_height, current_width = ref_video.shape[2:]
|
||||
|
||||
# Resize maintaining aspect ratio and center crop if needed
|
||||
@@ -745,11 +769,28 @@ class ValidationSampler:
|
||||
If the videos have different frame counts, the shorter one is padded with
|
||||
its last frame repeated.
|
||||
Args:
|
||||
left_video: Left video tensor [C, F1, H, W] in [0, 1]
|
||||
right_video: Right video tensor [C, F2, H, W] in [0, 1]
|
||||
left_video: Left video tensor [C, F1, H1, W1] in [0, 1]
|
||||
right_video: Right video tensor [C, F2, H2, W2] in [0, 1]
|
||||
Returns:
|
||||
Concatenated video tensor [C, max(F1,F2), H, W*2] in [0, 1]
|
||||
Concatenated video tensor [C, max(F1,F2), H2, W1_scaled+W2] in [0, 1]
|
||||
"""
|
||||
left_height, left_width = left_video.shape[2], left_video.shape[3]
|
||||
right_height = right_video.shape[2]
|
||||
|
||||
# Resize left video to match right video's height if needed
|
||||
if left_height != right_height:
|
||||
# Scale width proportionally to maintain aspect ratio
|
||||
scale = right_height / left_height
|
||||
new_width = int(left_width * scale)
|
||||
# Interpolate expects [N, C, H, W], we have [C, F, H, W]
|
||||
# Reshape to [C*F, 1, H, W] -> interpolate -> reshape back
|
||||
c, f, h, w = left_video.shape
|
||||
left_video = left_video.reshape(c * f, 1, h, w)
|
||||
left_video = torch.nn.functional.interpolate(
|
||||
left_video, size=(right_height, new_width), mode="bilinear", align_corners=False
|
||||
)
|
||||
left_video = left_video.reshape(c, f, right_height, new_width)
|
||||
|
||||
left_frames = left_video.shape[1]
|
||||
right_frames = right_video.shape[1]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user