indigo/LTX-2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Automated PR - 2026-03-05

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-pipelines/src/ltx_pipelines/a2vid_two_stage.py
+44 -44
View File
@@ -14,7 +14,6 @@ from ltx_core.model.upsampler import upsample_video
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import Audio, AudioLatentShape, LatentState, VideoPixelShape
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.args import default_2_stage_arg_parser
@@ -24,10 +23,10 @@ from ltx_pipelines.utils.constants import (
from ltx_pipelines.utils.helpers import (
assert_resolution,
cleanup_memory,
combined_image_conditionings,
denoise_video_only,
generate_enhanced_prompt,
encode_prompts,
get_device,
image_conditionings_by_replacing_latent,
multi_modal_guider_denoising_func,
simple_denoising_func,
)
@@ -69,7 +68,7 @@ class A2VidPipelineTwoStage:
quantization=quantization,
)
self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
self.stage_2_model_ledger = self.stage_1_model_ledger.with_additional_loras(
loras=distilled_lora,
)
@@ -103,16 +102,14 @@ class A2VidPipelineTwoStage:
stepper = EulerDiffusionStep()
dtype = torch.bfloat16
text_encoder = self.stage_1_model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(text_encoder, prompt, images[0][0] if len(images) > 0 else None)
context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
v_context_p, a_context_p = context_p
v_context_n, _ = context_n
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
ctx_p, ctx_n = encode_prompts(
[prompt, negative_prompt],
self.stage_1_model_ledger,
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
)
v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
v_context_n, _ = ctx_n.video_encoding, ctx_n.audio_encoding
# Encode audio.
decoded_audio = decode_audio_from_file(audio_path, self.device, audio_start_time, audio_max_duration)
@@ -120,8 +117,28 @@ class A2VidPipelineTwoStage:
audio_shape = AudioLatentShape.from_duration(batch=1, duration=num_frames / frame_rate, channels=8, mel_bins=16)
encoded_audio_latent = encoded_audio_latent[:, :, : audio_shape.frames]
# Stage 1: encode image conditionings with the VAE encoder, then free it
# before loading the transformer to reduce peak VRAM.
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
width=width // 2,
height=height // 2,
fps=frame_rate,
)
video_encoder = self.stage_1_model_ledger.video_encoder()
stage_1_conditionings = combined_image_conditionings(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
torch.cuda.synchronize()
del video_encoder
cleanup_memory()
# Stage 1: Initial low resolution video generation with audio conditioning.
transformer = self.stage_1_model_ledger.transformer()
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
@@ -147,24 +164,6 @@ class A2VidPipelineTwoStage:
),
)
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
width=width // 2,
height=height // 2,
fps=frame_rate,
)
video_encoder = self.stage_1_model_ledger.video_encoder()
stage_1_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state = denoise_video_only(
output_shape=stage_1_output_shape,
conditionings=stage_1_conditionings,
@@ -183,12 +182,23 @@ class A2VidPipelineTwoStage:
cleanup_memory()
# Stage 2: Upsample and refine the video at higher resolution with distilled LoRA.
video_encoder = self.stage_1_model_ledger.video_encoder()
upscaled_video_latent = upsample_video(
latent=video_state.latent[:1],
video_encoder=video_encoder,
upsampler=self.stage_2_model_ledger.spatial_upsampler(),
)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = combined_image_conditionings(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
del video_encoder
torch.cuda.synchronize()
cleanup_memory()
@@ -210,15 +220,6 @@ class A2VidPipelineTwoStage:
),
)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state = denoise_video_only(
output_shape=stage_2_output_shape,
conditionings=stage_2_conditionings,
@@ -236,7 +237,6 @@ class A2VidPipelineTwoStage:
torch.cuda.synchronize()
del transformer
del video_encoder
cleanup_memory()
decoded_video = vae_decode_video(
@@ -278,7 +278,7 @@ def main() -> None:
distilled_lora=args.distilled_lora,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/distilled.py
+12 -15
View File
@@ -12,7 +12,6 @@ from ltx_core.model.upsampler import upsample_video
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import Audio, LatentState, VideoPixelShape
from ltx_pipelines.utils import ModelLedger, euler_denoising_loop
from ltx_pipelines.utils.args import (
@@ -28,10 +27,10 @@ from ltx_pipelines.utils.constants import (
from ltx_pipelines.utils.helpers import (
assert_resolution,
cleanup_memory,
combined_image_conditionings,
denoise_audio_video,
generate_enhanced_prompt,
encode_prompts,
get_device,
image_conditionings_by_replacing_latent,
simple_denoising_func,
)
from ltx_pipelines.utils.media_io import encode_video
@@ -93,15 +92,13 @@ class DistilledPipeline:
stepper = EulerDiffusionStep()
dtype = torch.bfloat16
text_encoder = self.model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(text_encoder, prompt, images[0][0] if len(images) > 0 else None)
context_p = encode_text(text_encoder, prompts=[prompt])[0]
video_context, audio_context = context_p
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
(ctx_p,) = encode_prompts(
[prompt],
self.model_ledger,
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
)
video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding
# Stage 1: Initial low resolution video generation.
video_encoder = self.model_ledger.video_encoder()
@@ -130,7 +127,7 @@ class DistilledPipeline:
height=height // 2,
fps=frame_rate,
)
stage_1_conditionings = image_conditionings_by_replacing_latent(
stage_1_conditionings = combined_image_conditionings(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
@@ -161,7 +158,7 @@ class DistilledPipeline:
stage_2_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = image_conditionings_by_replacing_latent(
stage_2_conditionings = combined_image_conditionings(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
@@ -209,7 +206,7 @@ def main() -> None:
distilled_checkpoint_path=args.distilled_checkpoint_path,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py
+33 -36
View File
@@ -19,17 +19,16 @@ from ltx_core.model.upsampler import upsample_video
from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import Audio, LatentState, VideoLatentShape, VideoPixelShape
from ltx_pipelines.utils import (
ModelLedger,
assert_resolution,
cleanup_memory,
combined_image_conditionings,
denoise_audio_video,
encode_prompts,
euler_denoising_loop,
generate_enhanced_prompt,
get_device,
image_conditionings_by_replacing_latent,
simple_denoising_func,
)
from ltx_pipelines.utils.args import (
@@ -169,20 +168,37 @@ class ICLoraPipeline:
stepper = EulerDiffusionStep()
dtype = torch.bfloat16
text_encoder = self.stage_1_model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(
text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
)
video_context, audio_context = encode_text(text_encoder, prompts=[prompt])[0]
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
(ctx_p,) = encode_prompts(
[prompt],
self.stage_1_model_ledger,
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
enhance_prompt_seed=seed,
)
video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding
# Stage 1: Initial low resolution video generation.
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
width=width // 2,
height=height // 2,
fps=frame_rate,
)
# Encode conditionings before loading transformer to reduce peak VRAM
video_encoder = self.stage_1_model_ledger.video_encoder()
stage_1_conditionings = self._create_conditionings(
images=images,
video_conditioning=video_conditioning,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
num_frames=num_frames,
conditioning_attention_strength=conditioning_attention_strength,
conditioning_attention_mask=conditioning_attention_mask,
)
transformer = self.stage_1_model_ledger.transformer()
stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES).to(self.device)
@@ -201,25 +217,6 @@ class ICLoraPipeline:
),
)
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
width=width // 2,
height=height // 2,
fps=frame_rate,
)
stage_1_conditionings = self._create_conditionings(
images=images,
video_conditioning=video_conditioning,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
num_frames=num_frames,
conditioning_attention_strength=conditioning_attention_strength,
conditioning_attention_mask=conditioning_attention_mask,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_1_output_shape,
conditionings=stage_1_conditionings,
@@ -278,7 +275,7 @@ class ICLoraPipeline:
)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = image_conditionings_by_replacing_latent(
stage_2_conditionings = combined_image_conditionings(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
@@ -340,7 +337,7 @@ class ICLoraPipeline:
Returns:
List of conditioning items. IC-LoRA conditionings are appended last.
"""
conditionings = image_conditionings_by_replacing_latent(
conditionings = combined_image_conditionings(
images=images,
height=height,
width=width,
@@ -510,7 +507,7 @@ def main() -> None:
distilled_checkpoint_path=args.distilled_checkpoint_path,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py
+12 -16
View File
@@ -18,7 +18,6 @@ from ltx_core.model.upsampler import upsample_video
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import Audio, LatentState, VideoPixelShape
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.args import ImageConditioningInput, default_2_stage_arg_parser, detect_checkpoint_path
@@ -27,7 +26,7 @@ from ltx_pipelines.utils.helpers import (
assert_resolution,
cleanup_memory,
denoise_audio_video,
generate_enhanced_prompt,
encode_prompts,
get_device,
image_conditionings_by_adding_guiding_latent,
multi_modal_guider_factory_denoising_func,
@@ -71,7 +70,7 @@ class KeyframeInterpolationPipeline:
loras=loras,
quantization=quantization,
)
self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
self.stage_2_model_ledger = self.stage_1_model_ledger.with_additional_loras(
loras=distilled_lora,
)
self.pipeline_components = PipelineComponents(
@@ -102,18 +101,15 @@ class KeyframeInterpolationPipeline:
stepper = EulerDiffusionStep()
dtype = torch.bfloat16
text_encoder = self.stage_1_model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(
text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
)
context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
v_context_p, a_context_p = context_p
v_context_n, a_context_n = context_n
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
ctx_p, ctx_n = encode_prompts(
[prompt, negative_prompt],
self.stage_1_model_ledger,
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
enhance_prompt_seed=seed,
)
v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
# Stage 1: Initial low resolution video generation.
video_encoder = self.stage_1_model_ledger.video_encoder()
@@ -252,7 +248,7 @@ def main() -> None:
distilled_lora=args.distilled_lora,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/retake.py
+34 -22
View File
@@ -20,7 +20,6 @@ from ltx_core.model.audio_vae import encode_audio as vae_encode_audio
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.tools import LatentTools
from ltx_core.types import (
Audio,
@@ -30,10 +29,11 @@ from ltx_core.types import (
VideoPixelShape,
)
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES
from ltx_pipelines.utils.args import QuantizationAction
from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES, detect_params
from ltx_pipelines.utils.helpers import (
cleanup_memory,
generate_enhanced_prompt,
encode_prompts,
get_device,
multi_modal_guider_denoising_func,
noise_audio_state,
@@ -197,7 +197,6 @@ class RetakePipeline:
# Public entry point #
# --------------------------------------------------------------------- #
@torch.inference_mode()
def __call__( # noqa: PLR0913, PLR0915
self,
video_path: str,
@@ -214,6 +213,7 @@ class RetakePipeline:
regenerate_audio: bool = True,
enhance_prompt: bool = False,
distilled: bool = False,
tiling_config: TilingConfig | None = None,
) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
"""Regenerate ``[start_time, end_time]`` of the source video (retake).
Parameters
@@ -321,22 +321,17 @@ class RetakePipeline:
del audio_encoder
cleanup_memory()
text_encoder = self.model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(text_encoder, prompt, None, seed=effective_seed)
prompts_to_encode = [prompt] if distilled else [prompt, negative_prompt]
contexts = encode_prompts(
prompts_to_encode,
self.model_ledger,
enhance_first_prompt=enhance_prompt,
enhance_prompt_seed=effective_seed,
)
if distilled:
# Distilled mode: single prompt, no negative
context_p = encode_text(text_encoder, prompts=[prompt])[0]
v_context_p, a_context_p = context_p
else:
context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
v_context_p, a_context_p = context_p
v_context_n, a_context_n = context_n
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
v_context_p, a_context_p = contexts[0].video_encoding, contexts[0].audio_encoding
if not distilled:
v_context_n, a_context_n = contexts[1].video_encoding, contexts[1].audio_encoding
transformer = self.model_ledger.transformer()
@@ -412,7 +407,9 @@ class RetakePipeline:
del transformer
cleanup_memory()
decoded_video = vae_decode_video(video_state.latent, self.model_ledger.video_decoder(), generator=generator)
decoded_video = vae_decode_video(
video_state.latent, self.model_ledger.video_decoder(), tiling_config, generator
)
decoded_audio = vae_decode_audio(
audio_state.latent, self.model_ledger.audio_decoder(), self.model_ledger.vocoder()
)
@@ -420,6 +417,7 @@ class RetakePipeline:
return decoded_video, decoded_audio
@torch.inference_mode()
def main() -> None:
"""CLI entry point for retake (regenerate a time region)."""
logging.getLogger().setLevel(logging.INFO)
@@ -433,6 +431,15 @@ def main() -> None:
parser.add_argument("--gemma-root", type=str, required=True, help="Path to Gemma text encoder weights.")
parser.add_argument("--seed", type=int, default=42, help="Random seed. Use -1 for a random seed.")
parser.add_argument("--loras", nargs="*", default=[], help="LoRA paths (optional).")
parser.add_argument(
"--quantization",
dest="quantization",
action=QuantizationAction,
nargs="+",
metavar=("POLICY", "AMAX_PATH"),
default=None,
help="Quantization policy: fp8-cast or fp8-scaled-mm [AMAX_PATH].",
)
args = parser.parse_args()
if args.start_time >= args.end_time:
@@ -452,16 +459,21 @@ def main() -> None:
pipeline = RetakePipeline(
checkpoint_path=args.checkpoint_path,
gemma_root=args.gemma_root,
loras=args.loras or [],
loras=tuple(args.loras) if args.loras else (),
quantization=args.quantization,
)
params = detect_params(args.checkpoint_path)
tiling_config = TilingConfig.default()
video_iter, audio = pipeline(
video_path=args.video_path,
prompt=args.prompt,
start_time=args.start_time,
end_time=args.end_time,
seed=args.seed,
video_guider_params=params.video_guider_params,
audio_guider_params=params.audio_guider_params,
tiling_config=tiling_config,
)
tiling_config = TilingConfig.default()
video_chunks_number = get_video_chunks_number(num_frames, tiling_config)
encode_video(
video=video_iter,
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_one_stage.py
+25 -25
View File
@@ -16,17 +16,16 @@ from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import Audio, LatentState, VideoPixelShape
from ltx_pipelines.utils import (
ModelLedger,
assert_resolution,
cleanup_memory,
combined_image_conditionings,
denoise_audio_video,
encode_prompts,
euler_denoising_loop,
generate_enhanced_prompt,
get_device,
image_conditionings_by_replacing_latent,
multi_modal_guider_factory_denoising_func,
)
from ltx_pipelines.utils.args import ImageConditioningInput, default_1_stage_arg_parser, detect_checkpoint_path
@@ -91,21 +90,32 @@ class TI2VidOneStagePipeline:
stepper = EulerDiffusionStep()
dtype = torch.bfloat16
text_encoder = self.model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(
text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
)
context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
v_context_p, a_context_p = context_p
v_context_n, a_context_n = context_n
ctx_p, ctx_n = encode_prompts(
[prompt, negative_prompt],
self.model_ledger,
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
enhance_prompt_seed=seed,
)
v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
# Encode image conditionings with the VAE encoder, then free it
# before loading the transformer to reduce peak VRAM.
stage_1_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
video_encoder = self.model_ledger.video_encoder()
stage_1_conditionings = combined_image_conditionings(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
torch.cuda.synchronize()
del text_encoder
del video_encoder
cleanup_memory()
# Stage 1: Initial low resolution video generation.
video_encoder = self.model_ledger.video_encoder()
transformer = self.model_ledger.transformer()
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
@@ -135,16 +145,6 @@ class TI2VidOneStagePipeline:
),
)
stage_1_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_1_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_1_output_shape,
conditionings=stage_1_conditionings,
@@ -178,7 +178,7 @@ def main() -> None:
pipeline = TI2VidOneStagePipeline(
checkpoint_path=args.checkpoint_path,
gemma_root=args.gemma_root,
loras=args.lora,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
)
video, audio = pipeline(
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py
+43 -41
View File
@@ -18,17 +18,16 @@ from ltx_core.model.upsampler import upsample_video
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import Audio, LatentState, VideoPixelShape
from ltx_pipelines.utils import (
ModelLedger,
assert_resolution,
cleanup_memory,
combined_image_conditionings,
denoise_audio_video,
encode_prompts,
euler_denoising_loop,
generate_enhanced_prompt,
get_device,
image_conditionings_by_replacing_latent,
multi_modal_guider_factory_denoising_func,
simple_denoising_func,
)
@@ -71,7 +70,7 @@ class TI2VidTwoStagesPipeline:
quantization=quantization,
)
self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
self.stage_2_model_ledger = self.stage_1_model_ledger.with_additional_loras(
loras=distilled_lora,
)
@@ -103,21 +102,38 @@ class TI2VidTwoStagesPipeline:
stepper = EulerDiffusionStep()
dtype = torch.bfloat16
text_encoder = self.stage_1_model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(
text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
)
context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
v_context_p, a_context_p = context_p
v_context_n, a_context_n = context_n
ctx_p, ctx_n = encode_prompts(
[prompt, negative_prompt],
self.stage_1_model_ledger,
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
enhance_prompt_seed=seed,
)
v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
# Stage 1: encode image conditionings with the VAE encoder, then free it
# before loading the transformer to reduce peak VRAM.
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
width=width // 2,
height=height // 2,
fps=frame_rate,
)
video_encoder = self.stage_1_model_ledger.video_encoder()
stage_1_conditionings = combined_image_conditionings(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
torch.cuda.synchronize()
del text_encoder
del video_encoder
cleanup_memory()
# Stage 1: Initial low resolution video generation.
video_encoder = self.stage_1_model_ledger.video_encoder()
transformer = self.stage_1_model_ledger.transformer()
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
@@ -144,21 +160,6 @@ class TI2VidTwoStagesPipeline:
),
)
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
width=width // 2,
height=height // 2,
fps=frame_rate,
)
stage_1_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_1_output_shape,
conditionings=stage_1_conditionings,
@@ -176,12 +177,23 @@ class TI2VidTwoStagesPipeline:
cleanup_memory()
# Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
video_encoder = self.stage_1_model_ledger.video_encoder()
upscaled_video_latent = upsample_video(
latent=video_state.latent[:1],
video_encoder=video_encoder,
upsampler=self.stage_2_model_ledger.spatial_upsampler(),
)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = combined_image_conditionings(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
del video_encoder
torch.cuda.synchronize()
cleanup_memory()
@@ -203,15 +215,6 @@ class TI2VidTwoStagesPipeline:
),
)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_2_output_shape,
conditionings=stage_2_conditionings,
@@ -229,7 +232,6 @@ class TI2VidTwoStagesPipeline:
torch.cuda.synchronize()
del transformer
del video_encoder
cleanup_memory()
decoded_video = vae_decode_video(
@@ -253,7 +255,7 @@ def main() -> None:
distilled_lora=args.distilled_lora,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_res2s.py → packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_hq.py
+62 -48
View File
@@ -14,30 +14,29 @@ from ltx_core.model.upsampler import upsample_video
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.tools import VideoLatentShape
from ltx_core.types import Audio, LatentState, VideoPixelShape
from ltx_pipelines.utils import (
ModelLedger,
assert_resolution,
cleanup_memory,
combined_image_conditionings,
denoise_audio_video,
generate_enhanced_prompt,
encode_prompts,
get_device,
image_conditionings_by_replacing_latent,
multi_modal_guider_denoising_func,
res2s_audio_video_denoising_loop,
simple_denoising_func,
)
from ltx_pipelines.utils.args import ImageConditioningInput, default_2_stage_arg_parser, detect_checkpoint_path
from ltx_pipelines.utils.constants import STAGE_2_DISTILLED_SIGMA_VALUES, detect_params
from ltx_pipelines.utils.args import ImageConditioningInput, hq_2_stage_arg_parser
from ltx_pipelines.utils.constants import LTX_2_3_HQ_PARAMS, STAGE_2_DISTILLED_SIGMA_VALUES
from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.types import PipelineComponents
device = get_device()
class TI2VidTwoStagesRes2sPipeline:
class TI2VidTwoStagesHQPipeline:
"""
Two-stage text/image-to-video generation pipeline using the res_2s sampler.
Same structure as :class:`TI2VidTwoStagesPipeline`: stage 1 generates video at
@@ -53,26 +52,38 @@ class TI2VidTwoStagesRes2sPipeline:
self,
checkpoint_path: str,
distilled_lora: list[LoraPathStrengthAndSDOps],
distilled_lora_strength_stage_1: float,
distilled_lora_strength_stage_2: float,
spatial_upsampler_path: str,
gemma_root: str,
loras: list[LoraPathStrengthAndSDOps],
loras: tuple[LoraPathStrengthAndSDOps, ...],
device: str = device,
quantization: QuantizationPolicy | None = None,
):
self.device = device
self.dtype = torch.bfloat16
distilled_lora_stage_1 = LoraPathStrengthAndSDOps(
path=distilled_lora[0].path,
strength=distilled_lora_strength_stage_1,
sd_ops=distilled_lora[0].sd_ops,
)
distilled_lora_stage_2 = LoraPathStrengthAndSDOps(
path=distilled_lora[0].path,
strength=distilled_lora_strength_stage_2,
sd_ops=distilled_lora[0].sd_ops,
)
self.stage_1_model_ledger = ModelLedger(
dtype=self.dtype,
device=device,
checkpoint_path=checkpoint_path,
gemma_root_path=gemma_root,
spatial_upsampler_path=spatial_upsampler_path,
loras=loras,
loras=(*loras, distilled_lora_stage_1),
quantization=quantization,
)
self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
loras=distilled_lora,
loras=(*loras, distilled_lora_stage_2),
)
self.pipeline_components = PipelineComponents(
@@ -103,23 +114,18 @@ class TI2VidTwoStagesRes2sPipeline:
noiser = GaussianNoiser(generator=generator)
dtype = torch.bfloat16
text_encoder = self.stage_1_model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(
text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
)
context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
v_context_p, a_context_p = context_p
v_context_n, a_context_n = context_n
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
# Stage 1: Initial low resolution video generation.
video_encoder = self.stage_1_model_ledger.video_encoder()
transformer = self.stage_1_model_ledger.transformer()
ctx_p, ctx_n = encode_prompts(
[prompt, negative_prompt],
self.stage_1_model_ledger,
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
enhance_prompt_seed=seed,
)
v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
# Stage 1: encode image conditionings with the VAE encoder, then free it
# before loading the transformer to reduce peak VRAM.
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
@@ -127,6 +133,21 @@ class TI2VidTwoStagesRes2sPipeline:
height=height // 2,
fps=frame_rate,
)
video_encoder = self.stage_1_model_ledger.video_encoder()
stage_1_conditionings = combined_image_conditionings(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
torch.cuda.synchronize()
del video_encoder
cleanup_memory()
transformer = self.stage_1_model_ledger.transformer()
empty_latent = torch.empty(VideoLatentShape.from_pixel_shape(stage_1_output_shape).to_torch_shape())
stepper = Res2sDiffusionStep()
sigmas = (
@@ -158,14 +179,6 @@ class TI2VidTwoStagesRes2sPipeline:
),
)
stage_1_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_1_output_shape,
conditionings=stage_1_conditionings,
@@ -183,13 +196,24 @@ class TI2VidTwoStagesRes2sPipeline:
cleanup_memory()
# Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
video_encoder = self.stage_1_model_ledger.video_encoder()
upscaled_video_latent = upsample_video(
latent=video_state.latent[:1],
video_encoder=video_encoder,
upsampler=self.stage_2_model_ledger.spatial_upsampler(),
)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = combined_image_conditionings(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
torch.cuda.synchronize()
del video_encoder
cleanup_memory()
transformer = self.stage_2_model_ledger.transformer()
@@ -210,15 +234,6 @@ class TI2VidTwoStagesRes2sPipeline:
),
)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_2_output_shape,
conditionings=stage_2_conditionings,
@@ -236,7 +251,6 @@ class TI2VidTwoStagesRes2sPipeline:
torch.cuda.synchronize()
del transformer
del video_encoder
cleanup_memory()
decoded_video = vae_decode_video(
@@ -251,16 +265,16 @@ class TI2VidTwoStagesRes2sPipeline:
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
checkpoint_path = detect_checkpoint_path()
params = detect_params(checkpoint_path)
parser = default_2_stage_arg_parser(params=params)
parser = hq_2_stage_arg_parser(params=LTX_2_3_HQ_PARAMS)
args = parser.parse_args()
pipeline = TI2VidTwoStagesRes2sPipeline(
pipeline = TI2VidTwoStagesHQPipeline(
checkpoint_path=args.checkpoint_path,
distilled_lora=args.distilled_lora,
distilled_lora_strength_stage_1=args.distilled_lora_strength_stage_1,
distilled_lora_strength_stage_2=args.distilled_lora_strength_stage_2,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/utils/__init__.py
+4 -2
View File
@@ -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",
packages/ltx-pipelines/src/ltx_pipelines/utils/args.py
+18
View File
@@ -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)
packages/ltx-pipelines/src/ltx_pipelines/utils/constants.py
+21
View File
@@ -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
packages/ltx-pipelines/src/ltx_pipelines/utils/helpers.py
+79
View File
@@ -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,
packages/ltx-pipelines/src/ltx_pipelines/utils/model_ledger.py
+40 -11
View File
@@ -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(