Automated PR - 2026-03-05
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user