From 1eebbe4da69d5ac69d1f28ee796ee69f88bcd825 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sun, 26 Apr 2026 14:23:02 -0700 Subject: [PATCH] fix(ltx2-sr): close port gaps surfaced by alignment harness retries MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit LTX-2 SR pipeline now passes module load + stage graph build on the public side. Each retry of scripts/ltx2_sr_alignment.py surfaced one more gap; this commit closes them so the pipeline progresses to the upsample → stage-2 → decode tail. * harness: PipelineSelection (real class name; PipelineConfig was a typo) * FastVideoArgs: generic refine_* fields (preferred user-facing API) * LTX2Pipeline.initialize_pipeline: getattr-with-default for debug_model_sums / debug_model_detail (public args lack the debug fields; LTX-2 still flips the env-var toggles when set) * VAEConfig: use_temporal_scaling_frames default True (LTX-2 latent prep reads this in the schedule-shift gate) * LTX2LatentPreparationStage._randn_ltx2_video_latents: switch to randn_tensor so list-of-generators batches work (matches internal sampling order) * ModelRegistry: register LTX2LatentUpsampler in _UPSAMPLERS * component_loader: register spatial_upsampler / temporal_upsampler with UpsamplerLoader (was hitting the generic stub loader); UpsamplerLoader gains LTX-2 fallback when pipeline_config.upsampler_config isn't a multi-config tuple, plus the 'model.'-prefix-stripping weight load path the internal version uses 7 retries -> public side now reaches the upsample stage. Co-authored-by: Junda (David) Su <90978028+Davids048@users.noreply.github.com> Co-authored-by: Matthew Noto <99706358+RandNMR73@users.noreply.github.com> Co-authored-by: XOR-op <17672363+XOR-op@users.noreply.github.com> Co-authored-by: Zhang Peiyuan <42993249+jzhang38@users.noreply.github.com> --- fastvideo/configs/models/vaes/base.py | 4 ++ fastvideo/fastvideo_args.py | 15 +++++ fastvideo/models/loader/component_loader.py | 63 ++++++++++++++----- fastvideo/models/registry.py | 2 + .../pipelines/basic/ltx2/ltx2_pipeline.py | 19 ++++-- .../ltx2/stages/ltx2_latent_preparation.py | 6 +- scripts/ltx2_sr_alignment.py | 4 +- 7 files changed, 90 insertions(+), 23 deletions(-) diff --git a/fastvideo/configs/models/vaes/base.py b/fastvideo/configs/models/vaes/base.py index 23192b569..864b1be01 100644 --- a/fastvideo/configs/models/vaes/base.py +++ b/fastvideo/configs/models/vaes/base.py @@ -37,6 +37,10 @@ class VAEConfig(ModelConfig): use_tiling: bool = True use_temporal_tiling: bool = True use_parallel_tiling: bool = True + # When True, latent preparation skips the schedule shift on frames + # whose temporal index is below the model's first-frame conditioning + # threshold. LTX-2 reads this in the latent prep stage. + use_temporal_scaling_frames: bool = True def __post_init__(self): self.blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index d26c68a43..21dd17955 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -162,6 +162,21 @@ class FastVideoArgs: ltx2_vae_temporal_tile_overlap_in_frames: int | None = None ltx2_initial_latent_path: str | None = None ltx2_audio_latent_path: str | None = None + # Generic stage-2 refine surface (preferred user-facing API). The + # ltx2_refine_* fields below remain the runtime carriers; these + # generic ones let CLI / typed-config callers set the same values + # without binding to a specific model family. ``None`` here means + # "fall back to the model_index.json default and/or the + # ltx2_refine_* runtime carrier". + refine_enabled: bool | None = None + refine_upsampler_path: str | None = None + refine_transformer_path: str | None = None + refine_lora_path: str | None = None + refine_num_inference_steps: int | None = None + refine_guidance_scale: float | None = None + refine_add_noise: bool | None = None + refine_noise_path: str | None = None + refine_audio_noise_path: str | None = None # LTX-2 stage-2 spatial refinement (the SR pipeline). When enabled the # transformer runs once at half resolution, the latents are upsampled # by the LTX2 latent upsampler, then a short stage-2 distilled diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 5a2145fe8..b91776c86 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -99,6 +99,12 @@ class ComponentLoader(ABC): # NumberConditioners; not a pure text encoder, so it gets # its own loader. "conditioner": (ConditionerLoader, "fastvideo"), + # LTX-2 spatial / temporal upsamplers — share the + # UpsamplerLoader path with the upsampler/upsampler_2 keys + # so the SR pipeline picks up real weights instead of the + # generic config-only loader. + "spatial_upsampler": (UpsamplerLoader, "diffusers"), + "temporal_upsampler": (UpsamplerLoader, "diffusers"), } if module_type in module_loaders: @@ -1058,7 +1064,7 @@ class ConditionerLoader(ComponentLoader): class UpsamplerLoader(ComponentLoader): - """Loader for upsamplers.""" + """Loader for upsamplers (incl. LTX-2 spatial/temporal upsamplers).""" def load(self, model_path: str, fastvideo_args: FastVideoArgs): """Load the upsampler based on the model path, and inference args.""" @@ -1068,36 +1074,65 @@ class UpsamplerLoader(ComponentLoader): if class_name is None: raise ValueError( "Model config does not contain a _class_name attribute. " - "Only diffusers format is supported." - ) + "Only diffusers format is supported.") - try: - upsampler_cfg = deepcopy(fastvideo_args.pipeline_config.upsampler_config[0]) - upsampler_cfg.update_model_config(config_dict) - except Exception as e: - upsampler_cfg = deepcopy(fastvideo_args.pipeline_config.upsampler_config[1]) - upsampler_cfg.update_model_config(config_dict) + # The base PipelineConfig declares ``upsampler_config`` as a + # single ``UpsamplerConfig`` instance, but Hunyuan15 narrows it + # to a tuple of two configs (one per SR target). We only treat + # the attribute as a multi-config when it actually is one; + # otherwise the LTX-2 branch below handles the single-class + # path that takes the diffusers config dict directly. + upsampler_config_attr = getattr(fastvideo_args.pipeline_config, + "upsampler_config", None) + if isinstance(upsampler_config_attr, list | tuple): + try: + upsampler_cfg = deepcopy(upsampler_config_attr[0]) + upsampler_cfg.update_model_config(config_dict) + except Exception: + upsampler_cfg = deepcopy(upsampler_config_attr[1]) + upsampler_cfg.update_model_config(config_dict) + elif class_name == "LTX2LatentUpsampler": + # LTX-2 pipeline_config does not declare upsampler_config; the + # `LTX2LatentUpsampler` wrapper takes the raw diffusers config + # dict directly via LatentUpsamplerConfigurator. + upsampler_cfg = deepcopy(config_dict) + else: + raise AttributeError( + "pipeline_config.upsampler_config is missing; cannot build " + f"upsampler config for class {class_name}") model_cls, _ = ModelRegistry.resolve_model_cls(class_name) model = model_cls(upsampler_cfg) target_device = get_local_torch_device() - model = model.to(target_device, dtype=PRECISION_TO_TYPE[fastvideo_args.pipeline_config.upsampler_precision]) + upsampler_precision = getattr(fastvideo_args.pipeline_config, + "upsampler_precision", "bf16") + model = model.to(target_device, + dtype=PRECISION_TO_TYPE[upsampler_precision]) - # Find all safetensors files safetensors_list = glob.glob( os.path.join(str(model_path), "*.safetensors")) if not safetensors_list: raise ValueError(f"No safetensors files found in {model_path}") - + if len(safetensors_list) == 1: loaded = safetensors_load_file(safetensors_list[0]) else: loaded = {} for sf_file in safetensors_list: loaded.update(safetensors_load_file(sf_file)) - - model.load_state_dict(loaded, strict=True) + + # The LTX-2 latent upsampler wrapper exposes the actual conv + # stack at ``self.model``; checkpoint state_dicts may be saved + # without the ``model.`` prefix when the inner module was + # serialised directly. Strip / forward as needed so both layouts + # load cleanly. + target_module = getattr(model, "model", model) + if loaded and all(k.startswith("model.") for k in loaded): + stripped = {k[len("model."):]: v for k, v in loaded.items()} + target_module.load_state_dict(stripped, strict=True) + else: + target_module.load_state_dict(loaded, strict=True) return model.eval() diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index a298cf667..ac978d83e 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -115,6 +115,8 @@ _SCHEDULERS = { _UPSAMPLERS = { "SRTo720pUpsampler": ("upsamplers", "hunyuan15", "SRTo720pUpsampler"), "SRTo1080pUpsampler": ("upsamplers", "hunyuan15", "SRTo1080pUpsampler"), + "LTX2LatentUpsampler": + ("upsamplers", "ltx2_upsampler", "LTX2LatentUpsampler"), } _LEGACY_FAST_VIDEO_MODELS = { diff --git a/fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py b/fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py index 6ec34bce0..a07bedf02 100644 --- a/fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py +++ b/fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py @@ -142,10 +142,16 @@ class LTX2Pipeline(LoRAPipeline): ) def initialize_pipeline(self, fastvideo_args: FastVideoArgs): - if fastvideo_args.debug_model_sums: + # Optional debug-instrumentation env vars. The internal FastVideoArgs + # carries debug_model_sums / debug_model_detail toggles; the public + # args don't define them yet, so getattr-with-default keeps the + # pipeline runnable on both shapes without forcing the public args + # to add fields a public consumer wouldn't otherwise touch. + if getattr(fastvideo_args, "debug_model_sums", False): os.environ["LTX2_PIPELINE_DEBUG_LOG"] = "1" - if fastvideo_args.debug_model_sums_path: - os.environ["LTX2_PIPELINE_DEBUG_PATH"] = (fastvideo_args.debug_model_sums_path) + sums_path = getattr(fastvideo_args, "debug_model_sums_path", None) + if sums_path: + os.environ["LTX2_PIPELINE_DEBUG_PATH"] = sums_path else: logger.warning("debug_model_sums is enabled but " "debug_model_sums_path is not set; no model sums " @@ -153,10 +159,11 @@ class LTX2Pipeline(LoRAPipeline): else: os.environ.pop("LTX2_PIPELINE_DEBUG_LOG", None) os.environ.pop("LTX2_PIPELINE_DEBUG_PATH", None) - if fastvideo_args.debug_model_detail: + if getattr(fastvideo_args, "debug_model_detail", False): os.environ["LTX2_DEBUG_DETAIL"] = "1" - if fastvideo_args.debug_model_detail_path: - os.environ["LTX2_PIPELINE_DEBUG_DETAIL_PATH"] = (fastvideo_args.debug_model_detail_path) + detail_path = getattr(fastvideo_args, "debug_model_detail_path", None) + if detail_path: + os.environ["LTX2_PIPELINE_DEBUG_DETAIL_PATH"] = detail_path else: logger.warning("debug_model_detail is enabled but " "debug_model_detail_path is not set; no detailed " diff --git a/fastvideo/pipelines/basic/ltx2/stages/ltx2_latent_preparation.py b/fastvideo/pipelines/basic/ltx2/stages/ltx2_latent_preparation.py index eabf967d5..94a643abc 100644 --- a/fastvideo/pipelines/basic/ltx2/stages/ltx2_latent_preparation.py +++ b/fastvideo/pipelines/basic/ltx2/stages/ltx2_latent_preparation.py @@ -56,7 +56,11 @@ def _randn_ltx2_video_latents( patchifier.get_token_count(video_shape), shape[1] * patch_volume, ) - patch_noise = torch.randn( + # `randn_tensor` accepts both a single torch.Generator and a list + # (one per batch element); `torch.randn` only accepts the former. + # The internal version uses `randn_tensor` so the sampling order + # matches the official LTX-2 reference; mirror that here. + patch_noise = randn_tensor( patch_shape, generator=generator, device=device, diff --git a/scripts/ltx2_sr_alignment.py b/scripts/ltx2_sr_alignment.py index 3387afb0c..7327fdde8 100644 --- a/scripts/ltx2_sr_alignment.py +++ b/scripts/ltx2_sr_alignment.py @@ -140,12 +140,12 @@ def _run_typed(args: argparse.Namespace) -> RunResult: from fastvideo import VideoGenerator # noqa: PLC0415 from fastvideo.api import ( # noqa: PLC0415 ComponentConfig, EngineConfig, GenerationRequest, GeneratorConfig, - InputConfig, OutputConfig, PipelineConfig, SamplingConfig) + InputConfig, OutputConfig, PipelineSelection, SamplingConfig) config = GeneratorConfig( model_path=args.model, engine=EngineConfig(num_gpus=1, ), - pipeline=PipelineConfig( + pipeline=PipelineSelection( components=ComponentConfig(), preset_overrides={ "refine": {