component_loader ¶

@classmethod
def for_module_type(
    cls, module_type: str, transformers_or_diffusers: str
) -> "ComponentLoader":
    """
    Factory method to create a component loader for a specific module type.

    Args:
        module_type: Type of module (e.g., "vae", "text_encoder", "transformer", "scheduler")
        transformers_or_diffusers: Whether the module is from transformers or diffusers

    Returns:
        A component loader for the specified module type
    """
    # Map of module types to their loader classes and expected library
    module_loaders = {
        "scheduler": (SchedulerLoader, "diffusers"),
        "transformer": (TransformerLoader, "diffusers"),
        "sr_transformer": (TransformerLoader, "diffusers"),
        "transformer_2": (TransformerLoader, "diffusers"),
        "transformer_3": (TransformerLoader, "diffusers"),
        "vae": (VAELoader, "diffusers"),
        "light_vae": (VAELoader, "diffusers"),
        "audio_vae": (AudioDecoderLoader, "diffusers"),
        "audio_decoder": (AudioDecoderLoader, "diffusers"),
        "vocoder": (VocoderLoader, "diffusers"),
        "text_encoder": (TextEncoderLoader, "transformers"),
        "text_encoder_2": (TextEncoderLoader, "transformers"),
        "text_encoder_3": (TextEncoderLoader, "transformers"),
        "tokenizer": (TokenizerLoader, "transformers"),
        "tokenizer_2": (TokenizerLoader, "transformers"),
        "tokenizer_3": (TokenizerLoader, "transformers"),
        "image_processor": (ImageProcessorLoader, "transformers"),
        "feature_extractor": (ImageProcessorLoader, "transformers"),
        "image_encoder": (ImageEncoderLoader, "transformers"),
        "vision_language_encoder": (VisionLanguageEncoderLoader, "transformers"),
        "processor": (ProcessorLoader, "transformers"),
        "upsampler": (UpsamplerLoader, "diffusers"),
        "upsampler_2": (UpsamplerLoader, "diffusers"),
        # Stable Audio's `StableAudioMultiConditioner` bundles T5 +
        # 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:
        loader_cls, expected_library = module_loaders[module_type]
        # Allow fastvideo.* libraries for custom implementations (e.g. Cosmos2_5Pipeline)
        # that aren't available in diffusers/transformers yet
        is_fastvideo_module = transformers_or_diffusers.startswith("fastvideo.")
        if not is_fastvideo_module:
            # Assert that the library matches what's expected for this module type
            assert transformers_or_diffusers == expected_library, f"{module_type} must be loaded from {expected_library}, got {transformers_or_diffusers}"
        return loader_cls()

    # For unknown module types, use a generic loader
    logger.warning(
        "No specific loader found for module type: %s. Using generic loader.",
        module_type,
    )
    return GenericComponentLoader(transformers_or_diffusers)

@abstractmethod
def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """
    Load the component based on the model path, architecture, and inference args.

    Args:
        model_path: Path to the component model
        fastvideo_args: FastVideoArgs

    Returns:
        The loaded component
    """
    raise NotImplementedError

def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """Load a generic component based on the model path, and inference args."""
    logger.warning(
        "Using generic loader for %s with library %s",
        model_path,
        self.library,
    )

    if self.library == "transformers":
        from transformers import AutoModel

        model = AutoModel.from_pretrained(
            model_path,
            trust_remote_code=fastvideo_args.trust_remote_code,
            revision=fastvideo_args.revision,
        )
        logger.info(
            "Loaded generic transformers model: %s",
            model.__class__.__name__,
        )
        return model
    elif self.library == "diffusers":
        logger.warning(
            "Generic loading for diffusers components is not fully implemented"
        )

        model_config = get_diffusers_config(model=model_path)
        logger.info("Diffusers Model config: %s", model_config)
        # This is a placeholder - in a real implementation, you'd need to handle this properly
        return None
    else:
        raise ValueError(f"Unsupported library: {self.library}")

def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """Load the text encoders based on the model path, and inference args."""
    # model_config: PretrainedConfig = get_hf_config(
    #     model=model_path,
    #     trust_remote_code=fastvideo_args.trust_remote_code,
    #     revision=fastvideo_args.revision,
    #     model_override_args=None,
    # )
    with open(os.path.join(model_path, "config.json")) as f:
        model_config = json.load(f)
    model_config.pop("_name_or_path", None)
    model_config.pop("transformers_version", None)
    model_config.pop("torch_dtype", None)
    model_config.pop("model_type", None)
    logger.info("HF Model config: %s", model_config)

    encoder_config = fastvideo_args.pipeline_config.image_encoder_config
    encoder_config.update_model_arch(model_config)

    from fastvideo.platforms import current_platform

    if fastvideo_args.image_encoder_cpu_offload:
        target_device = (
            torch.device("mps")
            if current_platform.is_mps()
            else torch.device("cpu")
        )
    else:
        target_device = get_local_torch_device()
    # TODO(will): add support for other dtypes
    return self.load_model(
        model_path,
        encoder_config,
        target_device,
        fastvideo_args,
        fastvideo_args.pipeline_config.image_encoder_precision,
    )

def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """Load the image processor based on the model path, and inference args."""
    logger.info("Loading image processor from %s", model_path)

    image_processor = AutoImageProcessor.from_pretrained(
        model_path,
    )
    logger.info(
        "Loaded image processor: %s", image_processor.__class__.__name__
    )
    return image_processor

@staticmethod
def load_module(
    module_name: str,
    component_model_path: str,
    transformers_or_diffusers: str,
    fastvideo_args: FastVideoArgs,
):
    """
    Load a pipeline module.

    Args:
        module_name: Name of the module (e.g., "vae", "text_encoder", "transformer", "scheduler")
        component_model_path: Path to the component model
        transformers_or_diffusers: Whether the module is from transformers or diffusers
        pipeline_args: Inference arguments

    Returns:
        The loaded module
    """
    logger.info(
        "Loading %s using %s from %s",
        module_name,
        transformers_or_diffusers,
        component_model_path,
    )

    # Get the appropriate loader for this module type
    loader = ComponentLoader.for_module_type(
        module_name, transformers_or_diffusers
    )

    # Load the module
    return loader.load(component_model_path, fastvideo_args)

def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """Load the scheduler based on the model path, and inference args."""
    config = get_diffusers_config(model=model_path)

    class_name = config.pop("_class_name")
    assert class_name is not None, (
        "Model config does not contain a _class_name attribute. Only diffusers format is supported."
    )

    scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name)

    scheduler = scheduler_cls(**config)
    if fastvideo_args.pipeline_config.flow_shift is not None:
        scheduler.set_shift(fastvideo_args.pipeline_config.flow_shift)
    return scheduler

def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """Load the text encoders based on the model path, and inference args."""
    # model_config: PretrainedConfig = get_hf_config(
    #     model=model_path,
    #     trust_remote_code=fastvideo_args.trust_remote_code,
    #     revision=fastvideo_args.revision,
    #     model_override_args=None,
    # )
    model_config = get_diffusers_config(model=model_path)
    model_config.pop("_name_or_path", None)
    model_config.pop("transformers_version", None)
    model_config.pop("model_type", None)
    model_config.pop("tokenizer_class", None)
    model_config.pop("torch_dtype", None)
    repo_root = os.path.dirname(model_path)
    index_path = os.path.join(repo_root, "model_index.json")
    gemma_path = ""
    gemma_path_from_candidate = False
    if os.path.isfile(index_path):
        try:
            with open(index_path, encoding="utf-8") as f:
                model_index = json.load(f)
            gemma_path = model_index.get("gemma_model_path", "")
        except json.JSONDecodeError:
            gemma_path = ""
    if not gemma_path:
        candidate = os.path.normpath(os.path.join(model_path, "gemma"))
        if os.path.isdir(candidate):
            gemma_path = candidate
            gemma_path_from_candidate = True
            model_config["gemma_model_path"] = gemma_path
    if gemma_path and not gemma_path_from_candidate:
        if not os.path.isabs(gemma_path):
            model_config["gemma_model_path"] = os.path.normpath(
                os.path.join(repo_root, gemma_path)
            )
    transformer_config_path = os.path.join(
        repo_root, "transformer", "config.json"
    )
    if os.path.isfile(transformer_config_path):
        try:
            with open(transformer_config_path, encoding="utf-8") as f:
                transformer_config = json.load(f)
            if (
                "connector_double_precision_rope" not in model_config
                or not model_config["connector_double_precision_rope"]
            ):
                if transformer_config.get("double_precision_rope") is True:
                    model_config["connector_double_precision_rope"] = True
            if "connector_rope_type" not in model_config:
                rope_type = transformer_config.get("rope_type")
                if rope_type is not None:
                    model_config["connector_rope_type"] = rope_type
        except json.JSONDecodeError:
            pass
    logger.info("HF Model config: %s", model_config)

    base = os.path.basename(os.path.normpath(model_path))
    idx = 0
    if base.startswith("text_encoder_"):
        try:
            idx = int(base.split("_")[-1]) - 1
        except Exception:
            idx = 0
    encoder_configs = fastvideo_args.pipeline_config.text_encoder_configs
    encoder_precisions = fastvideo_args.pipeline_config.text_encoder_precisions
    if idx < 0 or idx >= len(encoder_configs):
        raise IndexError(
            f"text encoder index {idx} out of range for text_encoder_configs (len={len(encoder_configs)}), model_path={model_path}"
        )
    encoder_config = encoder_configs[idx]
    if (
        model_config.get("architectures") == ["CLIPModel"]
        and isinstance(model_config.get("text_config"), dict)
    ):
        valid_arch_fields = {
            f.name for f in dataclasses.fields(encoder_config.arch_config)
        }
        model_config = {
            key: value
            for key, value in deepcopy(model_config["text_config"]).items()
            if key in valid_arch_fields
        }
        model_config["architectures"] = ["CLIPTextModel"]
    encoder_config.update_model_arch(model_config)
    if idx < 0 or idx >= len(encoder_precisions):
        raise IndexError(
            f"text encoder index {idx} out of range for text_encoder_precisions (len={len(encoder_precisions)}), model_path={model_path}"
        )
    encoder_precision = encoder_precisions[idx]

    target_device = get_local_torch_device()
    # TODO(will): add support for other dtypes
    return self.load_model(
        model_path,
        encoder_config,
        target_device,
        fastvideo_args,
        encoder_precision,
        use_text_encoder_override=True,
    )

def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """Load the tokenizer based on the model path, and inference args."""
    logger.info("Loading tokenizer from %s", model_path)
    resolved_model_path = model_path

    # LTX2 checkpoints may not ship a top-level tokenizer/ directory.
    # In that case, tokenizer assets live under text_encoder/gemma/.
    if not os.path.isdir(resolved_model_path):
        ltx2_gemma_path = os.path.normpath(
            os.path.join(resolved_model_path, "..", "text_encoder",
                         "gemma"))
        if os.path.isdir(ltx2_gemma_path):
            logger.info(
                "Tokenizer directory %s missing; falling back to %s",
                resolved_model_path,
                ltx2_gemma_path,
            )
            resolved_model_path = ltx2_gemma_path

    # Cosmos2.5 stores an AutoProcessor config in `tokenizer/config.json` (not a tokenizer
    # config). Use its `_name_or_path` (e.g. Qwen/Qwen2.5-VL-7B-Instruct) as the source.
    tokenizer_cfg_path = os.path.join(resolved_model_path, "config.json")
    if os.path.exists(tokenizer_cfg_path):
        try:
            with open(tokenizer_cfg_path, "r") as f:
                tokenizer_cfg = json.load(f)
            if isinstance(tokenizer_cfg, dict) and (
                tokenizer_cfg.get("_class_name") == "AutoProcessor"
                or "processor_type" in tokenizer_cfg
            ):
                src = tokenizer_cfg.get("_name_or_path", "")
                if isinstance(src, str) and src.strip():
                    processor = AutoProcessor.from_pretrained(
                        src.strip(),
                        trust_remote_code=True,
                    )
                    logger.info(
                        "Loaded tokenizer/processor from %s: %s",
                        src,
                        processor.__class__.__name__,
                    )
                    return processor
        except Exception:
            # If parsing fails, fall through to AutoTokenizer below.
            pass

    # Only Flux2 full's Mistral3 (require_processor=True) must load via
    # AutoProcessor. Gate the processor_config.json shortcut on that flag so
    # existing encoders (e.g. HunyuanVideo 1.5 / Qwen2.5-VL) stay on the
    # historical AutoTokenizer path below even if their tokenizer dir happens
    # to ship a processor_config.json.
    require_processor = False
    if hasattr(fastvideo_args.pipeline_config, "text_encoder_configs"):
        try:
            require_processor = any(
                getattr(getattr(cfg, "arch_config", None), "require_processor", False)
                for cfg in fastvideo_args.pipeline_config.text_encoder_configs)
        except Exception:
            require_processor = False

    if require_processor and os.path.exists(os.path.join(resolved_model_path, "processor_config.json")):
        processor = AutoProcessor.from_pretrained(
            resolved_model_path,
            local_files_only=os.path.isdir(resolved_model_path),
            trust_remote_code=fastvideo_args.trust_remote_code,
        )
        logger.info(
            "Loaded tokenizer/processor from %s: %s",
            resolved_model_path,
            processor.__class__.__name__,
        )
        return processor

    try:
        tokenizer = AutoTokenizer.from_pretrained(
            resolved_model_path,  # "<path to model>/tokenizer"
            # in v0, this was same string as encoder_name "ClipTextModel"
            # TODO(will): pass these tokenizer kwargs from inference args? Maybe
            # other method of config?
            local_files_only=os.path.isdir(resolved_model_path),
        )
    except (OSError, ValueError):
        tokenizer = AutoProcessor.from_pretrained(
            resolved_model_path,
            local_files_only=os.path.isdir(resolved_model_path),
            trust_remote_code=fastvideo_args.trust_remote_code,
        )
        logger.info(
            "Loaded tokenizer/processor from %s: %s",
            resolved_model_path,
            tokenizer.__class__.__name__,
        )
        return tokenizer
    padding_side = None
    if hasattr(fastvideo_args.pipeline_config, "text_encoder_configs"):
        try:
            arch_config = fastvideo_args.pipeline_config.text_encoder_configs[
                0
            ].arch_config
            padding_side = getattr(arch_config, "padding_side", None)
        except Exception:
            padding_side = None
    if padding_side:
        tokenizer.padding_side = padding_side
    if tokenizer.pad_token is None and tokenizer.eos_token is not None:
        tokenizer.pad_token = tokenizer.eos_token
    logger.info("Loaded tokenizer: %s", tokenizer.__class__.__name__)
    return tokenizer

def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """Load the transformer based on the model path, and inference args."""
    config = get_diffusers_config(model=model_path)
    hf_config = deepcopy(config)
    cls_name = config.pop("_class_name")
    config.pop("_name_or_path", None)
    if cls_name is None:
        raise ValueError(
            "Model config does not contain a _class_name attribute. "
            "Only diffusers format is supported."
        )

    logger.info("transformer cls_name: %s", cls_name)
    if fastvideo_args.override_transformer_cls_name is not None:
        cls_name = fastvideo_args.override_transformer_cls_name
        logger.info("Overriding transformer cls_name to %s", cls_name)

    fastvideo_args.model_paths["transformer"] = model_path

    # Config from Diffusers supersedes fastvideo's model config
    dit_config = deepcopy(fastvideo_args.pipeline_config.dit_config)
    dit_config.update_model_arch(config)

    # Generator-only QAT for DMD distillation: the teacher (real_score) and
    # critic (fake_score) transformers load with this flag set and must stay
    # full precision. Drop the nvfp4_qat quant from their copied config, and
    # mask the global ATTN_QAT_TRAIN env so their attention falls back to dense
    # (the backend is read globally at build time). The generator loads without
    # the flag and keeps both.
    _qat_generator_only = hasattr(fastvideo_args, "_loading_teacher_critic_model")
    _qat_prev_attn_env = None
    if _qat_generator_only:
        dit_config.quant_config = None
        from fastvideo.attention.selector import _cached_get_attn_backend
        _qat_prev_attn_env = os.environ.pop("FASTVIDEO_ATTENTION_BACKEND", None)
        _cached_get_attn_backend.cache_clear()

    model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)

    # 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}")

    # arch_config can infer architecture from weight keys (e.g. Flux2 layer counts)
    update_fn = getattr(dit_config.arch_config, "update_from_weight_keys", None)
    if callable(update_fn):
        weight_keys = _collect_safetensors_keys(safetensors_list)
        update_fn(weight_keys)

    # Check if we should use custom initialization weights
    custom_weights_path = getattr(
        fastvideo_args, "init_weights_from_safetensors", None
    )
    use_custom_weights = (
        custom_weights_path
        and os.path.exists(custom_weights_path)
        and not hasattr(fastvideo_args, "_loading_teacher_critic_model")
    )

    if use_custom_weights:
        if "transformer_2" in model_path:
            custom_weights_path = getattr(
                fastvideo_args, "init_weights_from_safetensors_2", None
            )
        assert custom_weights_path is not None, (
            "Custom initialization weights must be provided"
        )
        if os.path.isdir(custom_weights_path):
            safetensors_list = glob.glob(
                os.path.join(str(custom_weights_path), "*.safetensors")
            )
        else:
            assert custom_weights_path.endswith(".safetensors"), (
                "Custom initialization weights must be a safetensors file"
            )
            safetensors_list = [custom_weights_path]

    logger.info(
        "Loading model from %s safetensors files: %s",
        len(safetensors_list),
        safetensors_list,
    )

    default_dtype = PRECISION_TO_TYPE[
        fastvideo_args.pipeline_config.dit_precision
    ]

    # Load the model using FSDP loader
    logger.info("Loading model from %s, default_dtype: %s", cls_name,
                default_dtype)
    assert fastvideo_args.hsdp_shard_dim is not None
    # Cosmos2.5 checkpoints can include extra entries not present in the
    # instantiated model (e.g. pos_embedder ranges / *_extra_state). Load
    # non-strictly for Cosmos2.5 only; keep upstream strict behavior for others.
    strict_load = not (
        cls_name.startswith("Cosmos25")
        or cls_name == "Cosmos25Transformer3DModel"
        or getattr(fastvideo_args.pipeline_config, "prefix", "") == "Cosmos25"
    )
    model = maybe_load_fsdp_model(
        model_cls=model_cls,
        init_params={"config": dit_config, "hf_config": hf_config},
        weight_dir_list=safetensors_list,
        device=get_local_torch_device(),
        hsdp_replicate_dim=fastvideo_args.hsdp_replicate_dim,
        hsdp_shard_dim=fastvideo_args.hsdp_shard_dim,
        strict=strict_load,
        cpu_offload=fastvideo_args.dit_cpu_offload,
        pin_cpu_memory=fastvideo_args.pin_cpu_memory,
        fsdp_inference=fastvideo_args.use_fsdp_inference,
        # TODO(will): make these configurable
        default_dtype=default_dtype,
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.float32,
        output_dtype=None,
        training_mode=fastvideo_args.training_mode,
        enable_torch_compile=fastvideo_args.enable_torch_compile,
        torch_compile_kwargs=fastvideo_args.torch_compile_kwargs,
    )

    if _qat_generator_only:
        from fastvideo.attention.selector import _cached_get_attn_backend
        if _qat_prev_attn_env is not None:
            os.environ["FASTVIDEO_ATTENTION_BACKEND"] = _qat_prev_attn_env
        _cached_get_attn_backend.cache_clear()

    total_params = sum(p.numel() for p in model.parameters())
    logger.info("Loaded model with %.2fB parameters", total_params / 1e9)

    assert next(model.parameters()).dtype == default_dtype, (
        "Model dtype does not match default dtype"
    )

    model = model.eval()

    if fastvideo_args.inference_mode and fastvideo_args.dit_layerwise_offload:
        # Check if model has nn.ModuleList for layerwise offload compatibility
        has_module_list = any(
            isinstance(m, nn.ModuleList) for m in model.children()
        )
        if has_module_list:
            enable_layerwise_offload(model)
        else:
            logger.warning(
                "Layerwise offload requested but model %s does not have "
                "nn.ModuleList structure. Skipping layerwise offload.",
                cls_name
            )
    return model

def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """Load the upsampler based on the model path, and inference args."""
    config_dict = get_diffusers_config(model=model_path)
    class_name = config_dict.pop("_class_name", None)

    if class_name is None:
        raise ValueError(
            "Model config does not contain a _class_name attribute. "
            "Only diffusers format is supported.")

    # 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()
    upsampler_precision = getattr(fastvideo_args.pipeline_config,
                                  "upsampler_precision", "bf16")
    model = model.to(target_device,
                     dtype=PRECISION_TO_TYPE[upsampler_precision])

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

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

def load(self, model_path: str, fastvideo_args: FastVideoArgs):
    """Load the VAE based on the model path, and inference args."""
    config = get_diffusers_config(model=model_path)
    class_name = config.pop("_class_name")
    config.pop("_name_or_path", None)
    assert class_name is not None, (
        "Model config does not contain a _class_name attribute. Only diffusers format is supported."
    )
    fastvideo_args.model_paths["vae"] = model_path

    from fastvideo.platforms import current_platform

    if fastvideo_args.vae_cpu_offload:
        target_device = (
            torch.device("mps")
            if current_platform.is_mps()
            else torch.device("cpu")
        )
    else:
        target_device = get_local_torch_device()

    with set_default_torch_dtype(
        PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
        if fastvideo_args.pipeline_config.vae_precision
        else torch.bfloat16
    ):
        pipeline_name = fastvideo_args.pipeline_config.__class__.__name__
        is_gen3c = pipeline_name.startswith("Gen3C")
        is_cosmos25 = pipeline_name == "Cosmos25Config"

        # GEN3C: prefer tokenizer-backed VAE checkpoint when available.
        # This aligns latent conditioning with the GEN3C temporal contract.
        if is_gen3c and class_name in (
                "AutoencoderKLWan", "AutoencoderKLGen3CTokenizer"):
            from fastvideo.models.vaes.gen3c_tokenizer_vae import (
                AutoencoderKLGen3CTokenizer)

            dtype = PRECISION_TO_TYPE[
                fastvideo_args.pipeline_config.vae_precision]
            num_frames = int(
                getattr(fastvideo_args.pipeline_config, "num_frames", 121))
            state_t = int(
                getattr(fastvideo_args.pipeline_config, "state_t", 16))
            if state_t > 1 and num_frames > 1:
                target_temporal = max(1,
                                      (num_frames - 1) // (state_t - 1))
            else:
                target_temporal = 8

            jit_dir = self._find_gen3c_jit_tokenizer_dir(model_path)
            if jit_dir is not None:
                vae = AutoencoderKLGen3CTokenizer.from_jit_tokenizer(
                    jit_dir,
                    device=target_device,
                    dtype=dtype,
                    target_temporal_compression=target_temporal,
                    pixel_chunk_duration=num_frames,
                )
                logger.info(
                    "Loaded GEN3C tokenizer VAE from JIT assets in %s (target temporal compression=%d)",
                    jit_dir,
                    target_temporal,
                )
                return vae.eval()

            tokenizer_ckpt = self._find_gen3c_tokenizer_checkpoint(
                model_path)
            if tokenizer_ckpt is not None:
                vae = AutoencoderKLGen3CTokenizer.from_tokenizer_checkpoint(
                    tokenizer_ckpt,
                    device=target_device,
                    dtype=dtype,
                    target_temporal_compression=target_temporal,
                    pixel_chunk_duration=num_frames,
                )
                logger.info(
                    "Loaded GEN3C tokenizer VAE from %s (target temporal compression=%d)",
                    tokenizer_ckpt,
                    target_temporal,
                )
                return vae.eval()
            logger.warning(
                "GEN3C tokenizer VAE checkpoint not found near %s; falling back to configured class %s.",
                model_path,
                class_name,
            )

        # Cosmos2.5 uses a Wan2.1 VAE stored as `tokenizer.safetensors` under the VAE folder.
        if class_name == "AutoencoderKLWan" and is_cosmos25:
            from fastvideo.models.vaes.cosmos25wanvae import Cosmos25WanVAE

            dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
            vae = Cosmos25WanVAE(device=target_device, dtype=dtype)

            weight_path = os.path.join(model_path, "tokenizer.safetensors")
            if not os.path.exists(weight_path):
                raise FileNotFoundError(
                    f"Missing Cosmos2.5 VAE weights: {weight_path}"
                )
            sd = safetensors_load_file(weight_path)
            vae.load_state_dict(sd, strict=False)
            return vae.eval()

        if class_name == "LingBotWorld2WanVAE":
            dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
            config.pop("_class_name", None)
            vae_config = fastvideo_args.pipeline_config.vae_config
            vae_config.update_model_arch(config)
            vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
            weight_path = os.path.join(model_path, "Wan2.1_VAE.pth")
            if not os.path.exists(weight_path):
                raise FileNotFoundError(
                    f"Missing LingBot World 2 VAE weights: {weight_path}"
                )
            vae = vae_cls(
                vae_config,
                checkpoint_path=weight_path,
                dtype=dtype,
            ).to(target_device)
            return vae.eval()

        # LTX-2 uses CausalVideoAutoencoder with nested "vae" config
        if class_name == "CausalVideoAutoencoder" and "vae" in config:
            vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
            vae = vae_cls(config).to(target_device)
            if hasattr(vae, "set_tiling_config"):
                vae_config = fastvideo_args.pipeline_config.vae_config
                vae.set_tiling_config(
                    spatial_tile_size_in_pixels=getattr(
                        vae_config, "ltx2_spatial_tile_size_in_pixels", 512),
                    spatial_tile_overlap_in_pixels=getattr(
                        vae_config, "ltx2_spatial_tile_overlap_in_pixels", 64),
                    temporal_tile_size_in_frames=getattr(
                        vae_config, "ltx2_temporal_tile_size_in_frames", 64),
                    temporal_tile_overlap_in_frames=getattr(
                        vae_config,
                        "ltx2_temporal_tile_overlap_in_frames", 24),
                )
        else:
            config.pop("_class_name", None)
            vae_config = fastvideo_args.pipeline_config.vae_config
            vae_config.update_model_arch(config)
            vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
            vae = vae_cls(vae_config).to(target_device)

    # 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}")
    # Common case: a single `.safetensors` checkpoint file.
    # Some models may be sharded into multiple files; in that case we merge.
    loaded = {}
    for sf_file in safetensors_list:
        loaded.update(safetensors_load_file(sf_file))

    # LTX-2 CausalVideoAutoencoder needs per_channel_statistics remapping
    if class_name == "CausalVideoAutoencoder" and "vae" in config:
        per_channel_prefixes = (
            "per_channel_statistics.",
            "vae.per_channel_statistics.",
        )
        remapped = {}
        for key, tensor in loaded.items():
            remapped[key] = tensor
            for prefix in per_channel_prefixes:
                if key.startswith(prefix):
                    suffix = key[len(prefix):]
                    remapped.setdefault(
                        f"encoder.per_channel_statistics.{suffix}",
                        tensor,
                    )
                    remapped.setdefault(
                        f"decoder.per_channel_statistics.{suffix}",
                        tensor,
                    )
                    break
        loaded = remapped

    # Diffusers-format AutoencoderKL checkpoints should match exactly; load
    # strictly so missing/unexpected keys are surfaced early.
    strict_load = class_name == "AutoencoderKL"
    vae.load_state_dict(loaded, strict=strict_load)
    if (class_name == "AutoencoderKLWan"
            and getattr(vae.config, "use_light_vae", False)
            and target_device.type == "cuda"
            and hasattr(vae, "optimize_memory_format")):
        vae.optimize_memory_format()

    return vae.eval()