SolitaryThinker
b2d55a7ba0
Replace the re-export stubs with real vendored copies of the fastvideo modeling
+ layer + supporting infra, for the kept diffusion models (wan21, wan_causal,
ltx2, flux2, matrixgame2). v2 now imports ZERO `fastvideo.*` — it is
self-contained. (bagel/qwen_omni load from vllm_omni, an external pkg, not
fastvideo; cosmos3's load_id was already dangling — both out of scope here.)
Vendored (cp + `sed fastvideo. -> v2.`):
- models/ the 5 models' nn.Module dits/vaes/encoders/audio/upsamplers + the
component loader/ + the lazy class registry (other families' rows are
dormant/lazy — only the 5 resolve).
- layers/ attention/ platforms/ distributed/ configs/ logging_utils/ hooks/
third_party/pynvml + top-level forward_context/fastvideo_args/envs/logger/
utils/version — copied verbatim (layers et al. 'as is').
- api/ slimmed to schema + results (the VideoGenerator's config dataclasses);
the fastvideo parser/presets/overrides (which pull the pipeline runtime) are
intentionally NOT vendored — v2 has its own runtime/loop.
Decoupling surgery (cut the loader's coupling to the fastvideo runtime):
- configs/pipeline_registry.py (vendored from fastvideo/registry.py, renamed to
avoid colliding with v2/registry.py): dropped the _register_presets() auto-call
and matrixgame3 (removed model); config-class resolution preserved.
- configs/pipelines/__init__.py: dropped the registry back-edge (fixes an import
cycle) — base.py imports the registry lazily where used.
- torch_backend: load_component now uses v2.models.loader.
The only remaining external 'fastvideo*' refs are `fastvideo_kernel` (the
separate optional CUDA-kernel pkg for sparse/MoBA attention) — guarded; the dense
TORCH_SDPA path v2 uses never imports it.
Vendored subtrees added to the pre-commit exclude (faithful copies, mirroring the
existing fastvideo/models exclusion — not re-linted, to stay re-syncable).
Verified: grep finds zero fastvideo-package imports in v2/; Wan2.1 T2V on H100 is
BIT-IDENTICAL to the fastvideo-backed path (same .npy SHA256, byte-for-byte); CPU
mini 156 (154 passed + 2 env-skipped on x86, torch present). Backup: branch v2_backup.
98 lines
3.9 KiB
Python
98 lines
3.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Utilities for selecting and loading models."""
|
|
import contextlib
|
|
import re
|
|
from collections import defaultdict
|
|
from collections.abc import Callable, Iterator
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
from v2.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def set_default_torch_dtype(dtype: torch.dtype):
|
|
"""Sets the default torch dtype to the given dtype."""
|
|
old_dtype = torch.get_default_dtype()
|
|
torch.set_default_dtype(dtype)
|
|
yield
|
|
torch.set_default_dtype(old_dtype)
|
|
|
|
|
|
def get_param_names_mapping(
|
|
mapping_dict: dict[str, str]) -> Callable[[str], tuple[str, Any, Any]]:
|
|
"""
|
|
Creates a mapping function that transforms parameter names using regex patterns.
|
|
|
|
Args:
|
|
mapping_dict (Dict[str, str]): Dictionary mapping regex patterns to replacement patterns
|
|
param_name (str): The parameter name to be transformed
|
|
|
|
Returns:
|
|
Callable[[str], str]: A function that maps parameter names from source to target format
|
|
"""
|
|
|
|
def mapping_fn(name: str) -> tuple[str, Any, Any]:
|
|
# Try to match and transform the name using the regex patterns in mapping_dict
|
|
for pattern, replacement in mapping_dict.items():
|
|
match = re.match(pattern, name)
|
|
if match:
|
|
merge_index = None
|
|
total_splitted_params = None
|
|
if isinstance(replacement, tuple):
|
|
merge_index = replacement[1]
|
|
total_splitted_params = replacement[2]
|
|
replacement = replacement[0]
|
|
name = re.sub(pattern, replacement, name)
|
|
return name, merge_index, total_splitted_params
|
|
|
|
# If no pattern matches, return the original name
|
|
return name, None, None
|
|
|
|
return mapping_fn
|
|
|
|
|
|
def hf_to_custom_state_dict(
|
|
hf_param_sd: dict[str, torch.Tensor] | Iterator[tuple[str, torch.Tensor]],
|
|
param_names_mapping: Callable[[str], tuple[str, Any, Any]]
|
|
) -> tuple[dict[str, torch.Tensor], dict[str, tuple[str, Any, Any]]]:
|
|
"""
|
|
Converts a Hugging Face parameter state dictionary to a custom parameter state dictionary.
|
|
|
|
Args:
|
|
hf_param_sd (Dict[str, torch.Tensor]): The Hugging Face parameter state dictionary
|
|
param_names_mapping (Callable[[str], tuple[str, Any, Any]]): A function that maps parameter names from source to target format
|
|
|
|
Returns:
|
|
custom_param_sd (Dict[str, torch.Tensor]): The custom formatted parameter state dict
|
|
reverse_param_names_mapping (Dict[str, Tuple[str, Any, Any]]): Maps back from custom to hf
|
|
"""
|
|
custom_param_sd = {}
|
|
to_merge_params = defaultdict(dict) # type: ignore
|
|
reverse_param_names_mapping = {}
|
|
if isinstance(hf_param_sd, dict):
|
|
hf_param_sd = hf_param_sd.items() # type: ignore
|
|
for source_param_name, full_tensor in hf_param_sd: # type: ignore
|
|
target_param_name, merge_index, num_params_to_merge = param_names_mapping(
|
|
source_param_name)
|
|
reverse_param_names_mapping[target_param_name] = (source_param_name,
|
|
merge_index,
|
|
num_params_to_merge)
|
|
if merge_index is not None:
|
|
to_merge_params[target_param_name][merge_index] = full_tensor
|
|
if len(to_merge_params[target_param_name]) == num_params_to_merge:
|
|
# cat at output dim according to the merge_index order
|
|
sorted_tensors = [
|
|
to_merge_params[target_param_name][i]
|
|
for i in range(num_params_to_merge)
|
|
]
|
|
full_tensor = torch.cat(sorted_tensors, dim=0)
|
|
del to_merge_params[target_param_name]
|
|
else:
|
|
continue
|
|
custom_param_sd[target_param_name] = full_tensor
|
|
return custom_param_sd, reverse_param_names_mapping
|