Davids048andClaude Fable 5.1
a5aa64ab6b
VideoGenerator.from_pretrained(model_path, config) takes the same nested
settings as from_config; the 24 flat keywords, from_pretrained_kwargs_to_config,
FROM_PRETRAINED_KWARGS, and every flat_field in the schema are deleted.
nvfp4_fa4 is the typed field engine.attention.nvfp4_fa4, applied by a
resolution step. 65 call sites and the docs use the nested form; the three
kwargs golden cases are config cases with byte-identical results.
Every value is decided during resolution; no runtime code calls
with_override. The device offload policy (unified memory, MPS, layerwise
conflicts, lazy module load) runs as resolution steps in the main process
(fastvideo/api/device_policy.py), and the worker runs on the config it
receives. The LTX-2 refine defaults from model_index.json and the MiniMax-H3
schedule from fastvideo_inference.json are resolution steps
(fastvideo/api/checkpoint_defaults.py) that read local checkpoints or download
only those files; the H3 pipeline validates the schedule against the loaded
schedulers but no longer writes it. The preprocessing entry point passes the
downloaded local path as run state instead of a config override. Teacher and
critic models load with override_transformer_cls_name as a load argument.
The test isolation turns the device policy off, as it blocks downloads, so
that the goldens do not depend on the machine.
Against the 3f6893a0 goldens, every field still matches except boundary_ratio
and ltx2_vae_tiling (DESIGN section 9) and the fields that no longer exist.
The launcher and trainer YAML verifications report 0 unexplained differences.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
1437 lines
67 KiB
Python
1437 lines
67 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.api.device_policy import UNIFIED_MEMORY_OFFLOAD_PATHS
|
|
from fastvideo.api.inference_resolution import torch_compile_kwargs
|
|
from fastvideo.api.resolution import ResolvedGeneratorConfig
|
|
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.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__)
|
|
|
|
|
|
def _text_encoder_quant(resolved_config: ResolvedGeneratorConfig) -> str | None:
|
|
"""The online text-encoder quantization name, or ``None`` when the ``engine.quantization`` section is unset."""
|
|
quantization = resolved_config.engine.quantization
|
|
return None if quantization is None else quantization.text_encoder_quant
|
|
|
|
|
|
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, resolved_config: ResolvedGeneratorConfig):
|
|
"""
|
|
Load the component based on the model path, architecture, and inference args.
|
|
|
|
Args:
|
|
model_path: Path to the component model
|
|
resolved_config: The resolved runtime config
|
|
|
|
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, resolved_config: ResolvedGeneratorConfig):
|
|
"""Load the text encoders based on the model path, and inference args."""
|
|
# model_config: PretrainedConfig = get_hf_config(
|
|
# model=model_path,
|
|
# trust_remote_code=resolved_config.trust_remote_code,
|
|
# revision=resolved_config.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 = resolved_config.pipeline_config.text_encoder_configs
|
|
encoder_precisions = resolved_config.engine.precision.text_encoders
|
|
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,
|
|
resolved_config,
|
|
encoder_precision,
|
|
use_text_encoder_override=True,
|
|
)
|
|
|
|
def load_model(
|
|
self,
|
|
model_path: str,
|
|
model_config: EncoderConfig,
|
|
target_device: torch.device,
|
|
resolved_config: ResolvedGeneratorConfig,
|
|
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()
|
|
if cpu_offload is None:
|
|
# offload_flag names the device-policy switch of this component; its value is the engine.offload field
|
|
# at the same path (text_encoder_cpu_offload -> engine.offload.text_encoder), which resolution already
|
|
# settled for the memory class of this device.
|
|
offload_field = UNIFIED_MEMORY_OFFLOAD_PATHS[offload_flag].rsplit(".", 1)[1]
|
|
cpu_offload = getattr(resolved_config.engine.offload, offload_field)
|
|
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,
|
|
resolved_config,
|
|
use_text_encoder_override,
|
|
)
|
|
checkpoint_quant_config = _configure_text_encoder_quantization(
|
|
model_config,
|
|
model_cls,
|
|
checkpoint_path,
|
|
)
|
|
text_encoder_quant = _text_encoder_quant(resolved_config)
|
|
text_encoder_weights = resolved_config.pipeline.components.text_encoder_weights
|
|
if checkpoint_quant_config is not None:
|
|
if text_encoder_quant is not None:
|
|
raise ValueError("Serialized checkpoint quantization is selected from checkpoint metadata; "
|
|
"engine.quantization.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 text_encoder_quant is not None:
|
|
if text_encoder_weights is None:
|
|
raise ValueError("engine.quantization.text_encoder_quant is set but "
|
|
"pipeline.components.text_encoder_weights is None")
|
|
quant_cls = get_quantization_config(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 text_encoder_weights 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 = resolved_config.engine.offload.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, resolved_config: ResolvedGeneratorConfig):
|
|
"""Load the text encoders based on the model path, and inference args."""
|
|
# model_config: PretrainedConfig = get_hf_config(
|
|
# model=model_path,
|
|
# trust_remote_code=resolved_config.trust_remote_code,
|
|
# revision=resolved_config.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(
|
|
resolved_config.pipeline_config, "image_encoder_configs", None)
|
|
encoder_precisions = getattr(
|
|
resolved_config.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 = resolved_config.pipeline_config.image_encoder_config
|
|
encoder_precision = resolved_config.engine.precision.image_encoder
|
|
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 resolved_config.engine.offload.image_encoder:
|
|
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,
|
|
resolved_config,
|
|
encoder_precision,
|
|
cpu_offload=resolved_config.engine.offload.image_encoder,
|
|
offload_flag="image_encoder_cpu_offload",
|
|
)
|
|
|
|
|
|
class VisionLanguageEncoderLoader(ComponentLoader):
|
|
"""Loader for vision-language autoregressive encoders."""
|
|
|
|
def load(self, model_path: str, resolved_config: ResolvedGeneratorConfig):
|
|
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=resolved_config.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, resolved_config: ResolvedGeneratorConfig):
|
|
from transformers import AutoProcessor
|
|
|
|
logger.info("Loading processor from %s", model_path)
|
|
processor = AutoProcessor.from_pretrained(
|
|
model_path,
|
|
trust_remote_code=resolved_config.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, resolved_config: ResolvedGeneratorConfig):
|
|
"""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, resolved_config: ResolvedGeneratorConfig):
|
|
"""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(resolved_config.pipeline_config, "text_encoder_configs"):
|
|
try:
|
|
require_processor = any(
|
|
getattr(getattr(cfg, "arch_config", None), "require_processor", False)
|
|
for cfg in resolved_config.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=resolved_config.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=resolved_config.trust_remote_code,
|
|
)
|
|
logger.info(
|
|
"Loaded tokenizer/processor from %s: %s",
|
|
resolved_model_path,
|
|
tokenizer.__class__.__name__,
|
|
)
|
|
return tokenizer
|
|
padding_side = None
|
|
if hasattr(resolved_config.pipeline_config, "text_encoder_configs"):
|
|
try:
|
|
arch_config = resolved_config.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, resolved_config: ResolvedGeneratorConfig):
|
|
"""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.")
|
|
|
|
from fastvideo.platforms import current_platform
|
|
|
|
if resolved_config.engine.offload.vae:
|
|
target_device = (torch.device("mps") if current_platform.is_mps() else torch.device("cpu"))
|
|
else:
|
|
target_device = get_local_torch_device()
|
|
|
|
vae_precision = resolved_config.engine.precision.vae
|
|
with set_default_torch_dtype(PRECISION_TO_TYPE[vae_precision] if vae_precision else torch.bfloat16):
|
|
pipeline_name = resolved_config.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[vae_precision]
|
|
num_frames = int(getattr(resolved_config.pipeline_config, "num_frames", 121))
|
|
state_t = int(getattr(resolved_config.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[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[vae_precision]
|
|
config.pop("_class_name", None)
|
|
vae_config = resolved_config.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 = resolved_config.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 = resolved_config.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, resolved_config: ResolvedGeneratorConfig):
|
|
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(resolved_config.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 resolved_config.engine.offload.vae:
|
|
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(resolved_config.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, resolved_config: ResolvedGeneratorConfig):
|
|
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(resolved_config.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,
|
|
resolved_config: ResolvedGeneratorConfig,
|
|
loading_teacher_critic_model: bool = False,
|
|
override_transformer_cls_name: str | None = None):
|
|
"""Load the transformer based on the model path, and inference args.
|
|
|
|
``loading_teacher_critic_model`` marks a DMD teacher (real_score) or critic (fake_score) transformer, which
|
|
loads its checkpoint weights in full precision with dense attention. ``override_transformer_cls_name`` is the
|
|
model class to build for this load; it takes precedence over
|
|
``pipeline.components.override_transformer_cls_name``.
|
|
"""
|
|
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 override_transformer_cls_name is None:
|
|
override_transformer_cls_name = resolved_config.pipeline.components.override_transformer_cls_name
|
|
if override_transformer_cls_name is not None:
|
|
cls_name = override_transformer_cls_name
|
|
logger.info("Overriding transformer cls_name to %s", cls_name)
|
|
|
|
# Config from Diffusers supersedes fastvideo's model config
|
|
dit_config = deepcopy(resolved_config.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 loading_teacher_critic_model 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 = 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
|
|
components = resolved_config.pipeline.components
|
|
custom_weights_path = components.transformer_weights
|
|
use_custom_weights = (custom_weights_path and os.path.exists(custom_weights_path)
|
|
and not loading_teacher_critic_model)
|
|
|
|
if use_custom_weights:
|
|
if "transformer_2" in model_path:
|
|
custom_weights_path = components.transformer_2_weights
|
|
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[resolved_config.engine.precision.dit]
|
|
|
|
# Load the model using FSDP loader
|
|
logger.info("Loading model from %s, default_dtype: %s", cls_name, default_dtype)
|
|
parallelism = resolved_config.engine.parallelism
|
|
assert parallelism.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(resolved_config.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=parallelism.hsdp_replicate_dim,
|
|
hsdp_shard_dim=parallelism.hsdp_shard_dim,
|
|
strict=strict_load,
|
|
cpu_offload=resolved_config.engine.offload.dit,
|
|
pin_cpu_memory=resolved_config.engine.offload.pin_cpu_memory,
|
|
fsdp_inference=resolved_config.engine.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=resolved_config.training_mode,
|
|
enable_torch_compile=resolved_config.engine.compile.enabled,
|
|
torch_compile_kwargs=torch_compile_kwargs(resolved_config),
|
|
inference_regional_compile=resolved_config.engine.compile.regional,
|
|
inference_vsa_tile_size=resolved_config.engine.attention.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=components.lora_path,
|
|
lora_strength=components.lora_strength,
|
|
)
|
|
|
|
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 resolved_config.inference_mode and resolved_config.engine.offload.dit_layerwise:
|
|
# 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, resolved_config: ResolvedGeneratorConfig):
|
|
"""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)
|
|
flow_shift = resolved_config.pipeline.flow_shift
|
|
if flow_shift is not None:
|
|
scheduler.set_shift(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, resolved_config: ResolvedGeneratorConfig):
|
|
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(resolved_config.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, resolved_config: ResolvedGeneratorConfig):
|
|
"""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(resolved_config.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(resolved_config.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, resolved_config: ResolvedGeneratorConfig):
|
|
"""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=resolved_config.trust_remote_code,
|
|
revision=resolved_config.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,
|
|
resolved_config: ResolvedGeneratorConfig,
|
|
loading_teacher_critic_model: bool = False,
|
|
override_transformer_cls_name: str | None = None,
|
|
):
|
|
"""
|
|
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
|
|
resolved_config: The resolved runtime config
|
|
loading_teacher_critic_model: Whether a transformer is a DMD teacher or critic; passed to
|
|
``TransformerLoader.load`` and ignored by the other loaders
|
|
override_transformer_cls_name: The model class to build for a transformer, over
|
|
``pipeline.components.override_transformer_cls_name``; passed to ``TransformerLoader.load`` and
|
|
ignored by the other loaders
|
|
|
|
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
|
|
# engine.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(resolved_config.engine.attention.backend)
|
|
attention_context = (_component_attention_backend_scope(resolved, component=module_name)
|
|
if resolved is not None else nullcontext())
|
|
|
|
with attention_context:
|
|
if isinstance(loader, TransformerLoader):
|
|
return loader.load(component_model_path,
|
|
resolved_config,
|
|
loading_teacher_critic_model=loading_teacher_critic_model,
|
|
override_transformer_cls_name=override_transformer_cls_name)
|
|
return loader.load(component_model_path, resolved_config)
|