Automated PR - 2026-03-05

This commit is contained in:
sync-bot
2026-03-05 15:47:20 +00:00
parent 3b6d09d7b6
commit d230aec5cd
29 changed files with 739 additions and 540 deletions
@@ -1,10 +1,11 @@
from ltx_pipelines.utils.helpers import (
assert_resolution,
cleanup_memory,
combined_image_conditionings,
denoise_audio_video,
encode_prompts,
generate_enhanced_prompt,
get_device,
image_conditionings_by_replacing_latent,
multi_modal_guider_denoising_func,
multi_modal_guider_factory_denoising_func,
simple_denoising_func,
@@ -20,12 +21,13 @@ __all__ = [
"ModelLedger",
"assert_resolution",
"cleanup_memory",
"combined_image_conditionings",
"denoise_audio_video",
"encode_prompts",
"euler_denoising_loop",
"generate_enhanced_prompt",
"get_device",
"gradient_estimating_euler_denoising_loop",
"image_conditionings_by_replacing_latent",
"multi_modal_guider_denoising_func",
"multi_modal_guider_factory_denoising_func",
"res2s_audio_video_denoising_loop",
@@ -8,6 +8,7 @@ from ltx_pipelines.utils.constants import (
DEFAULT_IMAGE_CRF,
DEFAULT_LORA_STRENGTH,
DEFAULT_NEGATIVE_PROMPT,
LTX_2_3_HQ_PARAMS,
LTX_2_3_PARAMS,
PipelineParams,
)
@@ -457,6 +458,23 @@ def default_2_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argpa
return parser
def hq_2_stage_arg_parser(params: PipelineParams = LTX_2_3_HQ_PARAMS) -> argparse.ArgumentParser:
parser = default_2_stage_arg_parser(params=params)
parser.add_argument(
"--distilled-lora-strength-stage-1",
type=float,
default=0.25,
help=(f"Strength of the distilled LoRA used in the first stage (default: {0.25})."),
)
parser.add_argument(
"--distilled-lora-strength-stage-2",
type=float,
default=0.5,
help=(f"Strength of the distilled LoRA used in the second stage (default: {0.5})."),
)
return parser
def default_2_stage_distilled_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser:
parser = basic_arg_parser(params=params, distilled=True)
parser.set_defaults(height=params.stage_2_height, width=params.stage_2_width)
@@ -71,6 +71,27 @@ LTX_2_3_PARAMS = replace(
video_guider_params=replace(LTX_2_PARAMS.video_guider_params, stg_blocks=[28]),
audio_guider_params=replace(LTX_2_PARAMS.audio_guider_params, stg_blocks=[28]),
)
LTX_2_3_HQ_PARAMS = PipelineParams(
num_inference_steps=15,
stage_1_height=1088 // 2,
stage_1_width=1920 // 2,
video_guider_params=MultiModalGuiderParams(
cfg_scale=3.0,
stg_scale=0.0,
rescale_scale=0.45,
modality_scale=3.0,
skip_step=0,
stg_blocks=[],
),
audio_guider_params=MultiModalGuiderParams(
cfg_scale=7.0,
stg_scale=0.0,
rescale_scale=1.0,
modality_scale=3.0,
skip_step=0,
stg_blocks=[],
),
)
DEFAULT_LORA_STRENGTH = 1.0
DEFAULT_IMAGE_CRF = 33
@@ -21,6 +21,7 @@ from ltx_core.guidance.perturbations import (
from ltx_core.model.transformer import Modality, X0Model
from ltx_core.model.video_vae import VideoEncoder
from ltx_core.text_encoders.gemma import GemmaTextEncoder
from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessorOutput
from ltx_core.tools import AudioLatentTools, LatentTools, VideoLatentTools
from ltx_core.types import AudioLatentShape, LatentState, VideoLatentShape, VideoPixelShape
from ltx_pipelines.utils.args import ImageConditioningInput
@@ -44,6 +45,84 @@ def cleanup_memory() -> None:
torch.cuda.synchronize()
def encode_prompts(
prompts: list[str],
model_ledger: object,
*,
enhance_prompt_image: str | None = None,
enhance_prompt_seed: int = 42,
enhance_first_prompt: bool = False,
) -> list[EmbeddingsProcessorOutput]:
"""Encode prompts through Gemma → embeddings processor, freeing each after use.
Loads the text encoder from *model_ledger*, optionally enhances the first
prompt, encodes all *prompts*, frees the text encoder, then loads the
embeddings processor to produce the final outputs. Because the text encoder
is loaded and freed entirely within this function, there are no lingering
references that could prevent GPU memory reclamation.
Args:
prompts: Text prompts to encode.
model_ledger: ModelLedger instance (used to load text encoder and embeddings processor).
enhance_prompt_image: Optional image path for prompt enhancement.
enhance_prompt_seed: Seed for prompt enhancement (default 42).
enhance_first_prompt: If True, enhance ``prompts[0]`` before encoding.
Returns:
List of EmbeddingsProcessorOutput, one per prompt.
"""
text_encoder = model_ledger.text_encoder()
if enhance_first_prompt:
prompts = list(prompts)
prompts[0] = generate_enhanced_prompt(text_encoder, prompts[0], enhance_prompt_image, seed=enhance_prompt_seed)
raw_outputs = [text_encoder.encode(p) for p in prompts]
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
embeddings_processor = model_ledger.gemma_embeddings_processor()
results: list[EmbeddingsProcessorOutput] = [
embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw_outputs
]
del embeddings_processor
cleanup_memory()
return results
def combined_image_conditionings(
images: list[ImageConditioningInput],
height: int,
width: int,
video_encoder: VideoEncoder,
dtype: torch.dtype,
device: torch.device,
) -> list[ConditioningItem]:
"""Create a list of conditionings by replacing the latent at the first frame with the encoded image if present
and using other encoded images as the keyframe conditionings."""
conditionings = []
for img in images:
image = load_image_conditioning(
image_path=img.path,
height=height,
width=width,
dtype=dtype,
device=device,
crf=img.crf,
)
encoded_image = video_encoder(image)
if img.frame_idx == 0:
conditioning = VideoConditionByLatentIndex(
latent=encoded_image,
strength=img.strength,
latent_idx=0,
)
else:
conditioning = VideoConditionByKeyframeIndex(
keyframes=encoded_image,
strength=img.strength,
frame_idx=img.frame_idx,
)
conditionings.append(conditioning)
return conditionings
def image_conditionings_by_replacing_latent(
images: list[ImageConditioningInput],
height: int,
@@ -33,8 +33,11 @@ from ltx_core.model.video_vae import (
)
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import (
AV_GEMMA_TEXT_ENCODER_KEY_OPS,
EMBEDDINGS_PROCESSOR_KEY_OPS,
GEMMA_LLM_KEY_OPS,
GEMMA_MODEL_OPS,
EmbeddingsProcessor,
EmbeddingsProcessorConfigurator,
GemmaTextEncoder,
GemmaTextEncoderConfigurator,
module_ops_from_gemma_root,
@@ -76,8 +79,8 @@ class ModelLedger:
:meth:`spatial_upsampler` method becomes available; otherwise calling it raises
a :class:`ValueError`.
loras:
Optional collection of LoRA configurations (paths, strengths, and key operations)
that are applied on top of the base transformer weights when building the model.
Tuple of LoRA configurations (path, strength, sd_ops) applied on top of the base
transformer weights. Use ``()`` for none.
registry:
Optional :class:`Registry` instance for weight caching across builders.
Defaults to :class:`DummyRegistry` which performs no cross-builder caching.
@@ -85,8 +88,9 @@ class ModelLedger:
Optional :class:`QuantizationPolicy` controlling how transformer weights
are stored and how matmul is executed. Defaults to None, which means no quantization.
### Creating Variants
Use :meth:`with_loras` to create a new ``ModelLedger`` instance that includes
additional LoRA configurations while sharing the same registry for weight caching.
Use :meth:`with_additional_loras` to create a new ``ModelLedger`` instance that
includes additional LoRA configurations or :meth:`with_loras` to replace existing
lora configurations while sharing the same registry for weight caching.
"""
def __init__(
@@ -96,7 +100,7 @@ class ModelLedger:
checkpoint_path: str | None = None,
gemma_root_path: str | None = None,
spatial_upsampler_path: str | None = None,
loras: LoraPathStrengthAndSDOps | None = None,
loras: tuple[LoraPathStrengthAndSDOps, ...] = (),
registry: Registry | None = None,
quantization: QuantizationPolicy | None = None,
):
@@ -105,7 +109,7 @@ class ModelLedger:
self.checkpoint_path = checkpoint_path
self.gemma_root_path = gemma_root_path
self.spatial_upsampler_path = spatial_upsampler_path
self.loras = loras or ()
self.loras = loras
self.registry = registry or DummyRegistry()
self.quantization = quantization
self.build_model_builders()
@@ -155,15 +159,23 @@ class ModelLedger:
registry=self.registry,
)
# Embeddings processor only needs the LTX checkpoint (no Gemma weights)
self.embeddings_processor_builder = Builder(
model_path=self.checkpoint_path,
model_class_configurator=EmbeddingsProcessorConfigurator,
model_sd_ops=EMBEDDINGS_PROCESSOR_KEY_OPS,
registry=self.registry,
)
if self.gemma_root_path is not None:
module_ops = module_ops_from_gemma_root(self.gemma_root_path)
model_folder = find_matching_file(self.gemma_root_path, "model*.safetensors").parent
weight_paths = [str(p) for p in model_folder.rglob("*.safetensors")]
self.text_encoder_builder = Builder(
model_path=(str(self.checkpoint_path), *weight_paths),
model_path=tuple(weight_paths),
model_class_configurator=GemmaTextEncoderConfigurator,
model_sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS,
model_sd_ops=GEMMA_LLM_KEY_OPS,
registry=self.registry,
module_ops=(GEMMA_MODEL_OPS, *module_ops),
)
@@ -181,14 +193,19 @@ class ModelLedger:
else:
return torch.device("cpu")
def with_loras(self, loras: LoraPathStrengthAndSDOps) -> "ModelLedger":
def with_additional_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> "ModelLedger":
"""Add new lora configurations to the existing ones."""
return self.with_loras((*self.loras, *loras))
def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> "ModelLedger":
"""Replace existing lora configurations with new ones."""
return ModelLedger(
dtype=self.dtype,
device=self.device,
checkpoint_path=self.checkpoint_path,
gemma_root_path=self.gemma_root_path,
spatial_upsampler_path=self.spatial_upsampler_path,
loras=(*self.loras, *loras),
loras=loras,
registry=self.registry,
quantization=self.quantization,
)
@@ -244,6 +261,18 @@ class ModelLedger:
return self.text_encoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
def gemma_embeddings_processor(self) -> EmbeddingsProcessor:
if not hasattr(self, "embeddings_processor_builder"):
raise ValueError(
"Embeddings processor not initialized. Please provide a checkpoint path to the ModelLedger constructor."
)
return (
self.embeddings_processor_builder.build(device=self._target_device(), dtype=self.dtype)
.to(self.device)
.eval()
)
def audio_encoder(self) -> AudioEncoder:
if not hasattr(self, "audio_encoder_builder"):
raise ValueError(