1409 lines
66 KiB
Python
1409 lines
66 KiB
Python
# 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,
|
|
)
|
|
|
|
# Strict check for unquantized models and for serialized checkpoint
|
|
# schemes, which track every loaded tensor. It runs before the
|
|
# post-load hooks so a missing tensor is reported by name here
|
|
# instead of as one anonymous "not loaded" failure inside a hook.
|
|
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}")
|
|
|
|
if checkpoint_quant_config is not None:
|
|
processed_linears = _process_quantized_text_encoder_weights(model, runtime_device)
|
|
logger.info("Validated %d serialized %s text-encoder linears", processed_linears,
|
|
checkpoint_quant_config.get_name())
|
|
|
|
# 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,
|
|
)
|
|
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, # "<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
|
|
|
|
|
|
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 `<subfolder>/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)
|