Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 William Lin
|
ba371e33ed |
@@ -1,34 +0,0 @@
|
||||
name: Close Stale Issues and PRs
|
||||
|
||||
on:
|
||||
schedule:
|
||||
# Daily at 1:30 AM UTC
|
||||
- cron: '30 1 * * *'
|
||||
|
||||
jobs:
|
||||
stale:
|
||||
if: github.repository == 'hao-ai-lab/FastVideo'
|
||||
permissions:
|
||||
issues: write
|
||||
pull-requests: write
|
||||
actions: write
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/stale@997185467fa4f803885201cee163a9f38240193d # v10.1.1
|
||||
with:
|
||||
operations-per-run: 500
|
||||
|
||||
exempt-draft-pr: true
|
||||
exempt-issue-labels: 'keep-open,pinned,security,Bug,RFC'
|
||||
exempt-pr-labels: 'keep-open,pinned'
|
||||
|
||||
labels-to-add-when-unstale: 'unstale'
|
||||
labels-to-remove-when-stale: 'unstale'
|
||||
|
||||
days-before-issue-stale: 90
|
||||
days-before-issue-close: 30
|
||||
stale-issue-label: 'stale'
|
||||
|
||||
days-before-pr-stale: 60
|
||||
days-before-pr-close: 14
|
||||
stale-pr-label: 'stale'
|
||||
@@ -461,12 +461,11 @@ The reusable implementation lives in
|
||||
| Workflow | Trigger | Purpose |
|
||||
|---|---|---|
|
||||
| `community-issue-labeler.yml` | Issue opened or edited | Adds scope/platform labels from issue keywords. |
|
||||
| `community-stale.yml` | Daily | Labels/closes stale items with exemptions. |
|
||||
|
||||
None of these automations post comments or messages on issues or pull
|
||||
requests; they only set labels and close stale items. Keep it that way: new
|
||||
automation must report through checks, commit statuses, or labels instead of
|
||||
PR/issue comments.
|
||||
requests; they only set labels. Issues and pull requests are never marked
|
||||
stale or closed automatically. Keep it that way: new automation must report
|
||||
through checks, commit statuses, or labels instead of PR/issue comments.
|
||||
|
||||
## File Reference
|
||||
|
||||
|
||||
@@ -1,87 +0,0 @@
|
||||
# Profiling model startup
|
||||
|
||||
Measure startup separately from inference latency before choosing what to retain
|
||||
between development runs. The initial instrumentation covers component totals
|
||||
and the native VAE safetensors loader, including Wan. It does not implement a
|
||||
weight donor or CUDA IPC.
|
||||
|
||||
## Run the Wan VAE profile
|
||||
|
||||
Use a prepared CUDA environment and a local Wan 2.1 1.3B Diffusers VAE component directory:
|
||||
|
||||
```bash
|
||||
python examples/benchmark/vae_startup.py \
|
||||
--model-path /path/to/checkpoint/vae \
|
||||
--output /tmp/vae-startup.json \
|
||||
--tensor-output /tmp/vae-output.pt \
|
||||
--synchronize
|
||||
```
|
||||
|
||||
The script measures torch and FastVideo/helper imports separately, device and
|
||||
process-group setup, VAE loading, and the first decode of a fixed seeded latent
|
||||
input. It preserves launcher-provided rendezvous settings; otherwise it chooses a
|
||||
local rendezvous port. It uses one GPU, FP32, and the default cached Wan VAE path.
|
||||
Imports and configuration can themselves initialize CUDA. The
|
||||
`cuda_initialized_before_device_setup` field records whether the explicit
|
||||
device-setup interval begins with an existing context; import time includes
|
||||
any initialization side effects.
|
||||
It verifies output shape and finiteness; those checks alone do not establish
|
||||
parity. Compare the saved tensors against an independent unchanged-source run.
|
||||
|
||||
Run the script in independent processes to measure repeated startup. Record the
|
||||
source revision, checkpoint revision, hardware, runtime versions, cache state,
|
||||
and process wall time alongside the JSON. The script excludes checkpoint
|
||||
downloads. Cached local files do not establish a cold filesystem-cache run.
|
||||
|
||||
## Collect component timings in a worker
|
||||
|
||||
```python
|
||||
from fastvideo.models.loader.startup import collect_startup_timings
|
||||
|
||||
with collect_startup_timings() as timings:
|
||||
pipeline = build_pipeline(...) # Your existing construction call.
|
||||
report = timings.to_dict()
|
||||
```
|
||||
|
||||
Collection is opt-in and local to the current context/process. Install a
|
||||
collector inside each spawned worker to measure its component loads; wrapping
|
||||
only a multiprocess driver does not collect worker timings. Lazy components are
|
||||
recorded when they actually load, if collection is still active.
|
||||
|
||||
Each phase has an ID, parent ID, start offset, duration, and success/error status.
|
||||
Parent durations include child durations: do not sum both. Recording duration
|
||||
also includes gaps between phases. A failed phase propagates its exception and
|
||||
records an error; the collector restores the previous context on exit.
|
||||
|
||||
Default durations measure host wall time, including asynchronous submission.
|
||||
For GPU completion timing, supply a synchronization callback. The example uses a
|
||||
callback that synchronizes only after CUDA is initialized, so the first callback
|
||||
does not move context initialization out of the device-setup interval.
|
||||
Synchronization changes overlap and adds overhead. Compare end-to-end latency
|
||||
with an uninstrumented run; do not interpret instrumented timings as free.
|
||||
|
||||
## Native VAE boundaries
|
||||
|
||||
`VAELoader.load()` retains the existing disk-loading behavior. Internally it now
|
||||
separates two operations:
|
||||
|
||||
- `build_model()` constructs the native module and moves its initial tensors to
|
||||
the target device. The caller supplies the resolved class, architecture config,
|
||||
and dtype context.
|
||||
- `load_weights()` reads safetensors, preserves checkpoint remapping and loading
|
||||
strictness, assigns weights, applies existing overlays/memory-format changes,
|
||||
and switches the module to evaluation mode.
|
||||
|
||||
Phases are `vae.construct`, `vae.device_transfer`, `vae.checkpoint_read`,
|
||||
`vae.weight_assignment`, and `vae.finalize`, nested inside `component.vae` when
|
||||
called through `PipelineComponentLoader`. Checkpoint reads can map files lazily;
|
||||
page faults and transfers can occur during assignment. These labels describe
|
||||
loader operations, not independent physical disk-I/O and PCIe measurements.
|
||||
|
||||
Checkpoint-backed special constructors retain their existing paths and only
|
||||
receive the component-total interval. Transformer/text-encoder internals,
|
||||
compilation, and kernel warm-up need further instrumentation. Reusing GPU tensor
|
||||
storage across processes will require a separate compatibility manifest,
|
||||
ownership/lifetime protocol, and loader implementation. The extracted `load_weights()` method still copies from disk; a future IPC
|
||||
loader must bind storage instead of calling `load_state_dict()` and claiming
|
||||
zero-copy attachment.
|
||||
@@ -1,86 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Profile a local Wan 2.1 1.3B VAE load and first decode on one CUDA GPU.
|
||||
|
||||
Run by filename so import timing begins before importing torch/FastVideo.
|
||||
The checkpoint must already be downloaded; this does not measure Hub downloads.
|
||||
"""
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
import time
|
||||
|
||||
|
||||
def main():
|
||||
started = time.perf_counter()
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model-path", required=True, help="Local Wan 2.1 1.3B Diffusers VAE component directory")
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
parser.add_argument("--tensor-output", type=Path, help="Optional decoded tensor for independent parity checks")
|
||||
parser.add_argument("--synchronize", action="store_true", help="Synchronize CUDA at phase boundaries")
|
||||
args = parser.parse_args()
|
||||
mark = time.perf_counter()
|
||||
import torch
|
||||
torch_import = time.perf_counter() - mark
|
||||
mark = time.perf_counter()
|
||||
import fastvideo.envs as envs
|
||||
from fastvideo.utils import get_open_port
|
||||
from fastvideo.configs.pipelines import PipelineConfig
|
||||
from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.models.loader.component_loader import PipelineComponentLoader
|
||||
from fastvideo.models.loader.startup import collect_startup_timings, startup_phase
|
||||
from fastvideo.models.wan.vae_config import WanVAEConfig
|
||||
fastvideo_import = time.perf_counter() - mark
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("This profile requires one CUDA GPU")
|
||||
mark = time.perf_counter()
|
||||
runtime = FastVideoArgs(model_path=args.model_path,
|
||||
pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision="fp32"))
|
||||
runtime.vae_cpu_offload = False
|
||||
configuration_seconds = time.perf_counter() - mark
|
||||
result = {"imports": {"torch_seconds": torch_import, "fastvideo_and_helpers_seconds": fastvideo_import},
|
||||
"model_path": args.model_path, "torch": torch.__version__,
|
||||
"configuration_seconds": configuration_seconds,
|
||||
"cuda_initialized_before_device_setup": torch.cuda.is_initialized()}
|
||||
distributed_ready = False
|
||||
def synchronize_if_initialized():
|
||||
# Do not move CUDA context initialization out of the device_setup phase.
|
||||
if torch.cuda.is_initialized():
|
||||
torch.cuda.synchronize()
|
||||
try:
|
||||
with collect_startup_timings(synchronize=synchronize_if_initialized if args.synchronize else None) as timings:
|
||||
with startup_phase("device_setup"):
|
||||
envs.setdefault_external("MASTER_ADDR", "127.0.0.1")
|
||||
envs.setdefault_external("MASTER_PORT", str(get_open_port()))
|
||||
maybe_init_distributed_environment_and_model_parallel(1, 1)
|
||||
distributed_ready = True
|
||||
vae = PipelineComponentLoader.load_module("vae", args.model_path, "diffusers", runtime)
|
||||
torch.cuda.synchronize()
|
||||
result["ready_seconds"] = time.perf_counter() - started
|
||||
latents = torch.randn((1, 16, 2, 12, 12), generator=torch.Generator().manual_seed(1024)).cuda()
|
||||
with startup_phase("first_forward"), torch.no_grad():
|
||||
output = vae.decode(latents)
|
||||
torch.cuda.synchronize()
|
||||
assert output.shape == (1, 3, 5, 96, 96) and torch.isfinite(output).all()
|
||||
result["verified_seconds"] = time.perf_counter() - started
|
||||
result["output_shape"] = list(output.shape)
|
||||
result["device"] = torch.cuda.get_device_name(0)
|
||||
if args.tensor_output:
|
||||
args.tensor_output.parent.mkdir(parents=True, exist_ok=True)
|
||||
torch.save(output.cpu(), args.tensor_output)
|
||||
result["status"] = "ok"
|
||||
except Exception as error:
|
||||
result["status"] = "error"
|
||||
result["error"] = repr(error)
|
||||
raise
|
||||
finally:
|
||||
if "timings" in locals():
|
||||
result["startup"] = timings.to_dict()
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.output.write_text(json.dumps(result, indent=2) + "\n")
|
||||
if distributed_ready:
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -39,7 +39,6 @@ from fastvideo.models.loader.text_encoder_quantization import (
|
||||
_process_quantized_text_encoder_weights,
|
||||
_resolve_text_encoder_checkpoint_path,
|
||||
)
|
||||
from fastvideo.models.loader.startup import startup_phase
|
||||
from fastvideo.models.loader.utils import set_default_torch_dtype
|
||||
from fastvideo.models.loader.weight_utils import (
|
||||
filter_duplicate_safetensors_files,
|
||||
@@ -935,7 +934,7 @@ class VAELoader(ComponentLoader):
|
||||
# LTX-2 uses CausalVideoAutoencoder with nested "vae" config
|
||||
if class_name == "CausalVideoAutoencoder" and "vae" in config:
|
||||
vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
|
||||
vae = self.build_model(vae_cls, config, target_device)
|
||||
vae = vae_cls(config).to(target_device)
|
||||
if hasattr(vae, "set_tiling_config"):
|
||||
vae_config = fastvideo_args.pipeline_config.vae_config
|
||||
vae.set_tiling_config(
|
||||
@@ -949,30 +948,8 @@ class VAELoader(ComponentLoader):
|
||||
vae_config = fastvideo_args.pipeline_config.vae_config
|
||||
vae_config.update_model_arch(config)
|
||||
vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
|
||||
vae = self.build_model(vae_cls, vae_config, target_device)
|
||||
vae = vae_cls(vae_config).to(target_device)
|
||||
|
||||
return self.load_weights(vae, model_path, class_name=class_name, config=config, target_device=target_device)
|
||||
|
||||
def build_model(self, vae_cls: type[nn.Module], config: Any, target_device: torch.device) -> nn.Module:
|
||||
"""Construct a native VAE without checkpoint weights; the caller owns the dtype context.
|
||||
|
||||
Checkpoint-backed special constructors in load() retain their existing
|
||||
behavior. This boundary covers the native safetensors path, including Wan.
|
||||
"""
|
||||
with startup_phase("vae.construct"):
|
||||
vae = vae_cls(config)
|
||||
with startup_phase("vae.device_transfer"):
|
||||
vae = vae.to(target_device)
|
||||
return vae
|
||||
|
||||
def load_weights(self, vae: nn.Module, model_path: str, *, class_name: str,
|
||||
config: dict[str, Any], target_device: torch.device) -> nn.Module:
|
||||
"""Hydrate a constructed native VAE from disk and finalize it for inference.
|
||||
|
||||
Keep checkpoint remapping, strictness, overlays and memory-format changes
|
||||
here so construction can be measured/reused independently. This remains a
|
||||
copying disk loader; it does not attach IPC storage.
|
||||
"""
|
||||
# Find all safetensors files
|
||||
safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors"))
|
||||
int8_convrot_path = None
|
||||
@@ -988,9 +965,8 @@ class VAELoader(ComponentLoader):
|
||||
# Common case: a single `.safetensors` checkpoint file.
|
||||
# Some models may be sharded into multiple files; in that case we merge.
|
||||
loaded = {}
|
||||
with startup_phase("vae.checkpoint_read"):
|
||||
for sf_file in safetensors_list:
|
||||
loaded.update(safetensors_load_file(sf_file))
|
||||
for sf_file in safetensors_list:
|
||||
loaded.update(safetensors_load_file(sf_file))
|
||||
|
||||
# LTX-2 CausalVideoAutoencoder needs per_channel_statistics remapping
|
||||
if class_name == "CausalVideoAutoencoder" and "vae" in config:
|
||||
@@ -1019,17 +995,15 @@ class VAELoader(ComponentLoader):
|
||||
# strictly so missing/unexpected keys are surfaced early.
|
||||
# Kandinsky6SRVAE: a partially loaded KVAE decodes plausible-looking garbage, so it is strict as well.
|
||||
strict_load = class_name in {"AutoencoderKL", "AutoencoderKLMiniMaxH3", "Kandinsky6SRVAE"}
|
||||
with startup_phase("vae.weight_assignment"):
|
||||
vae.load_state_dict(loaded, strict=strict_load)
|
||||
with startup_phase("vae.finalize"):
|
||||
if class_name == "AutoencoderKLMiniMaxH3" and int8_convrot_path is not None:
|
||||
from fastvideo.models.vaes.minimax_h3_int8_convrot import overlay_minimax_h3_int8_convrot_decoder
|
||||
overlay_minimax_h3_int8_convrot_decoder(vae, int8_convrot_path)
|
||||
if (class_name == "AutoencoderKLWan" and getattr(vae.config, "use_light_vae", False)
|
||||
and target_device.type == "cuda" and hasattr(vae, "optimize_memory_format")):
|
||||
vae.optimize_memory_format()
|
||||
vae.load_state_dict(loaded, strict=strict_load)
|
||||
if class_name == "AutoencoderKLMiniMaxH3" and int8_convrot_path is not None:
|
||||
from fastvideo.models.vaes.minimax_h3_int8_convrot import overlay_minimax_h3_int8_convrot_decoder
|
||||
overlay_minimax_h3_int8_convrot_decoder(vae, int8_convrot_path)
|
||||
if (class_name == "AutoencoderKLWan" and getattr(vae.config, "use_light_vae", False)
|
||||
and target_device.type == "cuda" and hasattr(vae, "optimize_memory_format")):
|
||||
vae.optimize_memory_format()
|
||||
|
||||
return vae.eval()
|
||||
return vae.eval()
|
||||
|
||||
|
||||
class AudioDecoderLoader(ComponentLoader):
|
||||
@@ -1654,5 +1628,5 @@ class PipelineComponentLoader:
|
||||
attention_context = (_component_attention_backend_scope(resolved, component=module_name)
|
||||
if resolved is not None else nullcontext())
|
||||
|
||||
with startup_phase(f"component.{module_name}"), attention_context:
|
||||
with attention_context:
|
||||
return loader.load(component_model_path, fastvideo_args)
|
||||
|
||||
@@ -1,84 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Opt-in, process-local startup timings without importing or initializing CUDA."""
|
||||
from collections.abc import Callable, Iterator
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass, field
|
||||
import os
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass
|
||||
class StartupTimings:
|
||||
"""Nested wall-clock intervals; parent durations include their children.
|
||||
|
||||
Without a synchronization callback these measure host submission time only.
|
||||
Passing a callback (for example torch.cuda.synchronize) waits at both phase
|
||||
boundaries and changes execution/overlap; use a separate uninstrumented run
|
||||
for an end-to-end latency comparison.
|
||||
"""
|
||||
synchronize: Callable[[], None] | None = field(default=None, repr=False)
|
||||
phases: list[dict[str, Any]] = field(default_factory=list)
|
||||
_start: float = field(default_factory=time.perf_counter, repr=False)
|
||||
_end: float | None = field(default=None, repr=False)
|
||||
_next_id: int = field(default=0, repr=False)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"schema_version": 1,
|
||||
"pid": os.getpid(),
|
||||
"clock": "perf_counter",
|
||||
"timing_mode": "synchronized_boundaries" if self.synchronize else "host",
|
||||
"recording_seconds": (self._end if self._end is not None else time.perf_counter()) - self._start,
|
||||
"phases": sorted(self.phases, key=lambda item: item["id"]),
|
||||
}
|
||||
|
||||
|
||||
_active: ContextVar[StartupTimings | None] = ContextVar("fastvideo_startup_timings", default=None)
|
||||
_stack: ContextVar[tuple[tuple[StartupTimings, int], ...]] = ContextVar("fastvideo_startup_stack", default=())
|
||||
|
||||
|
||||
@contextmanager
|
||||
def collect_startup_timings(*, synchronize: Callable[[], None] | None = None) -> Iterator[StartupTimings]:
|
||||
"""Collect only work in this context/process. Spawned workers need their own collector."""
|
||||
timings = StartupTimings(synchronize=synchronize)
|
||||
token = _active.set(timings)
|
||||
try:
|
||||
yield timings
|
||||
finally:
|
||||
timings._end = time.perf_counter()
|
||||
_active.reset(token)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def startup_phase(name: str) -> Iterator[None]:
|
||||
"""Record a named phase when collection is enabled; otherwise just execute it."""
|
||||
timings = _active.get()
|
||||
if timings is None:
|
||||
yield
|
||||
return
|
||||
if timings.synchronize is not None:
|
||||
timings.synchronize()
|
||||
phase_id = timings._next_id
|
||||
timings._next_id += 1
|
||||
stack = _stack.get()
|
||||
parent = next((index for owner, index in reversed(stack) if owner is timings), None)
|
||||
token = _stack.set((*stack, (timings, phase_id)))
|
||||
start = time.perf_counter()
|
||||
status = "error"
|
||||
try:
|
||||
yield
|
||||
if timings.synchronize is not None:
|
||||
timings.synchronize()
|
||||
status = "ok"
|
||||
finally:
|
||||
timings.phases.append({
|
||||
"id": phase_id,
|
||||
"parent_id": parent,
|
||||
"name": name,
|
||||
"start_seconds": start - timings._start,
|
||||
"duration_seconds": time.perf_counter() - start,
|
||||
"status": status,
|
||||
})
|
||||
_stack.reset(token)
|
||||
@@ -1,64 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
import pytest
|
||||
|
||||
from fastvideo.models.loader.startup import collect_startup_timings, startup_phase
|
||||
|
||||
|
||||
def test_nested_intervals_are_inclusive_and_serializable():
|
||||
import json
|
||||
syncs = []
|
||||
with collect_startup_timings(synchronize=lambda: syncs.append(True)) as timings:
|
||||
with startup_phase("load"):
|
||||
with startup_phase("read"):
|
||||
pass
|
||||
result = timings.to_dict()
|
||||
outer, inner = result["phases"]
|
||||
assert outer["parent_id"] is None
|
||||
assert inner["parent_id"] == outer["id"]
|
||||
assert outer["duration_seconds"] >= inner["duration_seconds"] >= 0
|
||||
assert result["recording_seconds"] >= outer["duration_seconds"]
|
||||
assert result["timing_mode"] == "synchronized_boundaries"
|
||||
assert len(syncs) == 4
|
||||
json.dumps(result)
|
||||
|
||||
|
||||
def test_error_is_recorded_and_collector_is_restored():
|
||||
with collect_startup_timings() as outer:
|
||||
with pytest.raises(ValueError, match="checkpoint"):
|
||||
with collect_startup_timings() as failed:
|
||||
with startup_phase("read"):
|
||||
raise ValueError("checkpoint")
|
||||
with startup_phase("recovery"):
|
||||
pass
|
||||
assert [p["name"] for p in outer.phases] == ["recovery"]
|
||||
assert failed.phases[0]["status"] == "error"
|
||||
assert outer.phases[0]["parent_id"] is None
|
||||
with startup_phase("disabled"):
|
||||
pass
|
||||
assert len(outer.phases) == 1
|
||||
|
||||
|
||||
def test_disabled_collection_does_not_read_the_clock(monkeypatch):
|
||||
from fastvideo.models.loader import startup
|
||||
def unexpected():
|
||||
raise AssertionError("disabled collection measured work")
|
||||
monkeypatch.setattr(startup.time, "perf_counter", unexpected)
|
||||
with startup_phase("disabled"):
|
||||
pass
|
||||
|
||||
|
||||
def test_sync_failure_does_not_leave_a_stale_parent():
|
||||
calls = 0
|
||||
def synchronize():
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
if calls == 2:
|
||||
raise RuntimeError("CUDA failure")
|
||||
with collect_startup_timings(synchronize=synchronize) as timings:
|
||||
with pytest.raises(RuntimeError, match="CUDA failure"):
|
||||
with startup_phase("failed"):
|
||||
pass
|
||||
with startup_phase("next"):
|
||||
pass
|
||||
assert timings.phases[0]["status"] == "error"
|
||||
assert timings.phases[1]["parent_id"] is None
|
||||
@@ -1,82 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Real safetensors hydration remains equivalent to the combined native loader."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from safetensors.torch import save_file
|
||||
|
||||
from fastvideo.models.loader import component_loader
|
||||
from fastvideo.models.loader.component_loader import PipelineComponentLoader, VAELoader
|
||||
from fastvideo.models.loader.startup import collect_startup_timings
|
||||
from fastvideo.models.loader.utils import set_default_torch_dtype
|
||||
|
||||
|
||||
class TinyVAE(torch.nn.Module):
|
||||
def __init__(self, config):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.weight = torch.nn.Parameter(torch.zeros(2))
|
||||
self.register_buffer("scale", torch.zeros(2))
|
||||
|
||||
def forward(self, x):
|
||||
return x * self.weight + self.scale
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sharded", [False, True])
|
||||
def test_native_load_matches_separate_construction_and_hydration(tmp_path, monkeypatch, sharded):
|
||||
config = SimpleNamespace(update_model_arch=lambda _: None)
|
||||
args = SimpleNamespace(
|
||||
model_paths={}, vae_cpu_offload=False,
|
||||
pipeline_config=SimpleNamespace(vae_config=config, vae_precision="fp32"),
|
||||
)
|
||||
monkeypatch.setattr(component_loader, "get_local_torch_device", lambda: torch.device("cpu"))
|
||||
monkeypatch.setattr(component_loader, "get_diffusers_config", lambda **_: {"_class_name": "AutoencoderKLWan"})
|
||||
monkeypatch.setattr(component_loader.ModelRegistry, "resolve_model_cls", lambda _: (TinyVAE, None))
|
||||
tensors = {"weight": torch.tensor([2., 3.]), "scale": torch.tensor([5., 7.])}
|
||||
if sharded:
|
||||
for key, tensor in tensors.items():
|
||||
save_file({key: tensor}, str(tmp_path / f"{key}.safetensors"))
|
||||
else:
|
||||
save_file(tensors, str(tmp_path / "model.safetensors"))
|
||||
loader = VAELoader()
|
||||
with collect_startup_timings() as timings:
|
||||
combined = PipelineComponentLoader.load_module("vae", str(tmp_path), "diffusers", args)
|
||||
with set_default_torch_dtype(torch.float32):
|
||||
separate = loader.build_model(TinyVAE, config, torch.device("cpu"))
|
||||
assert torch.count_nonzero(separate.weight) == 0
|
||||
separate = loader.load_weights(separate, str(tmp_path), class_name="AutoencoderKLWan", config={},
|
||||
target_device=torch.device("cpu"))
|
||||
x = torch.tensor([11., 13.])
|
||||
torch.testing.assert_close(combined(x), torch.tensor([27., 46.]), rtol=0, atol=0)
|
||||
torch.testing.assert_close(combined(x), separate(x), rtol=0, atol=0)
|
||||
assert not combined.training and not separate.training
|
||||
for name, tensor in separate.state_dict().items():
|
||||
torch.testing.assert_close(tensor, tensors[name], rtol=0, atol=0)
|
||||
assert {p["name"] for p in timings.phases} == {
|
||||
"component.vae", "vae.construct", "vae.device_transfer", "vae.checkpoint_read", "vae.weight_assignment",
|
||||
"vae.finalize",
|
||||
}
|
||||
|
||||
root_phase = next(p for p in timings.phases if p["name"] == "component.vae")
|
||||
assert all(p["parent_id"] == root_phase["id"] for p in timings.phases if p is not root_phase)
|
||||
|
||||
|
||||
def test_construction_does_not_read_checkpoint(monkeypatch):
|
||||
def unexpected(*args, **kwargs):
|
||||
raise AssertionError("construction accessed checkpoint")
|
||||
monkeypatch.setattr(component_loader, "safetensors_load_file", unexpected)
|
||||
vae = VAELoader().build_model(TinyVAE, None, torch.device("cpu"))
|
||||
assert torch.count_nonzero(vae.weight) == 0
|
||||
|
||||
|
||||
def test_hydration_preserves_strict_loading_and_missing_checkpoint_errors(tmp_path):
|
||||
loader = VAELoader()
|
||||
vae = TinyVAE(None)
|
||||
with pytest.raises(ValueError, match="No safetensors"):
|
||||
loader.load_weights(vae, str(tmp_path), class_name="AutoencoderKL", config={},
|
||||
target_device=torch.device("cpu"))
|
||||
save_file({"weight": torch.ones(2)}, str(tmp_path / "model.safetensors"))
|
||||
with pytest.raises(RuntimeError, match="scale"):
|
||||
loader.load_weights(vae, str(tmp_path), class_name="AutoencoderKL", config={},
|
||||
target_device=torch.device("cpu"))
|
||||
@@ -255,7 +255,6 @@ nav:
|
||||
- Testing: contributing/testing.md
|
||||
- Environment Variables: contributing/env_vars.md
|
||||
- Performance Benchmarks: contributing/performance_benchmarks.md
|
||||
- Startup Profiling: contributing/startup_profiling.md
|
||||
- Profiling: contributing/profiling.md
|
||||
- Activation Trace: contributing/activation_trace.md
|
||||
- Coding Agents: contributing/coding_agents.md
|
||||
|
||||
Reference in New Issue
Block a user