nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[refactor]: read typed config paths in the loaders

Browse Source 
- The component loaders read engine.offload.*, engine.precision.*,
  engine.parallelism.hsdp_*, engine.use_fsdp_inference, engine.compile.*,
  engine.attention.*, pipeline.components.*, and pipeline.flow_shift (the
  scheduler shift) from the resolved config, and compile settings through
  torch_compile_kwargs(). They no longer write model_paths (the pipeline's
  ComponentState records paths) or apply the unified-memory policy to the
  config; the text-encoder loader only queries the policy for its target
  device.
- PipelineComponentLoader.load_module and TransformerLoader.load take
  loading_teacher_critic_model instead of reading an attribute that
  training code set on the config.
- Loader, encoder, transformer, VAE, and device-policy tests build resolved
  configs; the env-var docs name typed paths.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_012M91wnVFmPEvJ9h39r5BH7
This commit is contained in:
Davids048
2026-10-03 06:46:52 +00:00
co-authored by Claude Opus 5.5
parent 6c2a250fd0
commit 0b0c692012
25 changed files with 545 additions and 279 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/contributing/env_vars.md
+8 -7
View File
@@ -122,8 +122,9 @@ in `examples/`, `scripts/`, `docs/`, `apps/`, and the tests in the same pull req
minor release.
To remove a variable that no code reads, delete its entry and add the name to `DEPRECATED_VARIABLES` in
`fastvideo/envs.py` with a reason. `FastVideoArgs` calls `envs.warn_deprecated_variables()`, which logs a warning for
each listed variable that is set. Delete the entry in the next minor release.
`fastvideo/envs.py` with a reason. The config resolution step `warn_deprecated_environment_variables` in
`fastvideo/api/inference_resolution.py` calls `envs.warn_deprecated_variables()`, which logs a warning for each listed
variable that is set. Delete the entry in the next minor release.
## What the contract test checks
@@ -188,7 +189,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
| `FASTVIDEO_LOGGING_LEVEL` | str | `INFO` | logging | Default logging level. |
| `FASTVIDEO_LOGGING_PREFIX` | str | `""` | logging | Prefix prepended to every log message. |
| `FASTVIDEO_STAGE_LOGGING` | bool | `0` | logging | Log the time that each pipeline stage takes. |
| `FASTVIDEO_ATTENTION_BACKEND` | str | unset | attention | Attention backend, as an AttentionBackendEnum name such as TORCH_SDPA, FLASH_ATTN, VIDEO_SPARSE_ATTN, SAGE_ATTN, or SAGE_ATTN_THREE. FastVideoArgs uses it when FastVideoArgs.attention_backend is unset. An unsupported name raises an error. |
| `FASTVIDEO_ATTENTION_BACKEND` | str | unset | attention | Attention backend, as an AttentionBackendEnum name such as TORCH_SDPA, FLASH_ATTN, VIDEO_SPARSE_ATTN, SAGE_ATTN, or SAGE_ATTN_THREE. Config resolution uses it when engine.attention.backend is unset. An unsupported name raises an error. |
| `FASTVIDEO_FA4` | bool | `0` | attention | The FLASH_ATTN backend uses FlashAttention-4 (flash_attn.cute) instead of FA3 or FA2. |
| `FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN` | bool | `0` | attention | MiniMax-H3 dense DiT self-attention uses the FlashAttention-4 packed-varlen entry point. This changes the floating-point reduction order, so it is an inference-only opt-in. |
| `FASTVIDEO_VSA_SM100A` | bool | `0` | attention | VIDEO_SPARSE_ATTN_H3 sends no-grad tile-64 forwards to the data-center Blackwell (sm_100a) kernel. fastvideo-kernel reads the same variable with the same rule. |
@@ -196,9 +197,9 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
| `FASTVIDEO_DISABLE_ATTENTION_COMPILE` | bool | `1` | attention | Keep attention forward out of torch.compile graphs (torch.compiler.disable). Set it to 0 to let attention constructed under that setting be traced. Setting it explicitly to true also blocks regional compile. |
| `FASTVIDEO_MLX_WINDOW` | int | `0` | attention | MLX FastWan windowed attention size in tokens. 0 uses full attention. |
| `FASTVIDEO_MLX_WINDOW_SINK` | int | `0` | attention | Number of sink tokens that MLX windowed attention always attends to. |
| `FASTVIDEO_INFERENCE_TORCH_COMPILE` | bool | `0` | performance | Compile each DiT transformer block with fullgraph torch.compile at inference. Same as FastVideoArgs.inference_torch_compile=True. |
| `FASTVIDEO_VAE_PARALLEL_DECODE` | bool | `0` | performance | MiniMax-H3 VAE decode splits its temporal chunks across the sequence-parallel ranks instead of running serially on the output rank. Same as FastVideoArgs.vae_parallel_decode=True. |
| `FASTVIDEO_VAE_PARALLEL_ENCODE` | bool | `0` | performance | MiniMax-H3 reference-video VAE encode splits its temporal chunks across the sequence-parallel ranks. Same as FastVideoArgs.vae_parallel_encode=True. |
| `FASTVIDEO_INFERENCE_TORCH_COMPILE` | bool | `0` | performance | Compile each DiT transformer block with fullgraph torch.compile at inference. Same as engine.compile.regional=True. |
| `FASTVIDEO_VAE_PARALLEL_DECODE` | bool | `0` | performance | MiniMax-H3 VAE decode splits its temporal chunks across the sequence-parallel ranks instead of running serially on the output rank. Same as pipeline.minimax_h3.vae_parallel_decode=True. |
| `FASTVIDEO_VAE_PARALLEL_ENCODE` | bool | `0` | performance | MiniMax-H3 reference-video VAE encode splits its temporal chunks across the sequence-parallel ranks. Same as pipeline.minimax_h3.vae_parallel_encode=True. |
| `FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY` | str | unset | performance | Collective that moves chunks in parallel VAE decode: gather (used when unset) or all_gather. |
| `FASTVIDEO_MINIMAX_H3_FUSIONS` | str | `""` | performance | MiniMax-H3 inference-only Triton fusions: all, 1, or a comma-separated subset of modulate,qknorm_rope,swiglu. Empty, 0, or none keeps the eager implementation. |
| `FASTVIDEO_FSDP2_AUTOWRAP` | bool | `0` | performance | FSDP2 shards modules by parameter count instead of the model's shard conditions. Not supported by self-forcing distillation. |
@@ -234,7 +235,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
| `FASTVIDEO_LTX2_GEMMA_LOG` | str | `""` | debug | Log file for LTX-2 Gemma text-encoder hidden states, used by parity tests. Deprecated names: `LTX2_FASTVIDEO_GEMMA_LOG`. |
| `FASTVIDEO_COSMOS25_LOG_KNOBS` | bool | `0` | debug | Log the Cosmos 2.5 latent-preparation conditioning inputs. |
| `FASTVIDEO_CFG_GATE_STEP` | float | `1.0` | sampling | CFG gating fraction in [0, 1]. Steps before len(timesteps) \* X run the conditional and unconditional forwards; later steps reuse the cached difference. 1.0 disables gating. |
| `FASTVIDEO_LTX2_USE_DISTILLED_SIGMAS` | bool | `1` | sampling | LTX-2 uses the distilled sigma schedule when FastVideoArgs.ltx2_use_distilled_sigmas is also true. Deprecated names: `LTX2_USE_DISTILLED_SIGMAS`. |
| `FASTVIDEO_LTX2_USE_DISTILLED_SIGMAS` | bool | `1` | sampling | LTX-2 uses the distilled sigma schedule when pipeline.ltx2.use_distilled_sigmas is also true. Deprecated names: `LTX2_USE_DISTILLED_SIGMAS`. |
| `FASTVIDEO_EVAL_CACHE` | path | computed | eval | Cache directory for evaluation models and datasets. Defaults to $FASTVIDEO_CACHE_ROOT/eval. |
| `FASTVIDEO_PHYSICS_IQ_BUCKET_URL` | str | `https://storage.googleapis.com/physics-iq-benchmark` | eval | Base URL of the Physics-IQ benchmark bucket. |
| `FASTVIDEO_VBENCH_FULL_INFO_JSON` | str | unset | eval | Path to VBench_full_info.json, used instead of the vendored copy. Deprecated names: `VBENCH_FULL_INFO_JSON`. |
fastvideo/configs/pipelines/base.py
+2 -2
View File
@@ -91,8 +91,8 @@ class PipelineConfig:
def __setattr__(self, name: str, value: Any) -> None:
if self.__dict__.get("_frozen") and name in self.__dataclass_fields__:
raise AttributeError(f"{type(self).__name__}.{name} was decided by config resolution and is read-only; "
f"use fastvideo_args.override(source, {{'pipeline_config.{name}': value}}) to "
"change it")
f"to change pipeline_config.{name}, override its typed path with "
"resolved_config.with_override(source, {path: value}) and use the returned config")
super().__setattr__(name, value)
def freeze(self) -> None:
fastvideo/envs.py
+5 -5
View File
@@ -314,7 +314,7 @@ FASTVIDEO_ATTENTION_BACKEND = EnvStr(
None,
category="attention",
doc="Attention backend, as an AttentionBackendEnum name such as TORCH_SDPA, FLASH_ATTN, VIDEO_SPARSE_ATTN, "
"SAGE_ATTN, or SAGE_ATTN_THREE. FastVideoArgs uses it when FastVideoArgs.attention_backend is unset. An "
"SAGE_ATTN, or SAGE_ATTN_THREE. Config resolution uses it when engine.attention.backend is unset. An "
"unsupported name raises an error.")
# FA4 is opt-in and never auto-selected just because it is installed. Below
# sm90, grad-enabled and GQA calls are routed to FA2 (FA4's backward asserts
@@ -358,17 +358,17 @@ FASTVIDEO_INFERENCE_TORCH_COMPILE = EnvBool(
False,
category="performance",
doc="Compile each DiT transformer block with fullgraph torch.compile at inference. Same as "
"FastVideoArgs.inference_torch_compile=True.")
"engine.compile.regional=True.")
FASTVIDEO_VAE_PARALLEL_DECODE = EnvBool(
False,
category="performance",
doc="MiniMax-H3 VAE decode splits its temporal chunks across the sequence-parallel ranks instead of running "
"serially on the output rank. Same as FastVideoArgs.vae_parallel_decode=True.")
"serially on the output rank. Same as pipeline.minimax_h3.vae_parallel_decode=True.")
FASTVIDEO_VAE_PARALLEL_ENCODE = EnvBool(
False,
category="performance",
doc="MiniMax-H3 reference-video VAE encode splits its temporal chunks across the sequence-parallel ranks. "
"Same as FastVideoArgs.vae_parallel_encode=True.")
"Same as pipeline.minimax_h3.vae_parallel_encode=True.")
FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY = EnvStr(
None,
category="performance",
@@ -496,7 +496,7 @@ FASTVIDEO_CFG_GATE_STEP = EnvFloat(
FASTVIDEO_LTX2_USE_DISTILLED_SIGMAS = EnvBool(
True,
category="sampling",
doc="LTX-2 uses the distilled sigma schedule when FastVideoArgs.ltx2_use_distilled_sigmas is also true.",
doc="LTX-2 uses the distilled sigma schedule when pipeline.ltx2.use_distilled_sigmas is also true.",
deprecated_names=("LTX2_USE_DISTILLED_SIGMAS", ))
# ================== Evaluation ==================
fastvideo/models/encoders/minimax_h3_qwen3_vl.py
+2 -2
View File
@@ -666,8 +666,8 @@ class MiniMaxH3Qwen3VLConditioner(TextEncoder[torch.Tensor]):
f"tokens={int(mask.sum())}, features={features.shape[0]}")
return mask
# no_grad, NOT inference_mode: with text_encoder_cpu_offload=True (the
# FastVideoArgs default) the loader FSDP2-shards this conditioner, and
# no_grad, NOT inference_mode: with engine.offload.text_encoder=True (the
# default) the loader FSDP2-shards this conditioner, and
# FSDP2's wait_for_unshard reads tensor._version via
# _unsafe_preserve_version_counter - inference tensors do not track
# version counters, so inference_mode crashes the first encode. no_grad
fastvideo/models/hf_transformer_utils.py
-1
View File
@@ -80,7 +80,6 @@ def get_diffusers_config(model: str, ) -> dict[str, Any]:
Args:
model: The model name or path.
fastvideo_args: Optional inference arguments to override in the config.
Returns:
The loaded configuration.
fastvideo/models/loader/component_loader.py
+108 -80
View File
@@ -19,7 +19,9 @@ 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 offload_disabled_on_unified_memory
from fastvideo.api.device_policy import UNIFIED_MEMORY_OFFLOAD_PATHS, offload_disabled_on_unified_memory
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,
@@ -28,7 +30,6 @@ from fastvideo.attention.selector import (
)
from fastvideo.configs.models import EncoderConfig
from fastvideo.distributed import get_local_torch_device
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.layers.quantization import get_quantization_config
from fastvideo.logger import init_logger
from fastvideo.models.hf_transformer_utils import get_diffusers_config
@@ -52,6 +53,12 @@ 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."""
@@ -59,13 +66,13 @@ class ComponentLoader(ABC):
self.device = device
@abstractmethod
def load(self, model_path: str, resolved_config: FastVideoArgs):
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
fastvideo_args: FastVideoArgs
resolved_config: The resolved runtime config
Returns:
The loaded component
@@ -248,12 +255,12 @@ class TextEncoderLoader(ComponentLoader):
for source in secondary_weights:
yield from self._get_weights_iterator(source, to_cpu)
def load(self, model_path: str, resolved_config: FastVideoArgs):
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=fastvideo_args.trust_remote_code,
# revision=fastvideo_args.revision,
# trust_remote_code=resolved_config.trust_remote_code,
# revision=resolved_config.revision,
# model_override_args=None,
# )
model_config = get_diffusers_config(model=model_path)
@@ -307,7 +314,7 @@ class TextEncoderLoader(ComponentLoader):
except Exception:
idx = 0
encoder_configs = resolved_config.pipeline_config.text_encoder_configs
encoder_precisions = resolved_config.pipeline_config.text_encoder_precisions
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}"
@@ -344,7 +351,7 @@ class TextEncoderLoader(ComponentLoader):
model_path: str,
model_config: EncoderConfig,
target_device: torch.device,
resolved_config: FastVideoArgs,
resolved_config: ResolvedGeneratorConfig,
dtype: str = "fp16",
use_text_encoder_override: bool = False, # prevent subclasses from misusing
cpu_offload: bool | None = None,
@@ -352,14 +359,16 @@ class TextEncoderLoader(ComponentLoader):
):
runtime_device = get_local_torch_device()
device_id = runtime_device.index if runtime_device.index is not None else 0
requested_cpu_offload = getattr(resolved_config, offload_flag) if cpu_offload is None else cpu_offload
# A resolved config already carries the worker's offload policy, so only ask whether the policy covers this
# component; a FastVideoArgs applies the policy to itself.
apply_policy = getattr(resolved_config, "disable_offload_on_unified_memory", None)
if apply_policy is not None:
disable_cpu_offload = apply_policy(device_id, offload_flag=offload_flag)
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).
offload_field = UNIFIED_MEMORY_OFFLOAD_PATHS[offload_flag].rsplit(".", 1)[1]
requested_cpu_offload = getattr(resolved_config.engine.offload, offload_field)
else:
disable_cpu_offload = offload_disabled_on_unified_memory(device_id, offload_flag)
requested_cpu_offload = cpu_offload
# The resolved config already carries the worker's offload policy, so only ask whether the policy covers
# this component.
disable_cpu_offload = offload_disabled_on_unified_memory(device_id, offload_flag)
if requested_cpu_offload and disable_cpu_offload:
# Direct loader callers can choose a CPU target before the worker
@@ -390,10 +399,13 @@ class TextEncoderLoader(ComponentLoader):
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 resolved_config.override_text_encoder_quant is not None:
if 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")
"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 "
@@ -404,10 +416,11 @@ class TextEncoderLoader(ComponentLoader):
checkpoint_quant_config.get_name(),
checkpoint_path,
)
elif use_text_encoder_override and resolved_config.override_text_encoder_quant is not None:
if resolved_config.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(resolved_config.override_text_encoder_quant)
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):
@@ -427,7 +440,7 @@ class TextEncoderLoader(ComponentLoader):
model = model_cls(model_config) # type: ignore
weights_to_load = {name for name, _ in model.named_parameters()}
if (use_text_encoder_override and resolved_config.override_text_encoder_safetensors is not None):
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,
@@ -476,7 +489,7 @@ class TextEncoderLoader(ComponentLoader):
from fastvideo.platforms import current_platform
if use_cpu_offload:
pin_cpu_memory = resolved_config.pin_cpu_memory and is_pin_memory_available()
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")
@@ -513,12 +526,12 @@ class TextEncoderLoader(ComponentLoader):
class ImageEncoderLoader(TextEncoderLoader):
def load(self, model_path: str, resolved_config: FastVideoArgs):
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=fastvideo_args.trust_remote_code,
# revision=fastvideo_args.revision,
# 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:
@@ -547,7 +560,7 @@ class ImageEncoderLoader(TextEncoderLoader):
f"image encoder index {index} requires pipeline_config."
"image_encoder_configs")
encoder_config = resolved_config.pipeline_config.image_encoder_config
encoder_precision = resolved_config.pipeline_config.image_encoder_precision
encoder_precision = resolved_config.engine.precision.image_encoder
else:
if index < 0 or index >= len(encoder_configs):
raise IndexError(
@@ -564,7 +577,7 @@ class ImageEncoderLoader(TextEncoderLoader):
from fastvideo.platforms import current_platform
if resolved_config.image_encoder_cpu_offload:
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()
@@ -575,7 +588,7 @@ class ImageEncoderLoader(TextEncoderLoader):
target_device,
resolved_config,
encoder_precision,
cpu_offload=resolved_config.image_encoder_cpu_offload,
cpu_offload=resolved_config.engine.offload.image_encoder,
offload_flag="image_encoder_cpu_offload",
)
@@ -583,7 +596,7 @@ class ImageEncoderLoader(TextEncoderLoader):
class VisionLanguageEncoderLoader(ComponentLoader):
"""Loader for vision-language autoregressive encoders."""
def load(self, model_path: str, resolved_config: FastVideoArgs):
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)
@@ -600,7 +613,7 @@ class VisionLanguageEncoderLoader(ComponentLoader):
class ProcessorLoader(ComponentLoader):
"""Loader for HF processors that pair with vision-language encoders."""
def load(self, model_path: str, resolved_config: FastVideoArgs):
def load(self, model_path: str, resolved_config: ResolvedGeneratorConfig):
from transformers import AutoProcessor
logger.info("Loading processor from %s", model_path)
@@ -615,7 +628,7 @@ class ProcessorLoader(ComponentLoader):
class ImageProcessorLoader(ComponentLoader):
"""Loader for image processor."""
def load(self, model_path: str, resolved_config: FastVideoArgs):
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)
@@ -627,7 +640,7 @@ class ImageProcessorLoader(ComponentLoader):
class TokenizerLoader(ComponentLoader):
"""Loader for tokenizers."""
def load(self, model_path: str, resolved_config: FastVideoArgs):
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
@@ -759,24 +772,23 @@ class VAELoader(ComponentLoader):
return directory
return None
def load(self, model_path: str, resolved_config: FastVideoArgs):
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.")
resolved_config.model_paths["vae"] = model_path
from fastvideo.platforms import current_platform
if resolved_config.vae_cpu_offload:
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()
with set_default_torch_dtype(PRECISION_TO_TYPE[resolved_config.pipeline_config.vae_precision] if resolved_config.
pipeline_config.vae_precision else torch.bfloat16):
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"
@@ -786,7 +798,7 @@ class VAELoader(ComponentLoader):
if is_gen3c and class_name in ("AutoencoderKLWan", "AutoencoderKLGen3CTokenizer"):
from fastvideo.models.vaes.gen3c_tokenizer_vae import (AutoencoderKLGen3CTokenizer)
dtype = PRECISION_TO_TYPE[resolved_config.pipeline_config.vae_precision]
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:
@@ -835,7 +847,7 @@ class VAELoader(ComponentLoader):
if class_name == "AutoencoderKLWan" and is_cosmos25:
from fastvideo.models.vaes.cosmos25wanvae import Cosmos25WanVAE
dtype = PRECISION_TO_TYPE[resolved_config.pipeline_config.vae_precision]
dtype = PRECISION_TO_TYPE[vae_precision]
vae = Cosmos25WanVAE(device=target_device, dtype=dtype)
weight_path = os.path.join(model_path, "tokenizer.safetensors")
@@ -846,7 +858,7 @@ class VAELoader(ComponentLoader):
return vae.eval()
if class_name == "LingBotWorld2WanVAE":
dtype = PRECISION_TO_TYPE[resolved_config.pipeline_config.vae_precision]
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)
@@ -927,7 +939,7 @@ class VAELoader(ComponentLoader):
class AudioDecoderLoader(ComponentLoader):
"""Loader for full audio VAEs and legacy decode-only audio components."""
def load(self, model_path: str, resolved_config: FastVideoArgs):
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)
@@ -942,7 +954,7 @@ class AudioDecoderLoader(ComponentLoader):
config.pop("_name_or_path", None)
audio_vae_config = deepcopy(configured_audio_vae)
audio_vae_config.update_model_arch(config)
if getattr(resolved_config, "vae_cpu_offload", False):
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)
@@ -990,7 +1002,7 @@ class AudioDecoderLoader(ComponentLoader):
class VocoderLoader(ComponentLoader):
"""Loader for native vocoders."""
def load(self, model_path: str, resolved_config: FastVideoArgs):
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"
@@ -1035,8 +1047,15 @@ def _collect_safetensors_keys(safetensors_list: list) -> set:
class TransformerLoader(ComponentLoader):
"""Loader for transformer."""
def load(self, model_path: str, resolved_config: FastVideoArgs):
"""Load the transformer based on the model path, and inference args."""
def load(self,
model_path: str,
resolved_config: ResolvedGeneratorConfig,
loading_teacher_critic_model: bool = False):
"""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.
"""
config = get_diffusers_config(model=model_path)
hf_config = deepcopy(config)
cls_name = config.pop("_class_name")
@@ -1046,25 +1065,24 @@ class TransformerLoader(ComponentLoader):
"Only diffusers format is supported.")
logger.info("transformer cls_name: %s", cls_name)
if resolved_config.override_transformer_cls_name is not None:
cls_name = resolved_config.override_transformer_cls_name
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)
resolved_config.model_paths["transformer"] = model_path
# 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 this flag set and must stay
# 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 = hasattr(resolved_config, "_loading_teacher_critic_model")
_qat_generator_only = loading_teacher_critic_model
if _qat_generator_only:
dit_config.quant_config = None
@@ -1082,13 +1100,14 @@ class TransformerLoader(ComponentLoader):
update_fn(weight_keys)
# Check if we should use custom initialization weights
custom_weights_path = getattr(resolved_config, "init_weights_from_safetensors", None)
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 hasattr(resolved_config, "_loading_teacher_critic_model"))
and not loading_teacher_critic_model)
if use_custom_weights:
if "transformer_2" in model_path:
custom_weights_path = getattr(resolved_config, "init_weights_from_safetensors_2", None)
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"))
@@ -1103,11 +1122,12 @@ class TransformerLoader(ComponentLoader):
safetensors_list,
)
default_dtype = PRECISION_TO_TYPE[resolved_config.pipeline_config.dit_precision]
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)
assert resolved_config.hsdp_shard_dim is not None
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.
@@ -1134,27 +1154,27 @@ class TransformerLoader(ComponentLoader):
},
weight_dir_list=safetensors_list,
device=get_local_torch_device(),
hsdp_replicate_dim=resolved_config.hsdp_replicate_dim,
hsdp_shard_dim=resolved_config.hsdp_shard_dim,
hsdp_replicate_dim=parallelism.hsdp_replicate_dim,
hsdp_shard_dim=parallelism.hsdp_shard_dim,
strict=strict_load,
cpu_offload=resolved_config.dit_cpu_offload,
pin_cpu_memory=resolved_config.pin_cpu_memory,
fsdp_inference=resolved_config.use_fsdp_inference,
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.enable_torch_compile,
torch_compile_kwargs=resolved_config.torch_compile_kwargs,
inference_regional_compile=resolved_config.inference_torch_compile,
inference_vsa_tile_size=resolved_config.VSA_tile_size,
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=getattr(resolved_config, "lora_path", None),
lora_strength=getattr(resolved_config, "lora_strength", 1.0),
lora_path=components.lora_path,
lora_strength=components.lora_strength,
)
total_params = sum(p.numel() for p in model.parameters())
@@ -1167,7 +1187,7 @@ class TransformerLoader(ComponentLoader):
model = model.eval()
if resolved_config.inference_mode and resolved_config.dit_layerwise_offload:
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:
@@ -1182,7 +1202,7 @@ class TransformerLoader(ComponentLoader):
class SchedulerLoader(ComponentLoader):
"""Loader for scheduler."""
def load(self, model_path: str, resolved_config: FastVideoArgs):
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)
@@ -1193,8 +1213,9 @@ class SchedulerLoader(ComponentLoader):
scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name)
scheduler = scheduler_cls(**config)
if resolved_config.pipeline_config.flow_shift is not None:
scheduler.set_shift(resolved_config.pipeline_config.flow_shift)
flow_shift = resolved_config.pipeline.flow_shift
if flow_shift is not None:
scheduler.set_shift(flow_shift)
return scheduler
@@ -1209,7 +1230,7 @@ class ConditionerLoader(ComponentLoader):
fetched sub-encoders (T5) don't trip the missing-key check.
"""
def load(self, model_path: str, resolved_config: FastVideoArgs):
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)
@@ -1254,7 +1275,7 @@ class ConditionerLoader(ComponentLoader):
class UpsamplerLoader(ComponentLoader):
"""Loader for upsamplers (incl. LTX-2 spatial/temporal upsamplers)."""
def load(self, model_path: str, resolved_config: FastVideoArgs):
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)
@@ -1326,7 +1347,7 @@ class GenericComponentLoader(ComponentLoader):
super().__init__()
self.library = library
def load(self, model_path: str, resolved_config: FastVideoArgs):
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",
@@ -1369,7 +1390,8 @@ class PipelineComponentLoader:
module_name: str,
component_model_path: str,
transformers_or_diffusers: str,
resolved_config: FastVideoArgs,
resolved_config: ResolvedGeneratorConfig,
loading_teacher_critic_model: bool = False,
):
"""
Load a pipeline module.
@@ -1378,7 +1400,9 @@ class PipelineComponentLoader:
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
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
Returns:
The loaded module
@@ -1396,7 +1420,7 @@ class PipelineComponentLoader:
# 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
# engine.attention.backend is the process-wide default, applied
# per component.
#
# The loaders record the decision on their own config rather than this
@@ -1406,9 +1430,13 @@ class PipelineComponentLoader:
if _active_component_attention_backend_scope() is not None:
attention_context = nullcontext()
else:
resolved = coerce_attn_backend(getattr(resolved_config, "attention_backend", None))
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)
return loader.load(component_model_path, resolved_config)
fastvideo/models/loader/text_encoder_quantization.py
+3 -3
View File
@@ -10,8 +10,8 @@ import torch
import torch.nn as nn
from safetensors.torch import safe_open
from fastvideo.api.resolution import ResolvedGeneratorConfig
from fastvideo.configs.models import EncoderConfig
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.layers.linear import LinearBase, UnquantizedLinearMethod
from fastvideo.layers.quantization.base_config import QuantizationConfig
from fastvideo.models.encoders.base import TextEncoder
@@ -19,10 +19,10 @@ from fastvideo.models.encoders.base import TextEncoder
def _resolve_text_encoder_checkpoint_path(
model_path: str,
resolved_config: FastVideoArgs,
resolved_config: ResolvedGeneratorConfig,
use_text_encoder_override: bool,
) -> str:
override = resolved_config.override_text_encoder_safetensors if use_text_encoder_override else None
override = resolved_config.pipeline.components.text_encoder_weights if use_text_encoder_override else None
checkpoint_path = override or model_path
if not os.path.exists(checkpoint_path):
raise FileNotFoundError(f"Text-encoder checkpoint does not exist: {checkpoint_path}")
fastvideo/registry.py
+1 -1
View File
@@ -59,7 +59,7 @@ from fastvideo.api.matrixgame2 import MatrixGame2SamplingParam
from fastvideo.api.matrixgame3 import MatrixGame3SamplingParam
from fastvideo.api.flux import FluxSamplingParam
from fastvideo.fastvideo_args import WorkloadType
from fastvideo.api.schema import WorkloadType
from fastvideo.logger import init_logger
from fastvideo.utils import (maybe_download_model_index, verify_model_config_and_directory)
fastvideo/tests/encoders/test_clip_encoder.py
+5 -2
View File
@@ -9,7 +9,7 @@ import gc
import fastvideo.envs as envs
from fastvideo.configs.pipelines import HunyuanConfig, PipelineConfig
from fastvideo.forward_context import set_forward_context
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.api.inference_resolution import resolve_inference_config
from fastvideo.logger import init_logger
from fastvideo.utils import maybe_download_model
from fastvideo.configs.models.encoders import CLIPTextConfig
@@ -38,7 +38,10 @@ def test_clip_encoder():
- Load models with the same weights and parameters
- Produce nearly identical outputs for the same input prompts
"""
args = FastVideoArgs(model_path="openai/clip-vit-large-patch14", pipeline_config=HunyuanConfig())
args = resolve_inference_config({
"model_path": BASE_MODEL_PATH,
"pipeline": {"experimental": {"pipeline_config": HunyuanConfig()}},
})
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
logger.info("Loading models from %s", args.model_path)
fastvideo/tests/encoders/test_hyt5_encoder.py
+9 -4
View File
@@ -14,7 +14,7 @@ from fastvideo.forward_context import set_forward_context
from fastvideo.logger import init_logger
from fastvideo.models.loader.component_loader import TextEncoderLoader
from fastvideo.utils import maybe_download_model, PRECISION_TO_TYPE
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.api.inference_resolution import resolve_inference_config
from fastvideo.configs.models.encoders import T5Config
logger = init_logger(__name__)
@@ -22,11 +22,12 @@ logger = init_logger(__name__)
envs.setdefault_external("MASTER_ADDR", "localhost")
envs.setdefault_external("MASTER_PORT", "29503")
BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
@pytest.fixture
def t5_model_paths_and_config():
base_model_path = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
model_path = maybe_download_model(base_model_path)
model_path = maybe_download_model(BASE_MODEL_PATH)
text_encoder_path = os.path.join(model_path, "text_encoder_2")
tokenizer_path = os.path.join(model_path, "tokenizer_2")
return text_encoder_path, tokenizer_path, Hunyuan15T2V480PConfig()
@@ -45,7 +46,11 @@ def test_t5_encoder(t5_model_paths_and_config):
model1 = T5EncoderModel.from_pretrained(text_encoder_path).to(precision).to(device).eval()
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
args = FastVideoArgs(model_path=text_encoder_path, pipeline_config=pipeline_config, pin_cpu_memory=False)
args = resolve_inference_config({
"model_path": BASE_MODEL_PATH,
"engine": {"offload": {"pin_cpu_memory": False}},
"pipeline": {"experimental": {"pipeline_config": pipeline_config}},
})
loader = TextEncoderLoader()
model2 = loader.load(text_encoder_path, args)
model2 = model2.to(precision)
fastvideo/tests/encoders/test_llama_encoder.py
+5 -2
View File
@@ -9,7 +9,7 @@ import gc
import fastvideo.envs as envs
from fastvideo.configs.pipelines import HunyuanConfig
from fastvideo.forward_context import set_forward_context
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.api.inference_resolution import resolve_inference_config
from fastvideo.logger import init_logger
from fastvideo.models.loader.component_loader import TextEncoderLoader
from fastvideo.utils import maybe_download_model
@@ -39,7 +39,10 @@ def test_llama_encoder():
- Load models with the same weights and parameters
- Produce nearly identical outputs for the same input prompts
"""
args = FastVideoArgs(model_path="meta-llama/Llama-2-7b-hf", pipeline_config=HunyuanConfig())
args = resolve_inference_config({
"model_path": BASE_MODEL_PATH,
"pipeline": {"experimental": {"pipeline_config": HunyuanConfig()}},
})
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
fastvideo/tests/encoders/test_qwen2_5_encoder.py
+9 -4
View File
@@ -12,7 +12,7 @@ from fastvideo.forward_context import set_forward_context
from fastvideo.logger import init_logger
from fastvideo.models.loader.component_loader import TextEncoderLoader
from fastvideo.utils import maybe_download_model, PRECISION_TO_TYPE
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.api.inference_resolution import resolve_inference_config
from fastvideo.configs.models.encoders import Qwen2_5_VLConfig
logger = init_logger(__name__)
@@ -20,11 +20,12 @@ logger = init_logger(__name__)
envs.setdefault_external("MASTER_ADDR", "localhost")
envs.setdefault_external("MASTER_PORT", "29505")
BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
@pytest.fixture
def qwen_model_path_and_config():
base_model_path = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
model_path = maybe_download_model(base_model_path)
model_path = maybe_download_model(BASE_MODEL_PATH)
text_encoder_path = os.path.join(model_path, "text_encoder")
tokenizer_path = os.path.join(model_path, "tokenizer")
return text_encoder_path, tokenizer_path, Hunyuan15T2V480PConfig()
@@ -48,7 +49,11 @@ def test_qwen2_5_encoder(qwen_model_path_and_config):
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
# Load FastVideo model
args = FastVideoArgs(model_path=text_encoder_path, pipeline_config=pipeline_config, pin_cpu_memory=False)
args = resolve_inference_config({
"model_path": BASE_MODEL_PATH,
"engine": {"offload": {"pin_cpu_memory": False}},
"pipeline": {"experimental": {"pipeline_config": pipeline_config}},
})
loader = TextEncoderLoader()
model2 = loader.load(text_encoder_path, args)
fastvideo/tests/encoders/test_t5_encoder.py
+16 -5
View File
@@ -15,7 +15,7 @@ from fastvideo.logger import init_logger
from fastvideo.models.loader.component_loader import TextEncoderLoader
from fastvideo.tests.utils import skip_if_gated_repo_inaccessible
from fastvideo.utils import maybe_download_model, PRECISION_TO_TYPE
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.api.inference_resolution import resolve_inference_config
from fastvideo.configs.models.encoders import T5Config, T5LargeConfig
logger = init_logger(__name__)
@@ -23,10 +23,13 @@ logger = init_logger(__name__)
envs.setdefault_external("MASTER_ADDR", "localhost")
envs.setdefault_external("MASTER_PORT", "29503")
WAN_BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
COSMOS_BASE_MODEL_PATH = "nvidia/Cosmos-Predict2-2B-Video2World"
@pytest.fixture
def t5_model_paths_and_config():
base_model_path = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
base_model_path = WAN_BASE_MODEL_PATH
model_path = maybe_download_model(base_model_path, local_dir=os.path.join('data', base_model_path))
text_encoder_path = os.path.join(model_path, "text_encoder")
tokenizer_path = os.path.join(model_path, "tokenizer")
@@ -35,7 +38,7 @@ def t5_model_paths_and_config():
@pytest.fixture
def t5_large_model_paths_and_config():
base_model_path = "nvidia/Cosmos-Predict2-2B-Video2World"
base_model_path = COSMOS_BASE_MODEL_PATH
local_dir = os.path.join('data', base_model_path)
skip_if_gated_repo_inaccessible(base_model_path, local_path=local_dir, test_name="Cosmos T5-large encoder test")
model_path = maybe_download_model(base_model_path, local_dir=local_dir)
@@ -57,7 +60,11 @@ def test_t5_encoder(t5_model_paths_and_config):
model1 = UMT5EncoderModel.from_pretrained(text_encoder_path).to(precision).to(device).eval()
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
args = FastVideoArgs(model_path=text_encoder_path, pipeline_config=pipeline_config, pin_cpu_memory=False)
args = resolve_inference_config({
"model_path": WAN_BASE_MODEL_PATH,
"engine": {"offload": {"pin_cpu_memory": False}},
"pipeline": {"experimental": {"pipeline_config": pipeline_config}},
})
loader = TextEncoderLoader()
model2 = loader.load(text_encoder_path, args)
model2 = model2.to(precision)
@@ -137,7 +144,11 @@ def test_t5_large_encoder(t5_large_model_paths_and_config):
model1 = T5EncoderModel.from_pretrained(text_encoder_path).to(precision).to(device).eval()
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
args = FastVideoArgs(model_path=text_encoder_path, pipeline_config=pipeline_config, pin_cpu_memory=False)
args = resolve_inference_config({
"model_path": COSMOS_BASE_MODEL_PATH,
"engine": {"offload": {"pin_cpu_memory": False}},
"pipeline": {"experimental": {"pipeline_config": pipeline_config}},
})
loader = TextEncoderLoader()
model2 = loader.load(text_encoder_path, args)
model2 = model2.to(precision)
fastvideo/tests/loader/test_text_encoder_unified_memory_offload.py
+20 -14
View File
@@ -9,8 +9,11 @@ import pytest
import torch
import torch.nn as nn
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.api.inference_resolution import resolve_inference_config
from fastvideo.models.loader.component_loader import ImageEncoderLoader, TextEncoderLoader
from fastvideo.tests.api.config_snapshot import isolated_environment
WAN_T2V = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
class _PassthroughEncoder(nn.Module):
@@ -28,6 +31,12 @@ def _model_config():
return SimpleNamespace(architectures=["PassthroughEncoder"], _fsdp_shard_conditions=[], quant_config=None)
def _resolved_config(**offload):
"""A resolved inference config for a registered model with the given ``engine.offload`` values."""
with isolated_environment():
return resolve_inference_config({"model_path": WAN_T2V, "engine": {"offload": offload}})
@pytest.mark.parametrize(
("cpu_offload", "requested_target"),
[
@@ -47,19 +56,21 @@ def test_unified_memory_uses_worker_device_before_model_construction(monkeypatch
"fastvideo.models.loader.component_loader.ModelRegistry.resolve_model_cls",
lambda architectures: (_PassthroughEncoder, None),
)
args = FastVideoArgs(model_path=str(tmp_path), text_encoder_cpu_offload=True)
resolved_config = _resolved_config(text_encoder=True)
model = TextEncoderLoader().load_model(
str(tmp_path),
_model_config(),
requested_target,
args,
resolved_config,
cpu_offload=cpu_offload,
)
assert isinstance(model, _PassthroughEncoder)
assert _PassthroughEncoder.loaded_device == torch.device("cuda:5")
assert args.text_encoder_cpu_offload is False
# The loader only queries the policy; the worker records the offload decisions on the config.
assert resolved_config.engine.offload.text_encoder is True
assert resolved_config.override_log == ()
probe.assert_called_once_with(5)
@@ -71,13 +82,13 @@ def test_discrete_memory_preserves_explicit_cpu_target(monkeypatch, tmp_path) ->
"fastvideo.models.loader.component_loader.ModelRegistry.resolve_model_cls",
lambda architectures: (_PassthroughEncoder, None),
)
args = FastVideoArgs(model_path=str(tmp_path), text_encoder_cpu_offload=False)
resolved_config = _resolved_config(text_encoder=False)
TextEncoderLoader().load_model(
str(tmp_path),
_model_config(),
torch.device("cuda:2"),
args,
resolved_config,
cpu_offload=True,
)
@@ -94,22 +105,17 @@ def test_image_encoder_explicit_offload_resets_cpu_target(monkeypatch, tmp_path)
"fastvideo.models.loader.component_loader.ModelRegistry.resolve_model_cls",
lambda architectures: (_PassthroughEncoder, None),
)
args = FastVideoArgs(
model_path=str(tmp_path),
text_encoder_cpu_offload=True,
image_encoder_cpu_offload=True,
)
resolved_config = _resolved_config(text_encoder=True, image_encoder=True)
ImageEncoderLoader().load_model(
str(tmp_path),
_model_config(),
torch.device("cpu"),
args,
resolved_config,
cpu_offload=True,
offload_flag="image_encoder_cpu_offload",
)
assert _PassthroughEncoder.loaded_device == torch.device("cuda:4")
assert args.text_encoder_cpu_offload is False
assert args.image_encoder_cpu_offload is False
assert resolved_config.override_log == ()
probe.assert_called_once_with(4)
fastvideo/tests/loader/test_transformer_loader_config.py
+122
View File
@@ -0,0 +1,122 @@
# SPDX-License-Identifier: Apache-2.0
"""CPU tests for the resolved-config values that TransformerLoader hands to the FSDP model loader."""
from __future__ import annotations
import json
import pytest
import torch
import torch.nn as nn
import fastvideo.models.loader.component_loader as component_loader
from fastvideo.api.inference_resolution import resolve_inference_config, torch_compile_kwargs
from fastvideo.models.loader.component_loader import PipelineComponentLoader, TransformerLoader
from fastvideo.tests.api.config_snapshot import isolated_environment
WAN_T2V = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
class _FakeTransformer(nn.Module):
"""Model class that the registry returns for the test checkpoint."""
def _resolve(raw: dict):
"""A resolved inference config for a registered Wan model with ``raw`` merged into its input."""
with isolated_environment():
return resolve_inference_config({"model_path": WAN_T2V, **raw})
@pytest.fixture
def transformer_dir(tmp_path):
"""A transformer component directory with a Diffusers config and one weight file."""
component = tmp_path / "transformer"
component.mkdir()
(component / "config.json").write_text(json.dumps({"_class_name": "WanTransformer3DModel"}))
(component / "model.safetensors").write_bytes(b"")
return component
@pytest.fixture
def fsdp_calls(monkeypatch):
"""Keyword arguments of each ``maybe_load_fsdp_model`` call; each call returns a one-layer model."""
calls = []
def fake_maybe_load_fsdp_model(**kwargs):
calls.append(kwargs)
return nn.Linear(1, 1).to(kwargs["default_dtype"])
monkeypatch.setattr(component_loader, "maybe_load_fsdp_model", fake_maybe_load_fsdp_model)
monkeypatch.setattr(component_loader.ModelRegistry, "resolve_model_cls", lambda name: (_FakeTransformer, None))
monkeypatch.setattr(component_loader, "get_local_torch_device", lambda: torch.device("cpu"))
return calls
def test_transformer_loader_passes_the_typed_engine_settings(transformer_dir, fsdp_calls) -> None:
resolved_config = _resolve({
"engine": {
"use_fsdp_inference": True,
"offload": {
"dit": False,
"pin_cpu_memory": False
},
"compile": {
"enabled": True,
"mode": "max-autotune",
"extras": {
"dynamic": False
}
},
"attention": {
"vsa_tile_size": 128
},
"parallelism": {
"hsdp_shard_dim": 1
},
},
"pipeline": {
"components": {
"lora_path": "/adapters/style",
"lora_strength": 0.5
}
},
})
TransformerLoader().load(str(transformer_dir), resolved_config)
(call, ) = fsdp_calls
assert call["weight_dir_list"] == [str(transformer_dir / "model.safetensors")]
assert call["default_dtype"] is torch.bfloat16
assert (call["hsdp_replicate_dim"], call["hsdp_shard_dim"]) == (1, 1)
assert (call["cpu_offload"], call["pin_cpu_memory"], call["fsdp_inference"]) == (False, False, True)
assert call["training_mode"] is False
assert call["enable_torch_compile"] is True
assert call["torch_compile_kwargs"] == torch_compile_kwargs(resolved_config) == {
"mode": "max-autotune",
"dynamic": False
}
assert call["inference_regional_compile"] is False
assert call["inference_vsa_tile_size"] == 128
assert (call["lora_path"], call["lora_strength"]) == ("/adapters/style", 0.5)
def test_teacher_critic_flag_loads_checkpoint_weights_without_quantization(monkeypatch, tmp_path, transformer_dir,
fsdp_calls) -> None:
generator_weights = tmp_path / "generator.safetensors"
generator_weights.write_bytes(b"")
resolved_config = _resolve({"pipeline": {"components": {"transformer_weights": str(generator_weights)}}})
quant_config = object()
monkeypatch.setattr(resolved_config.pipeline_config.dit_config, "quant_config", quant_config)
PipelineComponentLoader.load_module("transformer", str(transformer_dir), "diffusers", resolved_config)
PipelineComponentLoader.load_module("transformer",
str(transformer_dir),
"diffusers",
resolved_config,
loading_teacher_critic_model=True)
generator, teacher = fsdp_calls
assert generator["weight_dir_list"] == [str(generator_weights)]
assert generator["init_params"]["config"].quant_config is not None
assert teacher["weight_dir_list"] == [str(transformer_dir / "model.safetensors")]
assert teacher["init_params"]["config"].quant_config is None
assert resolved_config.pipeline_config.dit_config.quant_config is quant_config
fastvideo/tests/platforms/test_unified_memory_offload.py
+117 -84
View File
@@ -7,14 +7,39 @@ from unittest.mock import Mock
import pytest
from fastvideo.fastvideo_args import ExecutionMode, UNIFIED_MEMORY_OFFLOAD_FLAGS, FastVideoArgs
from fastvideo.api.device_policy import (UNIFIED_MEMORY_OFFLOAD_PATHS, disable_offload_on_unified_memory,
finalize_device_offload_policy)
from fastvideo.api.inference_resolution import resolve_inference_config
from fastvideo.api.schema import OffloadConfig
from fastvideo.api.training_schema import resolve_training_config
from fastvideo.tests.api.config_snapshot import isolated_environment
WAN_T2V = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
UNIFIED_MEMORY_OFFLOAD_FLAGS = tuple(UNIFIED_MEMORY_OFFLOAD_PATHS)
def _args_with_offloads(*enabled_flags: str, **overrides) -> FastVideoArgs:
kwargs = {flag: flag in enabled_flags for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS}
kwargs.update(model_path="unused/for-this-test")
kwargs.update(overrides)
return FastVideoArgs(**kwargs)
def _offload_field(flag: str) -> str:
"""The ``engine.offload`` field name of a device-policy offload flag."""
return UNIFIED_MEMORY_OFFLOAD_PATHS[flag].rsplit(".", 1)[1]
def _resolve(engine: dict | None = None):
"""A resolved inference config for a registered model with the given ``engine`` section."""
with isolated_environment():
return resolve_inference_config({"model_path": WAN_T2V, "engine": engine or {}})
def _with_offloads(*enabled_flags: str, **engine):
"""A resolved inference config whose host offload modes are on exactly for ``enabled_flags``."""
offload = {_offload_field(flag): flag in enabled_flags for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS}
return _resolve({"offload": offload, **engine})
def _enabled_flags(resolved_config) -> list[str]:
"""The host offload modes that are on in ``resolved_config``."""
return [
flag for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS if getattr(resolved_config.engine.offload, _offload_field(flag))
]
@pytest.fixture
@@ -26,85 +51,95 @@ def as_unified_cuda(monkeypatch):
return probe
def test_constructing_inference_args_defers_device_policy(monkeypatch) -> None:
def test_resolving_an_inference_config_defers_device_policy(monkeypatch) -> None:
probe = Mock(side_effect=AssertionError("device probe ran in the driver"))
monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", probe)
args = FastVideoArgs(model_path="unused/for-this-test", use_fsdp_inference=True)
resolved_config = _resolve({"use_fsdp_inference": True})
assert args.use_fsdp_inference is True
assert args.dit_layerwise_offload is True
assert args.dit_cpu_offload is True
assert resolved_config.engine.use_fsdp_inference is True
assert resolved_config.engine.offload.dit_layerwise is True
assert resolved_config.engine.offload.dit is True
probe.assert_not_called()
def test_training_construction_retains_offload_conflict_normalization(monkeypatch) -> None:
def test_training_resolution_retains_offload_conflict_normalization(monkeypatch) -> None:
monkeypatch.setattr("fastvideo.platforms.current_platform.is_mps", lambda: False)
args = FastVideoArgs(
model_path="unused/for-this-test",
mode=ExecutionMode.FINETUNING,
inference_mode=False,
sp_size=1,
hsdp_shard_dim=1,
use_fsdp_inference=True,
)
with isolated_environment():
resolved_config = resolve_training_config({
"model_path": WAN_T2V,
"engine": {
"use_fsdp_inference": True,
"offload": {
"dit": True,
"dit_layerwise": True
},
"parallelism": {
"sp_size": 1,
"hsdp_shard_dim": 1
},
},
})
assert args.dit_layerwise_offload is True
assert args.dit_cpu_offload is False
assert args.use_fsdp_inference is False
assert resolved_config.engine.offload.dit_layerwise is True
assert resolved_config.engine.offload.dit is False
assert resolved_config.engine.use_fsdp_inference is False
def test_policy_list_covers_every_declared_offload_flag() -> None:
declared = {field.name for field in dataclasses.fields(FastVideoArgs) if field.name.endswith("_offload")}
def test_policy_list_covers_every_host_offload_field() -> None:
declared = {field.name for field in dataclasses.fields(OffloadConfig)} - {"pin_cpu_memory", "lazy_module_load"}
assert declared == set(UNIFIED_MEMORY_OFFLOAD_FLAGS)
assert declared == {_offload_field(flag) for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS}
def test_unified_device_disables_every_offload_flag(as_unified_cuda) -> None:
args = FastVideoArgs(model_path="unused/for-this-test")
resolved_config = _resolve()
assert args.finalize_device_offload_policy(device_id=6) is True
decided = finalize_device_offload_policy(resolved_config, device_id=6)
as_unified_cuda.assert_called_once_with(6)
assert not [flag for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS if getattr(args, flag)]
assert not _enabled_flags(decided)
assert _enabled_flags(resolved_config) == list(UNIFIED_MEMORY_OFFLOAD_FLAGS)
@pytest.mark.parametrize("flag", UNIFIED_MEMORY_OFFLOAD_FLAGS)
def test_each_offload_flag_is_independently_disabled(as_unified_cuda, flag: str) -> None:
args = _args_with_offloads(flag)
assert getattr(args, flag) is True
resolved_config = _with_offloads(flag)
assert _enabled_flags(resolved_config) == [flag]
args.disable_offload_on_unified_memory(device_id=2)
decided = disable_offload_on_unified_memory(resolved_config, device_id=2)
assert getattr(args, flag) is False
assert not _enabled_flags(decided)
assert decided.override_log == (("device_policy:unified_memory", {UNIFIED_MEMORY_OFFLOAD_PATHS[flag]: False}), )
def test_discrete_device_classification_preserves_offload_requests(monkeypatch) -> None:
probe = Mock(return_value=False)
monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", probe)
args = FastVideoArgs(model_path="unused/for-this-test")
resolved_config = _resolve()
assert args.disable_offload_on_unified_memory(device_id=3) is False
assert disable_offload_on_unified_memory(resolved_config, device_id=3) is resolved_config
probe.assert_called_once_with(3)
assert all(getattr(args, flag) for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS)
assert _enabled_flags(resolved_config) == list(UNIFIED_MEMORY_OFFLOAD_FLAGS)
def test_discrete_device_finalization_retains_layerwise_precedence(monkeypatch) -> None:
monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", lambda device_id: False)
monkeypatch.setattr("fastvideo.platforms.current_platform.is_mps", lambda: False)
args = FastVideoArgs(model_path="unused/for-this-test", use_fsdp_inference=True)
resolved_config = _resolve({"use_fsdp_inference": True})
args.finalize_device_offload_policy(device_id=3)
decided = finalize_device_offload_policy(resolved_config, device_id=3)
assert args.dit_layerwise_offload is True
assert args.dit_cpu_offload is False
assert args.use_fsdp_inference is False
assert args.text_encoder_cpu_offload is True
assert args.image_encoder_cpu_offload is True
assert args.vae_cpu_offload is True
assert args.lazy_module_load is False
offload = decided.engine.offload
assert offload.dit_layerwise is True
assert offload.dit is False
assert decided.engine.use_fsdp_inference is False
assert offload.text_encoder is True
assert offload.image_encoder is True
assert offload.vae is True
assert offload.lazy_module_load is False
def test_workers_classify_their_own_device(monkeypatch) -> None:
@@ -117,65 +152,62 @@ def test_workers_classify_their_own_device(monkeypatch) -> None:
monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", has_unified_memory)
monkeypatch.setattr("fastvideo.platforms.current_platform.get_device_name", lambda device_id: "NVIDIA GB10")
monkeypatch.setattr("fastvideo.platforms.current_platform.is_mps", lambda: False)
device_zero_args = FastVideoArgs(model_path="unused/for-this-test")
device_one_args = FastVideoArgs(model_path="unused/for-this-test")
resolved_config = _resolve()
device_zero_args.finalize_device_offload_policy(device_id=0)
device_one_args.finalize_device_offload_policy(device_id=1)
device_zero = finalize_device_offload_policy(resolved_config, device_id=0)
device_one = finalize_device_offload_policy(resolved_config, device_id=1)
assert seen_device_ids == [0, 1]
assert device_zero_args.dit_layerwise_offload is True
assert device_zero_args.text_encoder_cpu_offload is True
assert not [flag for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS if getattr(device_one_args, flag)]
assert device_zero.engine.offload.dit_layerwise is True
assert device_zero.engine.offload.text_encoder is True
assert not _enabled_flags(device_one)
def test_repeated_policy_on_same_device_reuses_classification(as_unified_cuda) -> None:
args = FastVideoArgs(model_path="unused/for-this-test")
def test_policy_applied_to_its_own_result_changes_nothing(as_unified_cuda) -> None:
decided = finalize_device_offload_policy(_resolve(), device_id=1)
args.finalize_device_offload_policy(device_id=1)
args.finalize_device_offload_policy(device_id=1)
again = finalize_device_offload_policy(decided, device_id=1)
as_unified_cuda.assert_called_once_with(1)
assert again.override_log == decided.override_log
assert again.to_dict() == decided.to_dict()
def test_mps_clears_offload_and_keeps_its_fsdp_rule(monkeypatch) -> None:
monkeypatch.setattr("fastvideo.platforms.current_platform.is_mps", lambda: True)
monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", lambda device_id: True)
monkeypatch.setattr("fastvideo.platforms.current_platform.get_device_name", lambda device_id: "mps")
args = FastVideoArgs(model_path="unused/for-this-test", use_fsdp_inference=True)
resolved_config = _resolve({"use_fsdp_inference": True})
args.finalize_device_offload_policy()
decided = finalize_device_offload_policy(resolved_config)
assert args.use_fsdp_inference is False
assert not [flag for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS if getattr(args, flag)]
assert decided.engine.use_fsdp_inference is False
assert not _enabled_flags(decided)
def test_cuda_unified_memory_preserves_realistic_fsdp_request(as_unified_cuda) -> None:
args = FastVideoArgs(model_path="unused/for-this-test", use_fsdp_inference=True)
assert args.dit_layerwise_offload is True
assert args.use_fsdp_inference is True
resolved_config = _resolve({"use_fsdp_inference": True})
assert resolved_config.engine.offload.dit_layerwise is True
assert resolved_config.engine.use_fsdp_inference is True
args.finalize_device_offload_policy()
decided = finalize_device_offload_policy(resolved_config)
assert args.use_fsdp_inference is True
assert not [flag for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS if getattr(args, flag)]
assert decided.engine.use_fsdp_inference is True
assert not _enabled_flags(decided)
def test_pin_cpu_memory_is_not_a_host_offload_mode(as_unified_cuda) -> None:
args = FastVideoArgs(model_path="unused/for-this-test", pin_cpu_memory=True)
resolved_config = _resolve({"offload": {"pin_cpu_memory": True}})
args.finalize_device_offload_policy()
decided = finalize_device_offload_policy(resolved_config)
assert "pin_cpu_memory" not in UNIFIED_MEMORY_OFFLOAD_FLAGS
assert args.pin_cpu_memory is True
assert decided.engine.offload.pin_cpu_memory is True
def test_already_disabled_flags_stay_disabled(as_unified_cuda) -> None:
args = _args_with_offloads()
decided = finalize_device_offload_policy(_with_offloads())
args.finalize_device_offload_policy()
assert not [flag for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS if getattr(args, flag)]
assert not _enabled_flags(decided)
@pytest.mark.parametrize("name_error", [NotImplementedError, ValueError, RuntimeError])
@@ -186,22 +218,23 @@ def test_platform_without_device_name_uses_generic_name(monkeypatch, name_error:
raise name_error("device name unavailable")
monkeypatch.setattr("fastvideo.platforms.current_platform.get_device_name", unsupported_name)
args = _args_with_offloads("text_encoder_cpu_offload")
resolved_config = _with_offloads("text_encoder_cpu_offload")
assert args.disable_offload_on_unified_memory() is True
assert args.text_encoder_cpu_offload is False
decided = disable_offload_on_unified_memory(resolved_config, device_id=0)
assert decided.engine.offload.text_encoder is False
def test_unified_device_auto_enables_lazy_module_load(as_unified_cuda) -> None:
args = FastVideoArgs(model_path="unused/for-this-test")
resolved_config = _resolve()
assert args.lazy_module_load is None
assert args.finalize_device_offload_policy(device_id=6) is True
assert args.lazy_module_load is True
assert resolved_config.engine.offload.lazy_module_load is None
assert finalize_device_offload_policy(resolved_config, device_id=6).engine.offload.lazy_module_load is True
def test_explicit_false_lazy_module_load_stays_off_on_unified(as_unified_cuda) -> None:
args = FastVideoArgs(model_path="unused/for-this-test", lazy_module_load=False)
resolved_config = _resolve({"offload": {"lazy_module_load": False}})
args.finalize_device_offload_policy(device_id=6)
assert args.lazy_module_load is False
decided = finalize_device_offload_policy(resolved_config, device_id=6)
assert decided.engine.offload.lazy_module_load is False
fastvideo/tests/transformers/test_cosmos.py
+13 -5
View File
@@ -9,7 +9,7 @@ from diffusers.models.transformers.transformer_cosmos import CosmosTransformer3D
import fastvideo.envs as envs
from fastvideo.configs.pipelines import PipelineConfig
from fastvideo.forward_context import set_forward_context
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.api.inference_resolution import resolve_inference_config
from fastvideo.logger import init_logger
from fastvideo.models.loader.component_loader import TransformerLoader
from fastvideo.tests.utils import skip_if_gated_repo_inaccessible
@@ -43,10 +43,18 @@ def test_cosmos2_transformer():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
precision = torch.bfloat16
precision_str = "bf16"
args = FastVideoArgs(model_path=TRANSFORMER_PATH,
dit_cpu_offload=False,
use_fsdp_inference=False,
pipeline_config=PipelineConfig(dit_config=CosmosVideoConfig(), dit_precision=precision_str))
args = resolve_inference_config({
"model_path": BASE_MODEL_PATH,
"engine": {
"offload": {"dit": False},
"use_fsdp_inference": False,
},
"pipeline": {
"experimental": {
"pipeline_config": PipelineConfig(dit_config=CosmosVideoConfig(), dit_precision=precision_str)
}
},
})
loader = TransformerLoader()
model2 = loader.load(TRANSFORMER_PATH, args).to(device, dtype=precision)
fastvideo/tests/transformers/test_flux.py
+14 -8
View File
@@ -11,9 +11,9 @@ from diffusers import FluxTransformer2DModel as HFFluxTransformer2DModel
from torch.testing import assert_close
import fastvideo.envs as envs
from fastvideo.api.inference_resolution import resolve_inference_config
from fastvideo.configs.models.dits.flux import FluxDiTConfig
from fastvideo.configs.pipelines.base import PipelineConfig
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.forward_context import set_forward_context
from fastvideo.models.loader.component_loader import TransformerLoader
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
@@ -21,6 +21,8 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
envs.setdefault_external("MASTER_ADDR", "localhost")
envs.setdefault_external("MASTER_PORT", "29517")
# Registered model id that selects the FLUX.1-dev pipeline config during resolution.
BASE_MODEL_PATH = "black-forest-labs/FLUX.1-dev"
_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
_DEFAULT_FLUX_TRANSFORMER = os.path.join(
_REPO_ROOT,
@@ -71,13 +73,17 @@ def test_flux_transformer_parity_vs_diffusers() -> None:
device = torch.device("cuda:0")
precision = torch.bfloat16
args = FastVideoArgs(
model_path=transformer_path,
dit_cpu_offload=False,
dit_layerwise_offload=False,
pipeline_config=PipelineConfig(dit_config=FluxDiTConfig(), dit_precision="bf16"),
)
args.device = device
args = resolve_inference_config({
"model_path": BASE_MODEL_PATH,
"engine": {
"offload": {"dit": False, "dit_layerwise": False},
},
"pipeline": {
"experimental": {
"pipeline_config": PipelineConfig(dit_config=FluxDiTConfig(), dit_precision="bf16")
}
},
})
generator = torch.Generator(device=device).manual_seed(0)
torch.manual_seed(0)
fastvideo/tests/transformers/test_hunyuangamecraft.py
+15 -8
View File
@@ -10,13 +10,13 @@ import pytest
import torch
import fastvideo.envs as envs
from fastvideo.api.inference_resolution import resolve_inference_config
from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig
from fastvideo.configs.pipelines import PipelineConfig
from fastvideo.distributed.parallel_state import (
get_sp_parallel_rank,
get_sp_world_size,
)
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.forward_context import set_forward_context
from fastvideo.logger import init_logger
from fastvideo.models.loader.component_loader import TransformerLoader
@@ -29,6 +29,8 @@ envs.setdefault_external("MASTER_PORT", "29507")
# Path to converted weights (local path, not HuggingFace)
TRANSFORMER_PATH = "official_weights/hunyuan-gamecraft/transformer"
# Registered model id that selects the HunyuanGameCraft pipeline config during resolution.
BASE_MODEL_PATH = "FastVideo/HunyuanGameCraft-Diffusers"
# Reference latent computed from FastVideo HunyuanGameCraft model with:
# - seed=42, batch=1, frames=9, H=44, W=80, text_seq=32
@@ -57,13 +59,18 @@ def test_hunyuangamecraft_transformer(env_overrides):
precision = torch.bfloat16
precision_str = "bf16"
args = FastVideoArgs(
model_path=TRANSFORMER_PATH,
dit_cpu_offload=False,
use_fsdp_inference=False,
pipeline_config=PipelineConfig(dit_config=HunyuanGameCraftConfig(), dit_precision=precision_str),
)
args.device = device
args = resolve_inference_config({
"model_path": BASE_MODEL_PATH,
"engine": {
"offload": {"dit": False},
"use_fsdp_inference": False,
},
"pipeline": {
"experimental": {
"pipeline_config": PipelineConfig(dit_config=HunyuanGameCraftConfig(), dit_precision=precision_str)
}
},
})
loader = TransformerLoader()
model = loader.load(TRANSFORMER_PATH, args).to(device, dtype=precision)
fastvideo/tests/transformers/test_hyworld.py
+15 -9
View File
@@ -5,13 +5,13 @@ import pytest
import torch
import fastvideo.envs as envs
from fastvideo.api.inference_resolution import resolve_inference_config
from fastvideo.configs.models.dits import HYWorldConfig
from fastvideo.configs.pipelines import PipelineConfig
from fastvideo.distributed.parallel_state import (
get_sp_parallel_rank,
get_sp_world_size,
)
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.forward_context import set_forward_context
from fastvideo.logger import init_logger
from fastvideo.models.loader.component_loader import TransformerLoader
@@ -24,7 +24,8 @@ logger = init_logger(__name__)
envs.setdefault_external("MASTER_ADDR", "localhost")
envs.setdefault_external("MASTER_PORT", "29503")
MODEL_PATH = maybe_download_model("FastVideo/HY-WorldPlay-Bidirectional-Diffusers")
BASE_MODEL_PATH = "FastVideo/HY-WorldPlay-Bidirectional-Diffusers"
MODEL_PATH = maybe_download_model(BASE_MODEL_PATH)
TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
REFERENCE_LATENTS = {
AttentionBackendEnum.FLASH_ATTN: -197132.85557549074,
@@ -57,13 +58,18 @@ def test_hyworld_transformer():
precision = torch.bfloat16
precision_str = "bf16"
args = FastVideoArgs(
model_path=transformer_path,
dit_cpu_offload=False,
use_fsdp_inference=False,
pipeline_config=PipelineConfig(dit_config=HYWorldConfig(), dit_precision=precision_str),
)
args.device = device
args = resolve_inference_config({
"model_path": BASE_MODEL_PATH,
"engine": {
"offload": {"dit": False},
"use_fsdp_inference": False,
},
"pipeline": {
"experimental": {
"pipeline_config": PipelineConfig(dit_config=HYWorldConfig(), dit_precision=precision_str)
}
},
})
loader = TransformerLoader()
model = loader.load(transformer_path, args).to(device, dtype=precision)
fastvideo/tests/transformers/test_wanvideo.py
+13 -8
View File
@@ -9,9 +9,9 @@ from torch.testing import assert_close
import fastvideo.envs as envs
from fastvideo.configs.pipelines import PipelineConfig
from fastvideo.forward_context import set_forward_context
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.api.inference_resolution import resolve_inference_config
from fastvideo.models.loader.component_loader import TransformerLoader
from fastvideo.tests.golden_gate._wan_checkpoint import component_path
from fastvideo.tests.golden_gate._wan_checkpoint import WAN_REPO, component_path
from fastvideo.models.wan.config import WanVideoConfig
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
@@ -25,12 +25,17 @@ def test_wan_transformer():
transformer_path = str(component_path("transformer"))
device = torch.device("cuda:0")
precision = torch.bfloat16
args = FastVideoArgs(
model_path=transformer_path,
dit_cpu_offload=True,
pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision="bf16"),
)
args.device = device
args = resolve_inference_config({
"model_path": WAN_REPO,
"engine": {
"offload": {"dit": True},
},
"pipeline": {
"experimental": {
"pipeline_config": PipelineConfig(dit_config=WanVideoConfig(), dit_precision="bf16")
}
},
})
loader = TransformerLoader()
candidate = loader.load(transformer_path, args).to(dtype=precision).eval()
fastvideo/tests/vaes/test_hunyuan15_vae.py
+10 -5
View File
@@ -12,7 +12,7 @@ from fastvideo.configs.pipelines import PipelineConfig
from fastvideo.logger import init_logger
# from fastvideo.models.vaes.hunyuanvae import (
# AutoencoderKLHunyuanVideo as MyHunyuanVAE)
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.api.inference_resolution import resolve_inference_config
from fastvideo.models.loader.component_loader import VAELoader
from fastvideo.configs.models.vaes import Hunyuan15VAEConfig
from fastvideo.utils import maybe_download_model
@@ -33,10 +33,15 @@ def test_hunyuan_vae():
device = torch.device("cuda:0")
precision = torch.float32
precision_str = "fp32"
args = FastVideoArgs(model_path=VAE_PATH,
pipeline_config=PipelineConfig(vae_config=Hunyuan15VAEConfig(), vae_precision=precision_str))
args.device = device
args.vae_cpu_offload = False
args = resolve_inference_config({
"model_path": BASE_MODEL_PATH,
"engine": {"offload": {"vae": False}},
"pipeline": {
"experimental": {
"pipeline_config": PipelineConfig(vae_config=Hunyuan15VAEConfig(), vae_precision=precision_str)
}
},
})
model1 = AutoencoderKLHunyuanVideo15.from_pretrained(VAE_PATH, torch_dtype=precision).to(device).eval()
model1.enable_tiling()
fastvideo/tests/vaes/test_minimax_h3_parallel_vae.py
+21 -13
View File
@@ -296,19 +296,27 @@ def test_parallel_encode_validates_input() -> None:
encode_pixels_parallel(vae, torch.zeros(1, 4, 4, 16, 16), group)
def test_fastvideo_args_strategy_literals_match_module() -> None:
"""fastvideo_args mirrors the strategy literals to avoid importing model
modules at args construction; keep the two in sync."""
from fastvideo.fastvideo_args import FastVideoArgs
def test_resolved_strategy_literals_match_module() -> None:
"""Config resolution mirrors the strategy literals to avoid importing model
modules; its default and accepted values match this module."""
from fastvideo.api.errors import ConfigValidationError
from fastvideo.api.inference_resolution import resolve_inference_config
from fastvideo.tests.api.config_snapshot import isolated_environment
args = FastVideoArgs(model_path="test/parallel-vae")
assert args.vae_parallel_decode is False
assert args.vae_parallel_encode is False
assert args.vae_parallel_decode_strategy == DEFAULT_DECODE_GATHER_STRATEGY
assert args.vae_parallel_decode_strategy in DECODE_GATHER_STRATEGIES
model_path = "FastVideo/FastVideo-FastH3-8-Step-V2"
def resolve(minimax_h3):
with isolated_environment():
return resolve_inference_config({"model_path": model_path, "pipeline": {"minimax_h3": minimax_h3}})
options = resolve({}).pipeline.minimax_h3
assert options.vae_parallel_decode is False
assert options.vae_parallel_encode is False
assert options.vae_parallel_decode_strategy == DEFAULT_DECODE_GATHER_STRATEGY
assert options.vae_parallel_decode_strategy in DECODE_GATHER_STRATEGIES
for strategy in DECODE_GATHER_STRATEGIES:
assert FastVideoArgs(model_path="test/parallel-vae",
vae_parallel_decode_strategy=strategy).vae_parallel_decode_strategy == strategy
with pytest.raises(ValueError, match="vae_parallel_decode_strategy"):
FastVideoArgs(model_path="test/parallel-vae", vae_parallel_decode_strategy="scatter")
resolved = resolve({"vae_parallel_decode_strategy": strategy})
assert resolved.pipeline.minimax_h3.vae_parallel_decode_strategy == strategy
with pytest.raises(ConfigValidationError, match="vae_parallel_decode_strategy"):
resolve({"vae_parallel_decode_strategy": "scatter"})
fastvideo/tests/vaes/test_wan_vae.py
+11 -6
View File
@@ -8,10 +8,10 @@ from torch.testing import assert_close
import fastvideo.envs as envs
from fastvideo.configs.pipelines import PipelineConfig
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.api.inference_resolution import resolve_inference_config
from fastvideo.models.loader.component_loader import VAELoader
from fastvideo.models.wan.vae_config import WanVAEConfig
from fastvideo.tests.golden_gate._wan_checkpoint import component_path
from fastvideo.tests.golden_gate._wan_checkpoint import WAN_REPO, component_path
envs.setdefault_external("MASTER_ADDR", "localhost")
envs.setdefault_external("MASTER_PORT", "29503")
@@ -23,10 +23,15 @@ def test_wan_vae():
vae_path = str(component_path("vae"))
device = torch.device("cuda:0")
precision = torch.float32
args = FastVideoArgs(model_path=vae_path,
pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision="fp32"))
args.device = device
args.vae_cpu_offload = False
args = resolve_inference_config({
"model_path": WAN_REPO,
"engine": {"offload": {"vae": False}},
"pipeline": {
"experimental": {
"pipeline_config": PipelineConfig(vae_config=WanVAEConfig(), vae_precision="fp32")
}
},
})
candidate = VAELoader().load(vae_path, args)
assert candidate.use_feature_cache # Preserve the original Wan VAE algorithm.
fastvideo/utils.py
+1 -1
View File
@@ -46,7 +46,7 @@ logger = init_logger(__name__)
T = TypeVar("T")
# TODO(will): used to convert fastvideo_args.precision to torch.dtype. Find a
# TODO(will): used to convert a precision name (engine.precision.*) to torch.dtype. Find a
# cleaner way to do this.
PRECISION_TO_TYPE = {
"fp32": torch.float32,