diff --git a/modules/ltx/ltx_capabilities.py b/modules/ltx/ltx_capabilities.py index d345eaa2a..499f8b57b 100644 --- a/modules/ltx/ltx_capabilities.py +++ b/modules/ltx/ltx_capabilities.py @@ -9,8 +9,8 @@ class LTXCaps: name: str repo_cls_name: str family: str # '0.9' or '2.x' + variant: str # '0.9', '2.0', '2.3' (finer-grained sub-variant) is_distilled: bool - is_ltx_2_3: bool is_i2v: bool supports_input_media: bool supports_multi_condition: bool @@ -76,18 +76,22 @@ def get_caps(model_name: str) -> Optional[LTXCaps]: is_ltx2 = cls_name in LTX2_CLASSES family = '2.x' if is_ltx2 else '0.9' + # 2.x sub-variant detection: unknown 2.x mirrors fall through to '2.0' (conservative default). + if is_ltx2: + variant = '2.3' if '2.3' in model_name else '2.0' + else: + variant = '0.9' is_distilled = 'Distilled' in model_name is_i2v = 'I2V' in model_name or cls_name in ('LTXImageToVideoPipeline', 'LTX2ImageToVideoPipeline') is_condition_cls = cls_name in CONDITION_CLASSES supports_input_media = is_i2v or is_condition_cls - is_ltx_2_3 = is_ltx2 and '2.3' in model_name caps = LTXCaps( name=model_name, repo_cls_name=cls_name, family=family, + variant=variant, is_distilled=is_distilled, - is_ltx_2_3=is_ltx_2_3, is_i2v=is_i2v, supports_input_media=supports_input_media, supports_multi_condition=is_condition_cls, @@ -96,7 +100,7 @@ def get_caps(model_name: str) -> Optional[LTXCaps]: supports_stg=is_ltx2, supports_audio=is_ltx2, supports_frame_rate_kwarg=is_ltx2, - use_cross_timestep=is_ltx_2_3, + use_cross_timestep=(variant == '2.3'), default_cfg=3.0, default_steps=30 if is_ltx2 else 50, default_sampler_shift=-1.0, @@ -112,16 +116,16 @@ def get_caps(model_name: str) -> Optional[LTXCaps]: caps.default_steps = 8 if is_ltx2 and not is_distilled: - if is_ltx_2_3: + if variant == '2.3': caps.stage2_dev_lora_repo = 'CalamitousFelicitousness/LTX-2.3-distilled-lora-384-Diffusers' - elif '2.0' in model_name: + elif variant == '2.0': caps.stage2_dev_lora_repo = 'CalamitousFelicitousness/LTX-2.0-distilled-lora-384-Diffusers' caps.supports_canonical_stage2 = caps.stage2_dev_lora_repo is not None if is_ltx2: - if '2.3' in model_name: + if variant == '2.3': caps.stg_default_blocks = [28] - elif '2.0' in model_name: + elif variant == '2.0': caps.stg_default_blocks = [29] else: caps.stg_default_blocks = [28] diff --git a/modules/ltx/ltx_process.py b/modules/ltx/ltx_process.py index 9c64bafc8..18aab0d44 100644 --- a/modules/ltx/ltx_process.py +++ b/modules/ltx/ltx_process.py @@ -6,7 +6,7 @@ from PIL import Image from modules import shared, errors, timer, memstats, progress, processing, sd_models, sd_samplers, devices, extra_networks, call_queue from modules.logger import log from modules.ltx import ltx_capabilities -from modules.ltx.ltx_util import get_bucket, get_frames, load_model, load_upsample, get_conditions, get_generator, get_prompts, vae_decode +from modules.ltx.ltx_util import get_bucket, get_frames, load_model, load_upsample, load_upsample_2x, get_conditions, get_generator, get_prompts, vae_decode from modules.processing_callbacks import diffusers_callback from modules.video_models.video_vae import set_vae_params from modules.video_models.video_save import save_video @@ -20,6 +20,7 @@ upsample_repo_id_09 = 'a-r-r-o-w/LTX-Video-0.9.7-Latent-Spatial-Upsampler-diffus upsample_repo_id_20 = 'Lightricks/LTX-2' upsample_repo_id_23 = 'CalamitousFelicitousness/LTX-2.3-Spatial-Upsampler-x2-1.1-Diffusers' upsample_pipe = None +upsample_pipe_2x = None STAGE2_DEV_LORA_ADAPTER = 'ltx2_stage2_distilled' @@ -394,20 +395,10 @@ def run_ltx(task_id, latents = upsample_pipe(latents=latents, **up_args).frames[0] upsample_pipe = sd_models.apply_balanced_offload(upsample_pipe) else: - from diffusers.pipelines.ltx2.pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline - from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel - # Skip apply_balanced_offload on the upsampler; checkpoint_name differs from the main - # pipe so the shared OffloadHook (sd_offload.py:488) would rebuild and force a heavy - # re-init on the next refine. At ~2.3GB it fits on device; free after the pass. - upsample_repo = upsample_repo_id_23 if '2.3' in caps.name else upsample_repo_id_20 - log.info(f'Video load: cls={LTX2LatentUpsamplePipeline.__name__} family=2.x repo={upsample_repo} auto={auto_refine_upsample}') - latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained( - upsample_repo, - subfolder='latent_upsampler', - cache_dir=shared.opts.hfcache_dir, - torch_dtype=devices.dtype, - ).to(devices.device) - up_pipe = LTX2LatentUpsamplePipeline(vae=shared.sd_model.vae, latent_upsampler=latent_upsampler) + global upsample_pipe_2x # pylint: disable=global-statement + upsample_repo = upsample_repo_id_23 if caps.variant == '2.3' else upsample_repo_id_20 + upsample_pipe_2x = load_upsample_2x(upsample_pipe_2x, upsample_repo) + upsample_pipe_2x = sd_models.apply_balanced_offload(upsample_pipe_2x) # 2.x base pass returns denormalized latents; latents_normalized=False tells the # upsampler "already raw, do not denormalize again". up_args = { @@ -420,12 +411,10 @@ def run_ltx(task_id, } if latents.ndim == 4: latents = latents.unsqueeze(0) - log.debug(f'Video: op=upsample family=2.x latents={latents.shape} {up_args}') + log.debug(f'Video: op=upsample family=2.x latents={latents.shape} auto={auto_refine_upsample} {up_args}') yield None, 'LTX: Upsample in progress...' - latents = up_pipe(latents=latents, **up_args).frames[0] - latent_upsampler.to('cpu') - del up_pipe, latent_upsampler - devices.torch_gc(force=True, reason='ltx:upsample') + latents = upsample_pipe_2x(latents=latents, **up_args).frames[0] + upsample_pipe_2x = sd_models.apply_balanced_offload(upsample_pipe_2x) except AssertionError as e: yield from abort(e, ok=True, p=p) return diff --git a/modules/ltx/ltx_util.py b/modules/ltx/ltx_util.py index 310a865fb..b404879f7 100644 --- a/modules/ltx/ltx_util.py +++ b/modules/ltx/ltx_util.py @@ -40,7 +40,7 @@ def load_upsample(upsample_pipe, upsample_repo_id): if upsample_pipe is None: t0 = time.time() from diffusers.pipelines.ltx.pipeline_ltx_latent_upsample import LTXLatentUpsamplePipeline - log.info(f'Video load: cls={LTXLatentUpsamplePipeline.__class__.__name__} repo="{upsample_repo_id}"') + log.info(f'Video load: cls={LTXLatentUpsamplePipeline.__name__} repo="{upsample_repo_id}"') upsample_pipe = LTXLatentUpsamplePipeline.from_pretrained( upsample_repo_id, vae=shared.sd_model.vae, @@ -52,6 +52,34 @@ def load_upsample(upsample_pipe, upsample_repo_id): return upsample_pipe +def load_upsample_2x(upsample_pipe, upsample_repo_id): + # 2.x ships the upsampler as a bare nn.Module in a subfolder; no from_pretrained on the + # pipeline wrapper, so we load the model + construct the pipeline manually. + if upsample_pipe is None: + t0 = time.time() + from diffusers.pipelines.ltx2.pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline + from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel + from modules import sd_checkpoint + log.info(f'Video load: cls={LTX2LatentUpsamplePipeline.__name__} repo="{upsample_repo_id}"') + latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained( + upsample_repo_id, + subfolder='latent_upsampler', + cache_dir=shared.opts.hfcache_dir, + torch_dtype=devices.dtype, + ).to(devices.device) + upsample_pipe = LTX2LatentUpsamplePipeline( + vae=shared.sd_model.vae, + latent_upsampler=latent_upsampler, + ) + # Synthetic checkpoint_info gives this pipe its own OffloadHook cache slot, so routing + # it through apply_balanced_offload does not invalidate the main pipe's module map + # (sd_offload.py:488 keys on sd_checkpoint_info.name). + upsample_pipe.sd_checkpoint_info = sd_checkpoint.CheckpointInfo('ltx-upsampler-2.x') + t1 = time.time() + timer.process.add('load', t1 - t0) + return upsample_pipe + + def _condition_cls(family: str): if family == '2.x': try: diff --git a/modules/video_models/video_overrides.py b/modules/video_models/video_overrides.py index 34ee151d2..582acc92e 100644 --- a/modules/video_models/video_overrides.py +++ b/modules/video_models/video_overrides.py @@ -17,9 +17,11 @@ def load_override(selected: Model, **load_args): # LTX if 'LTXVideo 0.9.5 I2V' in selected.name: kwargs['vae'] = diffusers.AutoencoderKLLTXVideo.from_pretrained(selected.repo, subfolder="vae", torch_dtype=torch.float32, cache_dir=shared.opts.hfcache_dir, **load_args) - # OzzyGT LTX-2.3 mirrors ship connectors twice: sharded safetensors + .index.json plus a - # redundant unsharded diffusion_pytorch_model.safetensors of the same weights. Diffusers - # fetches both but loads sharded; skip the ~6.3 GB duplicate. + # OzzyGT LTX-2.3 mirrors pack connectors/ twice by design: sharded (*-00001-of-0000N + + # .index.json) and unsharded diffusion_pytorch_model.safetensors of the byte-identical + # weights. snapshot_download faithfully fetches both; diffusers' component loader picks + # sharded when the index is present. ignore_patterns skips the ~6.3 GB unsharded copy + # without reaching for a cleaner upstream mirror. ltx2_redundant_connector_repos = { 'OzzyGT/LTX-2.3', 'OzzyGT/LTX-2.3-sdnq-dynamic-int4',