# SPDX-License-Identifier: Apache-2.0 import dataclasses import glob import json import os import time from abc import ABC, abstractmethod from collections.abc import Generator, Iterable from contextlib import nullcontext from copy import deepcopy from typing import Any, cast import torch import torch.distributed as dist import torch.nn as nn from safetensors.torch import load_file as safetensors_load_file, safe_open from torch.distributed import init_device_mesh from transformers import AutoImageProcessor, AutoProcessor, AutoTokenizer from transformers.utils import SAFE_WEIGHTS_INDEX_NAME from fastvideo.attention.selector import ( _active_component_attention_backend_scope, _component_attention_backend_scope, coerce_attn_backend, record_resolved_attention_backend, ) from fastvideo.configs.models import EncoderConfig from fastvideo.distributed import get_local_torch_device from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.layers.quantization import get_quantization_config from fastvideo.logger import init_logger from fastvideo.models.hf_transformer_utils import get_diffusers_config from fastvideo.models.loader.fsdp_load import maybe_load_fsdp_model, shard_model from fastvideo.models.loader.text_encoder_quantization import ( _configure_text_encoder_quantization, _process_quantized_text_encoder_weights, _resolve_text_encoder_checkpoint_path, ) from fastvideo.models.loader.utils import set_default_torch_dtype from fastvideo.models.loader.weight_utils import ( filter_duplicate_safetensors_files, filter_files_not_needed_for_inference, pt_weights_iterator, safetensors_weights_iterator, ) from fastvideo.models.registry import ModelRegistry from fastvideo.utils import PRECISION_TO_TYPE, is_pin_memory_available from fastvideo.hooks.layerwise_offload import enable_layerwise_offload logger = init_logger(__name__) class ComponentLoader(ABC): """Base class for loading a specific type of model component.""" def __init__(self, device=None) -> None: self.device = device @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 @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"), "audio_scheduler": (SchedulerLoader, "diffusers"), "transformer": (TransformerLoader, "diffusers"), "transformer_ref": (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"), "image_encoder_2": (ImageEncoderLoader, "transformers"), "image_encoder_3": (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) class TextEncoderLoader(ComponentLoader): """Loader for text encoders.""" @dataclasses.dataclass class Source: """A source for weights.""" model_or_path: str """The model ID or path.""" prefix: str = "" """A prefix to prepend to all weights.""" fall_back_to_pt: bool = True """Whether .pt weights can be used.""" allow_patterns_overrides: list[str] | None = None """If defined, weights will load exclusively using these patterns.""" counter_before_loading_weights: float = 0.0 counter_after_loading_weights: float = 0.0 def _prepare_weights( self, model_name_or_path: str, fall_back_to_pt: bool, allow_patterns_overrides: list[str] | None, ) -> tuple[str, list[str], bool]: """Prepare weights for the model. If the model is not local, it will be downloaded.""" # model_name_or_path = (self._maybe_download_from_modelscope( # model_name_or_path, revision) or model_name_or_path) is_local = os.path.isdir(model_name_or_path) assert is_local, "Model path must be a local directory" use_safetensors = False index_file = SAFE_WEIGHTS_INDEX_NAME allow_patterns = ["*.safetensors", "*.bin"] if fall_back_to_pt: allow_patterns += ["*.pt"] if allow_patterns_overrides is not None: allow_patterns = allow_patterns_overrides hf_folder = model_name_or_path hf_weights_files: list[str] = [] for pattern in allow_patterns: hf_weights_files += glob.glob(os.path.join(hf_folder, pattern)) if len(hf_weights_files) > 0: if pattern == "*.safetensors": use_safetensors = True break if use_safetensors: hf_weights_files = filter_duplicate_safetensors_files(hf_weights_files, hf_folder, index_file) else: hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files) if len(hf_weights_files) == 0: raise RuntimeError(f"Cannot find any model weights with `{model_name_or_path}`") return hf_folder, hf_weights_files, use_safetensors def _get_weights_iterator(self, source: "Source", to_cpu: bool) -> Generator[tuple[str, torch.Tensor], None, None]: """Get an iterator for the model weights based on the load format.""" hf_folder, hf_weights_files, use_safetensors = self._prepare_weights( source.model_or_path, source.fall_back_to_pt, source.allow_patterns_overrides, ) if use_safetensors: weights_iterator = safetensors_weights_iterator(hf_weights_files, to_cpu=to_cpu) else: weights_iterator = pt_weights_iterator(hf_weights_files, to_cpu=to_cpu) if self.counter_before_loading_weights == 0.0: self.counter_before_loading_weights = time.perf_counter() # Apply the prefix. return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator) def _get_all_weights( self, model: nn.Module, model_path: str, to_cpu: bool, ) -> Generator[tuple[str, torch.Tensor], None, None]: primary_weights = TextEncoderLoader.Source( model_path, prefix="", fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True), allow_patterns_overrides=getattr(model, "allow_patterns_overrides", None), ) yield from self._get_weights_iterator(primary_weights, to_cpu) secondary_weights = cast( Iterable[TextEncoderLoader.Source], getattr(model, "secondary_weights", ()), ) for source in secondary_weights: yield from self._get_weights_iterator(source, to_cpu) 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) record_resolved_attention_backend(encoder_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_model( self, model_path: str, model_config: EncoderConfig, target_device: torch.device, fastvideo_args: FastVideoArgs, dtype: str = "fp16", use_text_encoder_override: bool = False, # prevent subclasses from misusing cpu_offload: bool | None = None, offload_flag: str = "text_encoder_cpu_offload", ): runtime_device = get_local_torch_device() device_id = runtime_device.index if runtime_device.index is not None else 0 requested_cpu_offload = getattr(fastvideo_args, offload_flag) if cpu_offload is None else cpu_offload disable_cpu_offload = fastvideo_args.disable_offload_on_unified_memory(device_id, offload_flag=offload_flag) if requested_cpu_offload and disable_cpu_offload: # Direct loader callers can choose a CPU target before the worker # applies its device-local policy. Reset both the request and the # target so the model is never constructed on the host first. logger.info("Disabling %s on unified-memory device %d", offload_flag, device_id) cpu_offload = False target_device = runtime_device else: cpu_offload = requested_cpu_offload use_cpu_offload = (cpu_offload and len(getattr(model_config, "_fsdp_shard_conditions", [])) > 0) from fastvideo.platforms import current_platform if cpu_offload: target_device = (torch.device("mps") if current_platform.is_mps() else torch.device("cpu")) with set_default_torch_dtype(PRECISION_TO_TYPE[dtype]): architectures = getattr(model_config, "architectures", []) model_cls, _ = ModelRegistry.resolve_model_cls(architectures) checkpoint_path = _resolve_text_encoder_checkpoint_path( model_path, fastvideo_args, use_text_encoder_override, ) checkpoint_quant_config = _configure_text_encoder_quantization( model_config, model_cls, checkpoint_path, ) if checkpoint_quant_config is not None: if fastvideo_args.override_text_encoder_quant is not None: raise ValueError("Serialized checkpoint quantization is selected from checkpoint metadata; " "override_text_encoder_quant is an online conversion option and must be unset") requested_dtype = PRECISION_TO_TYPE[dtype] if requested_dtype not in checkpoint_quant_config.get_supported_act_dtypes(): raise ValueError(f"Serialized {checkpoint_quant_config.get_name()} text encoder does not support " f"activation dtype {requested_dtype}") checkpoint_quant_config.validate_runtime(runtime_device) logger.info( "Selected serialized %s text-encoder checkpoint execution from %s", checkpoint_quant_config.get_name(), checkpoint_path, ) elif use_text_encoder_override and fastvideo_args.override_text_encoder_quant is not None: if fastvideo_args.override_text_encoder_safetensors is None: raise ValueError("override_text_encoder_quant is set but override_text_encoder_safetensors is None") quant_cls = get_quantization_config(fastvideo_args.override_text_encoder_quant) model_config.quant_config = quant_cls() if getattr(model_cls, "supports_hf_from_pretrained", False): model = model_cls.from_pretrained_local( # type: ignore[attr-defined] model_path, model_config, # type: ignore[arg-type] dtype=PRECISION_TO_TYPE[dtype], device=target_device, ) # HF passthrough encoders return before FastVideo's FSDP # wrapping path, so the text stage needs their placement to # put token tensors on the same device. model._fastvideo_input_device = target_device return model.eval() with target_device: model = model_cls(model_config) # type: ignore weights_to_load = {name for name, _ in model.named_parameters()} if (use_text_encoder_override and fastvideo_args.override_text_encoder_safetensors is not None): if os.path.isdir(checkpoint_path): override_weights = self._get_all_weights( model, checkpoint_path, to_cpu=bool(cpu_offload), ) else: if self.counter_before_loading_weights == 0.0: self.counter_before_loading_weights = time.perf_counter() override_weights = safetensors_weights_iterator( [checkpoint_path], to_cpu=use_cpu_offload, ) loaded_weights: set[str] = model.load_weights(override_weights) # type: ignore else: loaded_weights: set[str] = model.load_weights( self._get_all_weights( model, model_path, to_cpu=cpu_offload, )) # type: ignore self.counter_after_loading_weights = time.perf_counter() logger.info( "Loading weights took %.2f seconds", self.counter_after_loading_weights - self.counter_before_loading_weights, ) if checkpoint_quant_config is not None: processed_linears = _process_quantized_text_encoder_weights(model, runtime_device) logger.info("Validated %d serialized blockwise FP8 text-encoder linears", processed_linears) # Explicitly move model to target device after loading weights model = model.to(target_device) from fastvideo.platforms import current_platform if use_cpu_offload: pin_cpu_memory = fastvideo_args.pin_cpu_memory and is_pin_memory_available() # Disable FSDP for MPS as it's not compatible if current_platform.is_mps(): logger.info("Disabling FSDP sharding for MPS platform as it's not compatible") elif current_platform.is_npu(): mesh = init_device_mesh( "npu", mesh_shape=(1, dist.get_world_size()), mesh_dim_names=("offload", "replicate"), ) shard_model( model, cpu_offload=True, reshard_after_forward=True, mesh=mesh["offload"], fsdp_shard_conditions=model._fsdp_shard_conditions, pin_cpu_memory=pin_cpu_memory, ) else: mesh = init_device_mesh( "cuda", mesh_shape=(1, dist.get_world_size()), mesh_dim_names=("offload", "replicate"), ) shard_model( model, cpu_offload=True, reshard_after_forward=True, mesh=mesh["offload"], fsdp_shard_conditions=model._fsdp_shard_conditions, pin_cpu_memory=pin_cpu_memory, ) # We only enable strict check for non-quantized models # that have loaded weights tracking currently. # if loaded_weights is not None: weights_not_loaded = weights_to_load - loaded_weights if weights_not_loaded and (model_config.quant_config is None or checkpoint_quant_config is not None): raise ValueError("Following weights were not initialized from " f"checkpoint: {weights_not_loaded}") return model.eval() class ImageEncoderLoader(TextEncoderLoader): 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) base = os.path.basename(os.path.normpath(model_path)) index = 0 if base.startswith("image_encoder_"): try: index = int(base.split("_")[-1]) - 1 except ValueError: index = 0 encoder_configs = getattr( fastvideo_args.pipeline_config, "image_encoder_configs", None) encoder_precisions = getattr( fastvideo_args.pipeline_config, "image_encoder_precisions", None) if encoder_configs is None: if index != 0: raise IndexError( f"image encoder index {index} requires pipeline_config." "image_encoder_configs") encoder_config = fastvideo_args.pipeline_config.image_encoder_config encoder_precision = fastvideo_args.pipeline_config.image_encoder_precision else: if index < 0 or index >= len(encoder_configs): raise IndexError( f"image encoder index {index} out of range for " f"image_encoder_configs (len={len(encoder_configs)})") encoder_config = encoder_configs[index] if encoder_precisions is None or index >= len( encoder_precisions): raise IndexError( f"image encoder index {index} has no matching precision") encoder_precision = encoder_precisions[index] encoder_config.update_model_arch(model_config) record_resolved_attention_backend(encoder_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, encoder_precision, cpu_offload=fastvideo_args.image_encoder_cpu_offload, offload_flag="image_encoder_cpu_offload", ) class VisionLanguageEncoderLoader(ComponentLoader): """Loader for vision-language autoregressive encoders.""" def load(self, model_path: str, fastvideo_args: FastVideoArgs): from fastvideo.distributed.parallel_state import get_local_torch_device from fastvideo.models.encoders.glm_image_ar_loader import (GlmImageARLoader) logger.info("Loading vision-language encoder from %s", model_path) target_device = get_local_torch_device() loader = GlmImageARLoader( model_path, torch_dtype=torch.bfloat16, trust_remote_code=fastvideo_args.trust_remote_code, ).to(target_device).eval() return loader class ProcessorLoader(ComponentLoader): """Loader for HF processors that pair with vision-language encoders.""" def load(self, model_path: str, fastvideo_args: FastVideoArgs): from transformers import AutoProcessor logger.info("Loading processor from %s", model_path) processor = AutoProcessor.from_pretrained( model_path, trust_remote_code=fastvideo_args.trust_remote_code, ) logger.info("Loaded processor: %s", processor.__class__.__name__) return processor class ImageProcessorLoader(ComponentLoader): """Loader for image processor.""" 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 class TokenizerLoader(ComponentLoader): """Loader for tokenizers.""" 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, # "/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 class VAELoader(ComponentLoader): """Loader for VAE.""" @staticmethod def _find_gen3c_tokenizer_checkpoint(model_path: str) -> str | None: """Locate tokenizer-backed VAE checkpoint used by GEN3C integration.""" candidates = [ os.path.join(model_path, "tokenizer.pth"), os.path.join(os.path.dirname(model_path), "tokenizer", "tokenizer.pth"), ] for candidate in candidates: if os.path.exists(candidate): return candidate return None @staticmethod def _find_gen3c_jit_tokenizer_dir(model_path: str) -> str | None: """Locate official tokenizer JIT assets (encoder/decoder/mean_std).""" candidates = [ model_path, os.path.join(os.path.dirname(model_path), "tokenizer"), ] required = ("encoder.jit", "decoder.jit", "mean_std.pt") for directory in candidates: if all(os.path.exists(os.path.join(directory, name)) for name in required): return directory return None 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 in {"AutoencoderKL", "AutoencoderKLMiniMaxH3"} 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() class AudioDecoderLoader(ComponentLoader): """Loader for full audio VAEs and legacy decode-only audio components.""" def load(self, model_path: str, fastvideo_args: FastVideoArgs): config = get_diffusers_config(model=model_path) class_name = config.pop("_class_name", None) or "LTX2AudioDecoder" model_cls, _ = ModelRegistry.resolve_model_cls(class_name) target_device = get_local_torch_device() if class_name == "AutoencoderKLMiniMaxH3Audio": from fastvideo.platforms import current_platform configured_audio_vae = getattr(fastvideo_args.pipeline_config, "audio_vae_config", None) if configured_audio_vae is None: raise ValueError("MiniMax H3 requires audio_vae_config.") config.pop("_name_or_path", None) audio_vae_config = deepcopy(configured_audio_vae) audio_vae_config.update_model_arch(config) if getattr(fastvideo_args, "vae_cpu_offload", False): target_device = torch.device("mps") if current_platform.is_mps() else torch.device("cpu") with set_default_torch_dtype(torch.float32): audio_vae = model_cls(audio_vae_config).to(device=target_device, dtype=torch.float32) 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}") loaded: dict[str, torch.Tensor] = {} for sf_file in safetensors_list: loaded.update(safetensors_load_file(sf_file)) audio_vae.load_state_dict(loaded, strict=True) return audio_vae.eval() precision = getattr(fastvideo_args.pipeline_config, "audio_decoder_precision", "bf16") # MMAudio normalizes its magnitude-preserving convolution weights in # fp32 and only then casts the whole feature utility module to bf16. # Constructing/loading directly in bf16 quantizes the unnormalized # checkpoint first and changes the decoded mel trajectory. construction_precision = "fp32" if class_name == "MMAudioVAE" else precision construction_device = torch.device("cpu") if class_name == "MMAudioVAE" else target_device with set_default_torch_dtype(PRECISION_TO_TYPE[construction_precision]): audio_decoder = model_cls(config).to(construction_device) safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors")) loaded: dict[str, torch.Tensor] = {} for sf_file in safetensors_list: loaded.update(safetensors_load_file(sf_file)) if class_name == "MMAudioVAE": audio_decoder.load_state_dict(loaded, strict=True) audio_decoder.remove_weight_norm() return audio_decoder.to(device=target_device, dtype=PRECISION_TO_TYPE[precision]).eval() decoder_state = {} for name, tensor in loaded.items(): if name.startswith("decoder."): decoder_state[name.replace("decoder.", "")] = tensor elif name.startswith("per_channel_statistics."): decoder_state[name] = tensor target_module = getattr(audio_decoder, "model", audio_decoder) target_module.load_state_dict(decoder_state, strict=False) return audio_decoder.eval() class VocoderLoader(ComponentLoader): """Loader for native vocoders.""" def load(self, model_path: str, fastvideo_args: FastVideoArgs): config = get_diffusers_config(model=model_path) class_name = config.pop("_class_name", None) or "LTX2Vocoder" model_cls, _ = ModelRegistry.resolve_model_cls(class_name) target_device = get_local_torch_device() precision = getattr(fastvideo_args.pipeline_config, "vocoder_precision", "bf16") # Canonical BigVGAN likewise removes parametrized weight norm in fp32 # before the official MMAudio feature module is cast to bf16. construction_precision = "fp32" if class_name == "BigVGANV2" else precision construction_device = torch.device("cpu") if class_name == "BigVGANV2" else target_device with set_default_torch_dtype(PRECISION_TO_TYPE[construction_precision]): vocoder = model_cls(config).to(construction_device) safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors")) loaded: dict[str, torch.Tensor] = {} for sf_file in safetensors_list: loaded.update(safetensors_load_file(sf_file)) if class_name == "BigVGANV2": vocoder.load_state_dict(loaded, strict=True) vocoder.remove_weight_norm() return vocoder.to(device=target_device, dtype=PRECISION_TO_TYPE[precision]).eval() target_module = getattr(vocoder, "model", vocoder) target_module.load_state_dict(loaded, strict=False) return vocoder.eval() def _collect_safetensors_keys(safetensors_list: list) -> set: """Collect all weight keys from safetensors files.""" all_keys: set[str] = set() for path in safetensors_list: try: with safe_open(path, framework="pt") as f: all_keys.update(f.keys()) except Exception as e: logger.warning("Could not read keys from %s: %s", path, e) return all_keys class TransformerLoader(ComponentLoader): """Loader for transformer.""" 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 # build their attention under a scope that ignores any process-wide # ATTN_QAT_TRAIN request so it falls back to dense. The generator loads # without the flag and keeps both. The scope is exception-safe and # needs no env mutation or selector cache flush (the request is part # of the resolution cache key). _qat_generator_only = hasattr(fastvideo_args, "_loading_teacher_critic_model") if _qat_generator_only: dit_config.quant_config = None 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") attention_context = (_component_attention_backend_scope(None, component="transformer") if _qat_generator_only else nullcontext()) with attention_context: # dit_config is what the model is handed and keeps as `self.config`, # so recording here makes the decision readable from the loaded # transformer — and records the narrowed one for teacher/critic. resolved = record_resolved_attention_backend(dit_config) # Every worker records its resolved backend so distributed profile # snapshots can prove that all ranks use the requested kernels. logger.info("Worker %s transformer attention backend: %s", os.environ.get("RANK", "0"), resolved.name if resolved else "automatic selection", local_main_process_only=False) 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, inference_regional_compile=fastvideo_args.inference_torch_compile, inference_vsa_tile_size=fastvideo_args.VSA_tile_size, # Only the whole-parameter half of the adapter is applied here, while # tensors are still unsharded; LoRAPipeline merges the low-rank half # once the module tree exists. lora_path=getattr(fastvideo_args, "lora_path", None), lora_strength=getattr(fastvideo_args, "lora_strength", 1.0), ) total_params = sum(p.numel() for p in model.parameters()) logger.info("Loaded model with %.2fB parameters", total_params / 1e9) first_name, first_parameter = next(model.named_parameters()) dtype_selector = getattr(model, "_get_parameter_dtype", None) expected_dtype = dtype_selector(first_name, default_dtype) if callable(dtype_selector) else default_dtype assert first_parameter.dtype == expected_dtype, "Model dtype does not match the configured parameter 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 class SchedulerLoader(ComponentLoader): """Loader for scheduler.""" 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 class ConditionerLoader(ComponentLoader): """Loader for multi-conditioner components (e.g. Stable Audio's `StableAudioMultiConditioner`, which bundles T5 + NumberConditioners and is neither a pure text encoder nor a Diffusers-shaped module). Reads `/config.json` to resolve the class via `ModelRegistry`, instantiates with no args (the class pulls its own defaults from its FastVideo config), then loads `diffusion_pytorch_model.safetensors` non-strictly so externally fetched sub-encoders (T5) don't trip the missing-key check. """ def load(self, model_path: str, fastvideo_args: FastVideoArgs): config = get_diffusers_config(model=model_path) class_name = config.pop("_class_name", None) config.pop("_name_or_path", None) if class_name is None: raise ValueError(f"Conditioner config at {model_path} is missing the " f"`_class_name` attribute required to resolve a model class.") model_cls, _ = ModelRegistry.resolve_model_cls(class_name) target_device = get_local_torch_device() precision = getattr(fastvideo_args.pipeline_config, "precision", "fp16") target_dtype = PRECISION_TO_TYPE.get(precision, torch.float16) # Without this merge the model falls back to its dataclass # defaults (e.g. SA-1.0's 3-conditioner spec — wrong for SA-small). from dataclasses import fields as _fields from fastvideo.configs.models.encoders import ( StableAudioConditionerConfig, ) if model_cls.__name__ == "StableAudioMultiConditioner": cond_config = StableAudioConditionerConfig() # `update_model_arch` is strict (raises on unknown keys); the # converter writes a few non-arch keys (`_class_name`, # `_diffusers_version`, `_name_or_path`) that must be filtered # out first. valid = {f.name for f in _fields(cond_config.arch_config)} cond_config.update_model_arch({k: v for k, v in config.items() if k in valid}) with set_default_torch_dtype(target_dtype): model = model_cls(cond_config) else: with set_default_torch_dtype(target_dtype): model = model_cls() weights = os.path.join(str(model_path), "diffusion_pytorch_model.safetensors") if not os.path.isfile(weights): raise FileNotFoundError(f"Conditioner weights not found: {weights}") state = safetensors_load_file(weights) # Non-strict: T5 weights live outside this checkpoint (fetched in # the conditioner's `__init__` from the standard HF repo). model.load_state_dict(state, strict=False) return model.to(device=target_device, dtype=target_dtype).eval() class UpsamplerLoader(ComponentLoader): """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.""" 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() class GenericComponentLoader(ComponentLoader): """Generic loader for components that don't have a specific loader.""" def __init__(self, library="transformers") -> None: super().__init__() self.library = library 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}") class PipelineComponentLoader: """ Utility class for loading pipeline components. This replaces the chain of if-else statements in load_pipeline_module. """ @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) # Resolve this component's attention backend ONCE, here, from the # explicit request and the environment. A per-role request already # resolved upstream (the train stack's moduleloader) wins; otherwise # fastvideo_args.attention_backend is the process-wide default, applied # per component. # # The loaders record the decision on their own config rather than this # function doing it after the fact, because a loader may narrow it for # one component: the DMD teacher/critic transformers build dense inside # a nested scope (see `record_resolved_attention_backend`). if _active_component_attention_backend_scope() is not None: attention_context = nullcontext() else: resolved = coerce_attn_backend(getattr(fastvideo_args, "attention_backend", None)) attention_context = (_component_attention_backend_scope(resolved, component=module_name) if resolved is not None else nullcontext()) with attention_context: return loader.load(component_model_path, fastvideo_args)