SolitaryThinkerandClaude Opus 5.5
d692d1cc3b
These ports were marked pending removal in #1820. Each has a newer, maintained counterpart in the repo, and their model-specific hooks leaked into shared stages, request fields, CLI args, and loaders. Remove each model's code, configs, presets, examples, converters, tests, cookbook pages, CI planner rules, and registry entries, along with the shared-code hooks only they used: - camera kwargs (GameCraft camera_states, DreamX y_camera) in the shared denoising loop - LongCat refine branches in latent preparation, the decoding crop, the --bsa-* CLI args, its train model, and third_party/longcat_video - Z-Image cfg_normalization/cfg_truncation and GameCraft/DreamX camera/action fields in SamplingParam, ForwardBatch, and the schema - the GLM-Image vision-language loader and the TurboDiffusion rCM scheduler and SLA auto-select in the train entrypoint - the Wan definition groups that only preserved DreamX detector order Removed env vars and their legacy aliases move to DEPRECATED_VARIABLES per the env-var policy. Matrix Game 2.0 stays: the new trainer's world model scenarios depend on it. The SLA and BSA attention backends stay. The TurboDiffusion int8 kernels in fastvideo-kernel are left for a separate kernel PR. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
1539 lines
74 KiB
Python
1539 lines
74 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 fastvideo.envs as envs
|
|
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, UpsamplerConfig
|
|
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"),
|
|
"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"),
|
|
# Kandinsky6 SR latent-upscaler bank (x2 / x4 entries).
|
|
# (the official Diffusers save_pretrained writes the library as "kandinsky6", the Hub repo as "diffusers")
|
|
"latent_upscaler": (UpsamplerLoader, ("diffusers", "kandinsky6")),
|
|
}
|
|
|
|
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
|
|
allowed = (expected_library, ) if isinstance(expected_library, str) else tuple(expected_library)
|
|
assert transformers_or_diffusers in allowed, 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)
|
|
|
|
prepare_layerwise = getattr(model, "prepare_layerwise_offload", None)
|
|
if envs.FASTVIDEO_H3_ENCODER_LAYERWISE.get() and callable(prepare_layerwise):
|
|
if target_device.type != "cpu":
|
|
raise ValueError("Layerwise H3 encoder requires text_encoder_cpu_offload=True")
|
|
prepare_layerwise(runtime_device)
|
|
use_cpu_offload = False
|
|
logger.info("Enabled text-only layerwise H3 encoder with CPU token embeddings")
|
|
|
|
from fastvideo.platforms import current_platform
|
|
|
|
if use_cpu_offload and checkpoint_quant_config is not None:
|
|
logger.info(
|
|
"Skipping FSDP CPU offload for serialized %s text encoder; "
|
|
"packed uint8 weights are not FSDP-shardable",
|
|
checkpoint_quant_config.get_name(),
|
|
)
|
|
elif 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 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", None)
|
|
config.pop("_name_or_path", None)
|
|
if class_name is None:
|
|
class_name = getattr(fastvideo_args, "_model_index_class_names", {}).get("vae")
|
|
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"))
|
|
int8_convrot_path = None
|
|
if class_name == "AutoencoderKLMiniMaxH3":
|
|
from fastvideo.models.vaes.minimax_h3_int8_convrot import (
|
|
dense_vae_safetensors,
|
|
find_int8_convrot_vae_path,
|
|
)
|
|
int8_convrot_path = find_int8_convrot_vae_path(model_path)
|
|
safetensors_list = dense_vae_safetensors(safetensors_list)
|
|
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.
|
|
# Kandinsky6SRVAE: a partially loaded KVAE decodes plausible-looking garbage, so it is strict as well.
|
|
strict_load = class_name in {"AutoencoderKL", "AutoencoderKLMiniMaxH3", "Kandinsky6SRVAE"}
|
|
vae.load_state_dict(loaded, strict=strict_load)
|
|
if class_name == "AutoencoderKLMiniMaxH3" and int8_convrot_path is not None:
|
|
from fastvideo.models.vaes.minimax_h3_int8_convrot import overlay_minimax_h3_int8_convrot_decoder
|
|
overlay_minimax_h3_int8_convrot_decoder(vae, int8_convrot_path)
|
|
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 getattr(fastvideo_args, "_model_index_class_names", {}).get("audio_vae")
|
|
or "LTX2AudioDecoder")
|
|
# The K6 Diffusers MMAudioVAE's config.json still nests a `vocoder_config` (checked here, not
|
|
# used for construction -- Kandinsky6AudioVAE itself has no vocoder submodule; the pipeline
|
|
# loads a separate `vocoder` component instead, see Kandinsky6AudioDecodingStage). The
|
|
# checkpoint's audio_vae/*.safetensors correspondingly still bundles a redundant `vocoder.*`
|
|
# copy of those weights for backward compatibility with pre-split code; filtered out below.
|
|
bundled_mmaudio = class_name == "MMAudioVAE" and "vocoder_config" in config
|
|
if bundled_mmaudio:
|
|
class_name = "Kandinsky6AudioVAE"
|
|
config["need_vae_encoder"] = True
|
|
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.
|
|
# Kandinsky6AudioVAE wraps MMAudioVAE internally (self.vae), so the
|
|
# same precision constraint applies to it too.
|
|
_needs_fp32_construction = class_name in ("MMAudioVAE", "Kandinsky6AudioVAE")
|
|
construction_precision = "fp32" if _needs_fp32_construction else precision
|
|
construction_device = torch.device("cpu") if _needs_fp32_construction 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 in ("MMAudioVAE", "Kandinsky6AudioVAE"):
|
|
if bundled_mmaudio:
|
|
loaded = {k: v for k, v in loaded.items() if not k.startswith("vocoder.")}
|
|
audio_decoder.load_state_dict(loaded, strict=True)
|
|
if bundled_mmaudio:
|
|
# Diffusers removes MPConv normalization before saving this export.
|
|
audio_decoder.vae._weights_normalized = True
|
|
else:
|
|
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()
|
|
|
|
|
|
def _mmaudio_vocoder_to_bigvgan(class_name: str, config: dict) -> tuple[str, dict]:
|
|
"""Diffusers' standalone MMAudioVocoder is BigVGAN-v2. Its config.json omits a few hyperparameters
|
|
that Diffusers' fixed-architecture reimplementation hardcodes instead of exposing (resblock type "1",
|
|
snakebeta activation with log-scale, no output bias/tanh); fill them in. A no-op for any other
|
|
class_name.
|
|
"""
|
|
if class_name != "MMAudioVocoder":
|
|
return class_name, config
|
|
config = dict(config)
|
|
config.setdefault("resblock", "1")
|
|
config.setdefault("activation", "snakebeta")
|
|
config.setdefault("snake_logscale", True)
|
|
config.setdefault("use_bias_at_final", False)
|
|
config.setdefault("use_tanh_at_final", False)
|
|
# The checkpoint's vocoder/*.safetensors is saved with weight_norm already removed (plain
|
|
# `conv_pre.weight`, not `conv_pre.parametrizations.weight.original0/1`). BigVGANV2's constructor
|
|
# applies weight_norm parametrization by default; strip it during construction (before
|
|
# VocoderLoader's subsequent load_state_dict(strict=True) below) so the fresh module's keys match
|
|
# what's on disk.
|
|
config.setdefault("weight_norm_removed", True)
|
|
return "BigVGANV2", config
|
|
|
|
|
|
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"
|
|
class_name, config = _mmaudio_vocoder_to_bigvgan(class_name, config)
|
|
|
|
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", None)
|
|
config.pop("_name_or_path", None)
|
|
if cls_name is None:
|
|
# Some checkpoints' component config.json omits _class_name even
|
|
# though model_index.json declares it (see composed_pipeline_base
|
|
# .load_modules); fall back to that before giving up.
|
|
cls_name = getattr(fastvideo_args, "_model_index_class_names", {}).get("transformer")
|
|
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}")
|
|
from fastvideo.layers.quantization.nvfp4_config import dense_transformer_safetensors
|
|
safetensors_list = dense_transformer_safetensors(safetensors_list)
|
|
|
|
# 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)
|
|
# Layerwise offload keeps every block's weights in pinned host memory, so load them on the CPU and
|
|
# attach the hooks before anything moves to the GPU; loading on the GPU first would need the whole
|
|
# DiT resident once, which is exactly what offload exists to avoid on small cards.
|
|
# Scoped to H3: other models keep main's GPU load, which their quantization paths expect.
|
|
layerwise_load = (cls_name.startswith("MiniMaxH3") and fastvideo_args.inference_mode
|
|
and fastvideo_args.dit_layerwise_offload and not fastvideo_args.use_fsdp_inference)
|
|
# The AdaLN host cache also needs the projection weights to stay off the device from the start.
|
|
adaln_table = envs.FASTVIDEO_H3_ADALN_TABLE.get() or None
|
|
adaln_host_cache = (fastvideo_args.inference_mode and not fastvideo_args.use_fsdp_inference
|
|
and (envs.FASTVIDEO_H3_ADALN_CACHE.get() or adaln_table is not None))
|
|
layerwise_load = layerwise_load or adaln_host_cache
|
|
model = maybe_load_fsdp_model(
|
|
model_cls=model_cls,
|
|
init_params={
|
|
"config": dit_config,
|
|
"hf_config": hf_config
|
|
},
|
|
weight_dir_list=safetensors_list,
|
|
device=torch.device("cpu") if layerwise_load else 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 adaln_host_cache and hasattr(model, "enable_adaln_host_cache"):
|
|
model.enable_adaln_host_cache(adaln_table)
|
|
logger.info("AdaLN modulation: %s", "precomputed tables from " + adaln_table if adaln_table
|
|
else "projections in pinned host memory behind a per-timestep cache")
|
|
if layerwise_load and not fastvideo_args.dit_layerwise_offload:
|
|
model = model.to(get_local_torch_device())
|
|
|
|
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)
|
|
# Blocks now hold placeholders; the remaining (non-block) weights and buffers belong on the GPU.
|
|
model = model.to(get_local_torch_device())
|
|
else:
|
|
logger.warning(
|
|
"Layerwise offload requested but model %s does not have "
|
|
"nn.ModuleList structure. Skipping layerwise offload.", cls_name)
|
|
if layerwise_load:
|
|
model = model.to(get_local_torch_device())
|
|
# FASTVIDEO_H3_SPLICE_TRANSFORMER=<transformer dir>: a second checkpoint of the same architecture
|
|
# runs denoising steps FASTVIDEO_H3_SPLICE_FROM_STEP (default 4) onward.
|
|
# Only the primary ``transformer`` component splices; the spliced load itself never does.
|
|
splice_path = envs.FASTVIDEO_H3_SPLICE_TRANSFORMER.get()
|
|
if (splice_path and not getattr(self, "_loading_splice", False) and hasattr(model, "attach_step_splice")
|
|
and os.path.basename(os.path.normpath(model_path)) == "transformer"):
|
|
self._loading_splice = True
|
|
try:
|
|
late = self.load(splice_path, fastvideo_args)
|
|
finally:
|
|
self._loading_splice = False
|
|
from_step = envs.FASTVIDEO_H3_SPLICE_FROM_STEP.get()
|
|
model.attach_step_splice(late, from_step)
|
|
logger.info("Step splice: steps >= %d run the transformer from %s", from_step, splice_path)
|
|
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", None)
|
|
if class_name is None:
|
|
class_name = getattr(fastvideo_args, "_model_index_class_names", {}).get("scheduler")
|
|
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)
|
|
|
|
if getattr(scheduler_cls, "is_piflow", False):
|
|
overrides = {
|
|
name: value
|
|
for name in ("eps", "final_step_size_scale", "num_policy_substeps")
|
|
if (value := getattr(fastvideo_args.pipeline_config, "piflow_" + name, None)) is not None
|
|
}
|
|
# Retain the checkpoint shift, not the base K6 pipeline default.
|
|
return scheduler_cls.from_config(config, **overrides)
|
|
|
|
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)
|
|
# Some bundles (e.g. Kandinsky6 SR's latent_upscaler) declare the class only in model_index.json, whose
|
|
# entries are keyed by the component directory name.
|
|
component_name = os.path.basename(os.path.normpath(str(model_path)))
|
|
class_name = (config_dict.pop("_class_name", None)
|
|
or getattr(fastvideo_args, "_model_index_class_names", {}).get(component_name))
|
|
|
|
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)
|
|
single_config = False
|
|
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)
|
|
elif isinstance(upsampler_config_attr, UpsamplerConfig) and type(upsampler_config_attr) is not UpsamplerConfig:
|
|
# A pipeline that declares one concrete upsampler config: fill it from the component's config.json.
|
|
config_dict.pop("_name_or_path", None)
|
|
upsampler_cfg = deepcopy(upsampler_config_attr)
|
|
upsampler_cfg.update_model_config(config_dict)
|
|
single_config = True
|
|
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)
|
|
upsampler_precision = getattr(fastvideo_args.pipeline_config, "upsampler_precision", "bf16")
|
|
# Build directly in the target precision: large upsamplers (the Kandinsky6 SR bank has 3.7B parameters)
|
|
# would otherwise be materialised and initialised in fp32 first.
|
|
with (set_default_torch_dtype(PRECISION_TO_TYPE[upsampler_precision]) if single_config else nullcontext()):
|
|
model = model_cls(upsampler_cfg)
|
|
|
|
target_device = get_local_torch_device()
|
|
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)
|
|
|
|
if single_config:
|
|
model.requires_grad_(False)
|
|
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)
|