fix(ltx2-sr): close port gaps surfaced by alignment harness retries

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

@@ -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

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

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()
 

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 = {

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 "

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,

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": {