diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md index d083b460a..7e9091b51 100644 --- a/docs/contributing/env_vars.md +++ b/docs/contributing/env_vars.md @@ -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`. | diff --git a/fastvideo/configs/pipelines/base.py b/fastvideo/configs/pipelines/base.py index 129782b1a..c12864cd0 100644 --- a/fastvideo/configs/pipelines/base.py +++ b/fastvideo/configs/pipelines/base.py @@ -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: diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 08197588c..5fcb25876 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -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 ================== diff --git a/fastvideo/models/encoders/minimax_h3_qwen3_vl.py b/fastvideo/models/encoders/minimax_h3_qwen3_vl.py index c348140d1..974d55e07 100644 --- a/fastvideo/models/encoders/minimax_h3_qwen3_vl.py +++ b/fastvideo/models/encoders/minimax_h3_qwen3_vl.py @@ -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 diff --git a/fastvideo/models/hf_transformer_utils.py b/fastvideo/models/hf_transformer_utils.py index 5f26179d9..4beb3d0d7 100644 --- a/fastvideo/models/hf_transformer_utils.py +++ b/fastvideo/models/hf_transformer_utils.py @@ -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. diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 72e38b81b..906f0af4e 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -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) diff --git a/fastvideo/models/loader/text_encoder_quantization.py b/fastvideo/models/loader/text_encoder_quantization.py index d62fa176a..a2f6b666d 100644 --- a/fastvideo/models/loader/text_encoder_quantization.py +++ b/fastvideo/models/loader/text_encoder_quantization.py @@ -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}") diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 6e150a2f3..2f86e2e3b 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -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) diff --git a/fastvideo/tests/encoders/test_clip_encoder.py b/fastvideo/tests/encoders/test_clip_encoder.py index e988175d8..6a7696d2a 100644 --- a/fastvideo/tests/encoders/test_clip_encoder.py +++ b/fastvideo/tests/encoders/test_clip_encoder.py @@ -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) diff --git a/fastvideo/tests/encoders/test_hyt5_encoder.py b/fastvideo/tests/encoders/test_hyt5_encoder.py index 568fce3c5..4f19986a7 100644 --- a/fastvideo/tests/encoders/test_hyt5_encoder.py +++ b/fastvideo/tests/encoders/test_hyt5_encoder.py @@ -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) diff --git a/fastvideo/tests/encoders/test_llama_encoder.py b/fastvideo/tests/encoders/test_llama_encoder.py index 6b940160f..05597f252 100644 --- a/fastvideo/tests/encoders/test_llama_encoder.py +++ b/fastvideo/tests/encoders/test_llama_encoder.py @@ -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") diff --git a/fastvideo/tests/encoders/test_qwen2_5_encoder.py b/fastvideo/tests/encoders/test_qwen2_5_encoder.py index b6a09bb5d..2206cf9b8 100644 --- a/fastvideo/tests/encoders/test_qwen2_5_encoder.py +++ b/fastvideo/tests/encoders/test_qwen2_5_encoder.py @@ -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) diff --git a/fastvideo/tests/encoders/test_t5_encoder.py b/fastvideo/tests/encoders/test_t5_encoder.py index 01d40348c..c49aca877 100644 --- a/fastvideo/tests/encoders/test_t5_encoder.py +++ b/fastvideo/tests/encoders/test_t5_encoder.py @@ -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) diff --git a/fastvideo/tests/loader/test_text_encoder_unified_memory_offload.py b/fastvideo/tests/loader/test_text_encoder_unified_memory_offload.py index c955358a3..a96a2c275 100644 --- a/fastvideo/tests/loader/test_text_encoder_unified_memory_offload.py +++ b/fastvideo/tests/loader/test_text_encoder_unified_memory_offload.py @@ -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) diff --git a/fastvideo/tests/loader/test_transformer_loader_config.py b/fastvideo/tests/loader/test_transformer_loader_config.py new file mode 100644 index 000000000..d501e4fd7 --- /dev/null +++ b/fastvideo/tests/loader/test_transformer_loader_config.py @@ -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 diff --git a/fastvideo/tests/platforms/test_unified_memory_offload.py b/fastvideo/tests/platforms/test_unified_memory_offload.py index cc35c3001..959624e07 100644 --- a/fastvideo/tests/platforms/test_unified_memory_offload.py +++ b/fastvideo/tests/platforms/test_unified_memory_offload.py @@ -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 diff --git a/fastvideo/tests/transformers/test_cosmos.py b/fastvideo/tests/transformers/test_cosmos.py index eb2bdaae6..52dd72c8d 100644 --- a/fastvideo/tests/transformers/test_cosmos.py +++ b/fastvideo/tests/transformers/test_cosmos.py @@ -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) diff --git a/fastvideo/tests/transformers/test_flux.py b/fastvideo/tests/transformers/test_flux.py index 5917fd481..d3f932749 100644 --- a/fastvideo/tests/transformers/test_flux.py +++ b/fastvideo/tests/transformers/test_flux.py @@ -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) diff --git a/fastvideo/tests/transformers/test_hunyuangamecraft.py b/fastvideo/tests/transformers/test_hunyuangamecraft.py index 66864a79f..fd32d8a36 100644 --- a/fastvideo/tests/transformers/test_hunyuangamecraft.py +++ b/fastvideo/tests/transformers/test_hunyuangamecraft.py @@ -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) diff --git a/fastvideo/tests/transformers/test_hyworld.py b/fastvideo/tests/transformers/test_hyworld.py index 14ffea7b9..0d43c9484 100644 --- a/fastvideo/tests/transformers/test_hyworld.py +++ b/fastvideo/tests/transformers/test_hyworld.py @@ -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) diff --git a/fastvideo/tests/transformers/test_wanvideo.py b/fastvideo/tests/transformers/test_wanvideo.py index 1827cd667..31edacbfd 100644 --- a/fastvideo/tests/transformers/test_wanvideo.py +++ b/fastvideo/tests/transformers/test_wanvideo.py @@ -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() diff --git a/fastvideo/tests/vaes/test_hunyuan15_vae.py b/fastvideo/tests/vaes/test_hunyuan15_vae.py index 707139426..2b15a69fa 100644 --- a/fastvideo/tests/vaes/test_hunyuan15_vae.py +++ b/fastvideo/tests/vaes/test_hunyuan15_vae.py @@ -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() diff --git a/fastvideo/tests/vaes/test_minimax_h3_parallel_vae.py b/fastvideo/tests/vaes/test_minimax_h3_parallel_vae.py index a881dcb21..a2e0a88bf 100644 --- a/fastvideo/tests/vaes/test_minimax_h3_parallel_vae.py +++ b/fastvideo/tests/vaes/test_minimax_h3_parallel_vae.py @@ -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"}) diff --git a/fastvideo/tests/vaes/test_wan_vae.py b/fastvideo/tests/vaes/test_wan_vae.py index b12d7ecd9..15487ff86 100644 --- a/fastvideo/tests/vaes/test_wan_vae.py +++ b/fastvideo/tests/vaes/test_wan_vae.py @@ -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. diff --git a/fastvideo/utils.py b/fastvideo/utils.py index b42c3f606..68e36aae2 100644 --- a/fastvideo/utils.py +++ b/fastvideo/utils.py @@ -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,