diff --git a/flashdreams/flashdreams/core/checkpoint/load.py b/flashdreams/flashdreams/core/checkpoint/load.py index 4d11f7642..33da01b0f 100644 --- a/flashdreams/flashdreams/core/checkpoint/load.py +++ b/flashdreams/flashdreams/core/checkpoint/load.py @@ -20,14 +20,16 @@ import io import json import os +import time from collections.abc import Callable, Mapping -from concurrent.futures import ProcessPoolExecutor +from concurrent.futures import ThreadPoolExecutor from typing import Literal, overload from urllib.parse import unquote, urlparse import torch from huggingface_hub import hf_hub_download, try_to_load_from_cache from loguru import logger +from safetensors import safe_open from safetensors.torch import load as load_safetensors from safetensors.torch import load_file as load_safetensors_file from safetensors.torch import save_file as save_safetensors @@ -239,7 +241,7 @@ def _safetensors_device(map_location: str | torch.device) -> str: def _hf_hub_download_shard_task( args: tuple[str, str, str | None, str], ) -> tuple[str, str]: - """Picklable worker: download one shard; used by ProcessPoolExecutor.""" + """Download or resolve one Hugging Face shard.""" repo_id, shard_file, subfolder, revision = args settings: dict[str, object] = { "repo": repo_id, @@ -275,7 +277,7 @@ def _parallel_hf_hub_download_shards( subfolder: str | None, revision: str, ) -> dict[str, str]: - """Download unique shard files in parallel processes; returns shard -> local path.""" + """Download unique shard files in parallel workers; returns shard -> local path.""" if not shard_files: return {} if len(shard_files) == 1: @@ -297,10 +299,10 @@ def _parallel_hf_hub_download_shards( work = [(repo_id, s, subfolder, revision) for s in shard_files] logger.info( f"Downloading {len(shard_files)} Hugging Face safetensors shards " - f"with up to {max_workers} parallel processes" + f"with up to {max_workers} parallel workers" ) shard_to_path: dict[str, str] = {} - with ProcessPoolExecutor(max_workers=max_workers) as pool: + with ThreadPoolExecutor(max_workers=max_workers) as pool: for shard_file, path in pool.map(_hf_hub_download_shard_task, work): shard_to_path[shard_file] = path return shard_to_path @@ -737,12 +739,363 @@ def _load_checkpoint_from_local( ) -> dict[str, torch.Tensor]: """Load checkpoint from local filesystem.""" if ext == ".safetensors": - with open(path, "rb") as f: - return load_safetensors(f.read()) + return load_safetensors_file(path, device=_safetensors_device(map_location)) else: return torch.load(path, map_location=map_location, weights_only=False) +def _copy_checkpoint_tensor(destination: torch.Tensor, source: torch.Tensor) -> int: + """Copy one checkpoint tensor into ``destination`` with bounded staging. + + GPU destinations use a CPU tensor in the destination dtype as the staging + buffer so checkpoint loading never holds a merged state dict or an extra + GPU copy of the tensor. + """ + checkpoint_bytes = source.numel() * source.element_size() + if destination.device.type != "cpu": + staged = source.to(dtype=destination.dtype) + if staged.data_ptr() == source.data_ptr(): + staged = staged.clone() + destination.copy_(staged) + if destination.device.type == "cuda": + # Keep the CPU staging buffer alive until CUDA has consumed it. + torch.cuda.synchronize(destination.device) + del staged + return checkpoint_bytes + + destination.copy_(source.to(device=destination.device, dtype=destination.dtype)) + return checkpoint_bytes + + +def _stream_safetensors_into_model( + model: torch.nn.Module, + path: str, +) -> torch.nn.Module: + """Copy a safetensors checkpoint into a model with bounded host residency. + + Args: + model: Materialized destination model. + path: Local safetensors checkpoint path. + + Returns: + The destination model with checkpoint weights loaded. + + Raises: + RuntimeError: Checkpoint keys or tensor shapes do not match the model. + """ + model_state = model.state_dict() + + with safe_open(path, framework="pt", device="cpu", backend="mmap") as source: + checkpoint_keys = set(source.keys()) + model_keys = set(model_state) + missing = sorted(model_keys - checkpoint_keys) + unexpected = sorted(checkpoint_keys - model_keys) + if missing or unexpected: + details = [] + if missing: + details.append(f"Missing key(s): {', '.join(missing[:20])}") + if unexpected: + details.append(f"Unexpected key(s): {', '.join(unexpected[:20])}") + raise RuntimeError( + f"Checkpoint does not match {type(model).__name__}: " + + "; ".join(details) + ) + + for name, destination in model_state.items(): + source_shape = tuple(source.get_slice(name).get_shape()) + if source_shape != tuple(destination.shape): + raise RuntimeError( + f"Checkpoint tensor {name!r} has shape {source_shape}, " + f"expected {tuple(destination.shape)}" + ) + + with torch.no_grad(): + for name, destination in model_state.items(): + tensor = source.get_tensor(name) + try: + _copy_checkpoint_tensor(destination, tensor) + finally: + del tensor + + return model + + +def _stream_sharded_safetensors_into_model( + model: torch.nn.Module, + *, + weight_map: Mapping[str, str], + resolve_shard_path: Callable[[str], str], +) -> torch.nn.Module: + """Copy a sharded safetensors checkpoint into a model one shard at a time.""" + model_state = model.state_dict() + checkpoint_keys = set(weight_map) + model_keys = set(model_state) + missing = sorted(model_keys - checkpoint_keys) + unexpected = sorted(checkpoint_keys - model_keys) + if missing or unexpected: + details = [] + if missing: + details.append(f"Missing key(s): {', '.join(missing[:20])}") + if unexpected: + details.append(f"Unexpected key(s): {', '.join(unexpected[:20])}") + raise RuntimeError( + f"Checkpoint does not match {type(model).__name__}: " + "; ".join(details) + ) + + keys_by_shard: dict[str, list[str]] = {} + for tensor_name, shard_file in weight_map.items(): + keys_by_shard.setdefault(shard_file, []).append(tensor_name) + + shard_files = sorted(keys_by_shard) + destination_devices = sorted( + {str(tensor.device) for tensor in model_state.values()} + ) + destination_dtypes = sorted({str(tensor.dtype) for tensor in model_state.values()}) + logger.info( + "Streaming sharded safetensors into {}: {} shard(s), {} tensor(s), " + "destination devices={}, dtypes={}", + type(model).__name__, + len(shard_files), + len(weight_map), + destination_devices, + destination_dtypes, + ) + + for shard_index, shard_file in enumerate(shard_files, start=1): + shard_path = resolve_shard_path(shard_file) + tensor_names = keys_by_shard[shard_file] + shard_size_gib = os.path.getsize(shard_path) / 1024**3 + started = time.perf_counter() + logger.info( + "Validating safetensors shard {}/{}: {} tensors, {:.2f} GiB, {}", + shard_index, + len(shard_files), + len(tensor_names), + shard_size_gib, + shard_file, + ) + with safe_open( + shard_path, framework="pt", device="cpu", backend="mmap" + ) as source: + shard_keys = set(source.keys()) + for name in tensor_names: + if name not in shard_keys: + raise KeyError( + f"Key {name!r} missing from shard {shard_file!r} " + f"(path {shard_path!r})" + ) + source_shape = tuple(source.get_slice(name).get_shape()) + destination = model_state[name] + if source_shape != tuple(destination.shape): + raise RuntimeError( + f"Checkpoint tensor {name!r} has shape {source_shape}, " + f"expected {tuple(destination.shape)}" + ) + logger.info( + "Validated safetensors shard {}/{} in {:.1f}s: {}", + shard_index, + len(shard_files), + time.perf_counter() - started, + shard_file, + ) + + total_copied_bytes = 0 + total_started = time.perf_counter() + with torch.no_grad(): + for shard_index, shard_file in enumerate(shard_files, start=1): + shard_path = resolve_shard_path(shard_file) + tensor_names = keys_by_shard[shard_file] + shard_copied_bytes = 0 + started = time.perf_counter() + logger.info( + "Streaming safetensors shard {}/{} into model: {} tensors, {}", + shard_index, + len(shard_files), + len(tensor_names), + shard_file, + ) + with safe_open( + shard_path, framework="pt", device="cpu", backend="mmap" + ) as source: + for name in tensor_names: + tensor = source.get_tensor(name) + try: + tensor_bytes = _copy_checkpoint_tensor( + model_state[name], tensor + ) + shard_copied_bytes += tensor_bytes + total_copied_bytes += tensor_bytes + finally: + del tensor + elapsed = time.perf_counter() - started + throughput = ( + shard_copied_bytes / 1024**3 / elapsed if elapsed > 0 else float("inf") + ) + logger.info( + "Streamed safetensors shard {}/{} in {:.1f}s: {:.2f} GiB copied " + "({:.2f} GiB/s), {}", + shard_index, + len(shard_files), + elapsed, + shard_copied_bytes / 1024**3, + throughput, + shard_file, + ) + + elapsed = time.perf_counter() - total_started + throughput = total_copied_bytes / 1024**3 / elapsed if elapsed > 0 else float("inf") + logger.info( + "Finished streaming {} safetensors shard(s) in {:.1f}s: {:.2f} GiB copied " + "({:.2f} GiB/s)", + len(shard_files), + elapsed, + total_copied_bytes / 1024**3, + throughput, + ) + + return model + + +def _stream_sharded_safetensors_index_into_model( + checkpoint_path: str, + *, + model: torch.nn.Module, + checkpoint_min_free_gb: float | None, +) -> torch.nn.Module | None: + """Stream a safetensors index checkpoint into ``model`` without merging.""" + if checkpoint_path.startswith("s3://"): + return None + + if _is_huggingface_checkpoint_url(checkpoint_path): + repo_id, index_filename, subfolder, revision = ( + _parse_huggingface_checkpoint_url(checkpoint_path) + ) + logger.info( + f"Streaming sharded safetensors checkpoint from Hugging Face: " + f"{checkpoint_path}" + ) + settings: dict[str, object] = { + "repo": repo_id, + "filename": index_filename, + "revision": revision, + } + _preflight_checkpoint_cache_requirement( + label="Hugging Face sharded checkpoint cache", + min_free_gb=checkpoint_min_free_gb, + settings=settings, + ) + min_bytes = _preflight_hf_cache( + label="Hugging Face checkpoint index cache", + settings=settings, + ) + try: + index_local = hf_hub_download( + repo_id=repo_id, + filename=index_filename, + subfolder=subfolder, + revision=revision, + ) + except Exception as exc: + _raise_hf_cache_disk_error( + exc, + label="Hugging Face checkpoint index cache", + required_bytes=min_bytes, + settings=settings, + ) + raise + with open(index_local) as f: + index = json.load(f) + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict) or not weight_map: + raise ValueError( + f"Invalid or empty weight_map in safetensors index: {index_local}" + ) + + unique_shards = sorted(set(weight_map.values())) + shard_to_path = _parallel_hf_hub_download_shards( + repo_id=repo_id, + shard_files=unique_shards, + subfolder=subfolder, + revision=revision, + ) + + def resolve_shard_path(shard_file: str) -> str: + return shard_to_path[shard_file] + + return _stream_sharded_safetensors_into_model( + model, + weight_map=weight_map, + resolve_shard_path=resolve_shard_path, + ) + + if not os.path.isfile(checkpoint_path): + raise FileNotFoundError( + f"Sharded safetensors index not found: {checkpoint_path}" + ) + logger.info( + f"Streaming sharded safetensors checkpoint from local index: {checkpoint_path}" + ) + with open(checkpoint_path) as f: + index = json.load(f) + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict) or not weight_map: + raise ValueError( + f"Invalid or empty weight_map in safetensors index: {checkpoint_path}" + ) + base_dir = os.path.dirname(os.path.abspath(checkpoint_path)) + + def resolve_shard_path(shard_file: str) -> str: + return os.path.join(base_dir, shard_file) + + return _stream_sharded_safetensors_into_model( + model, + weight_map=weight_map, + resolve_shard_path=resolve_shard_path, + ) + + +def _resolve_streamable_safetensors_path( + checkpoint_path: str, + *, + local_cache_dir: str, + checkpoint_min_free_gb: float | None, +) -> str | None: + """Resolve a locally available safetensors file for streaming model loads. + + Args: + checkpoint_path: Local path, S3 URI, or Hugging Face URL. + local_cache_dir: Directory for S3 and merged-safetensors caches. + checkpoint_min_free_gb: Optional Hugging Face cache-space requirement. + + Returns: + Local safetensors path, or ``None`` when materialization is still required. + """ + if _is_sharded_safetensors_index_checkpoint(checkpoint_path): + if checkpoint_path.startswith("s3://"): + return None + cache_path = _sharded_safetensors_merge_cache_path( + checkpoint_path, local_cache_dir + ) + if os.path.exists(cache_path): + logger.info(f"Streaming merged sharded checkpoint from cache: {cache_path}") + return cache_path + return None + + if _get_checkpoint_extension(checkpoint_path) != ".safetensors": + return None + if _is_huggingface_checkpoint_url(checkpoint_path): + return _download_checkpoint_from_huggingface_url( + checkpoint_path, + checkpoint_min_free_gb=checkpoint_min_free_gb, + ) + if checkpoint_path.startswith("s3://"): + cache_path = os.path.join( + local_cache_dir, checkpoint_path.removeprefix("s3://") + ) + return cache_path if os.path.exists(cache_path) else None + return checkpoint_path + + def _load_checkpoint_from_s3( s3_path: str, ext: str, @@ -837,8 +1190,9 @@ def load_checkpoint( Args: checkpoint_path: ``s3://`` URI, local path, or HF URL. Single-file or DCP directory. - model: Model to load weights into. Required for DCP. Optional for - single-file: when provided, ``load_state_dict`` is called. + model: Model to load weights into. Required for DCP. Cached + safetensors are streamed into a provided model; other single-file + formats use ``load_state_dict``. checkpoint_type: ``"auto"``, ``"single"``, or ``"distributed"``. local_cache_dir: Directory for caches. credential_path: S3 credentials path. @@ -873,6 +1227,25 @@ def load_checkpoint( checkpoint_type = "distributed" if checkpoint_type == "single": + if model is not None: + if _is_sharded_safetensors_index_checkpoint(checkpoint_path): + streamed_model = _stream_sharded_safetensors_index_into_model( + checkpoint_path, + model=model, + checkpoint_min_free_gb=checkpoint_min_free_gb, + ) + if streamed_model is not None: + logger.info(f"Streamed checkpoint into model: {checkpoint_path}") + return streamed_model + stream_path = _resolve_streamable_safetensors_path( + checkpoint_path, + local_cache_dir=local_cache_dir, + checkpoint_min_free_gb=checkpoint_min_free_gb, + ) + if stream_path is not None: + _stream_safetensors_into_model(model, stream_path) + logger.info(f"Streamed checkpoint into model: {checkpoint_path}") + return model state_dict = load_single_checkpoint( checkpoint_path=checkpoint_path, local_cache_dir=local_cache_dir, diff --git a/flashdreams/flashdreams/recipes/wan/transformer/wan21.py b/flashdreams/flashdreams/recipes/wan/transformer/wan21.py index c1ae0ab04..0022f1fc2 100644 --- a/flashdreams/flashdreams/recipes/wan/transformer/wan21.py +++ b/flashdreams/flashdreams/recipes/wan/transformer/wan21.py @@ -157,6 +157,16 @@ class Wan21TransformerConfig(TransformerConfig): """Pre-load state-dict remap (e.g. Self-Forcing's ``generator_ema.model.…`` layout).""" + stream_checkpoint: bool = False + """Load cached safetensors directly into the model with bounded host residency.""" + + init_device: str | None = None + """Optional device used for initial network parameter allocation. + + Large streaming-checkpoint models can set this to the final runtime device + so the module is not first materialized as fp32 CPU tensors. + """ + batch_shape: tuple[int, ...] = (1,) """Batch dims of the latent (excluding the L, D dims).""" @@ -273,19 +283,29 @@ def __init__(self, config: Wan21TransformerConfig) -> None: self._output_height: int | None = None self._output_width: int | None = None - self.network = config.network.setup() - self.network = self.network.to(dtype=config.dtype) + self.network = self._setup_network(config) self.network.eval() self.network.set_context_parallel_group(cp_group=self._cp_group) if config.checkpoint_path is not None: - state_dict = load_checkpoint( - config.checkpoint_path, - checkpoint_min_free_gb=config.checkpoint_min_free_gb, - ) - if config.state_dict_transform is not None: - state_dict = config.state_dict_transform(state_dict) - self.network.load_state_dict(state_dict) + if config.stream_checkpoint: + if config.state_dict_transform is not None: + raise ValueError( + "stream_checkpoint does not support state_dict_transform" + ) + load_checkpoint( + config.checkpoint_path, + model=self.network, + checkpoint_min_free_gb=config.checkpoint_min_free_gb, + ) + else: + state_dict = load_checkpoint( + config.checkpoint_path, + checkpoint_min_free_gb=config.checkpoint_min_free_gb, + ) + if config.state_dict_transform is not None: + state_dict = config.state_dict_transform(state_dict) + self.network.load_state_dict(state_dict) self.network.update_parameters_after_loading_checkpoint() if config.compile_network: @@ -308,6 +328,23 @@ def __init__(self, config: Wan21TransformerConfig) -> None: self._cuda_graph_dispatch.uncond_call or self.network ) + @staticmethod + def _setup_network(config: Wan21TransformerConfig) -> WanDiTNetwork: + init_device = ( + None if config.init_device is None else torch.device(config.init_device) + ) + if init_device is None: + return config.network.setup().to(dtype=config.dtype) + + previous_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(config.dtype) + with torch.device(init_device): + network = config.network.setup() + finally: + torch.set_default_dtype(previous_dtype) + return network.to(device=init_device, dtype=config.dtype) + @property def latent_shape(self) -> tuple[int, ...]: """Per-rank post-patchify latent shape ``[*batch_shape, L/cp, D]``. diff --git a/flashdreams/tests/test_checkpoint_loading.py b/flashdreams/tests/test_checkpoint_loading.py new file mode 100644 index 000000000..716e2cd28 --- /dev/null +++ b/flashdreams/tests/test_checkpoint_loading.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Checkpoint loading behavior tests.""" + +import importlib +import json + +import pytest +import torch +from safetensors.torch import save_file as save_safetensors_file + +pytestmark = pytest.mark.ci_cpu + + +def test_local_safetensors_uses_file_backed_loader(monkeypatch, tmp_path) -> None: + """Load local safetensors without materializing the file as bytes.""" + checkpoint_load = importlib.import_module("flashdreams.core.checkpoint.load") + checkpoint_path = tmp_path / "weights.safetensors" + expected = {"weight": torch.ones(2)} + calls: list[tuple[str, str]] = [] + + def fake_load_file(path: str, *, device: str) -> dict[str, torch.Tensor]: + calls.append((path, device)) + return expected + + def reject_bytes_load(_data: bytes) -> dict[str, torch.Tensor]: + pytest.fail("safetensors checkpoints must use the file-backed loader") + + monkeypatch.setattr(checkpoint_load, "load_safetensors_file", fake_load_file) + monkeypatch.setattr(checkpoint_load, "load_safetensors", reject_bytes_load) + + actual = checkpoint_load.load_single_checkpoint( + str(checkpoint_path), map_location=torch.device("cpu") + ) + + assert actual is expected + assert calls == [(str(checkpoint_path), "cpu")] + + +def test_safetensors_model_load_streams_without_full_state_dict( + monkeypatch, tmp_path +) -> None: + """Stream safetensors tensors directly into a materialized model.""" + checkpoint_load = importlib.import_module("flashdreams.core.checkpoint.load") + checkpoint_path = tmp_path / "weights.safetensors" + expected = torch.arange(6, dtype=torch.float32).view(2, 3) + save_safetensors_file({"weight": expected}, checkpoint_path) + model = torch.nn.Linear(3, 2, bias=False) + + def reject_full_load(*_args, **_kwargs) -> None: + pytest.fail("model loads must not materialize the complete state dict") + + monkeypatch.setattr(checkpoint_load, "load_safetensors_file", reject_full_load) + + actual = checkpoint_load.load_checkpoint(str(checkpoint_path), model=model) + + assert actual is model + torch.testing.assert_close(model.weight, expected) + + +def test_sharded_safetensors_model_load_streams_without_merged_state_dict( + monkeypatch, tmp_path +) -> None: + """Stream indexed safetensors shards into a model without merging first.""" + checkpoint_load = importlib.import_module("flashdreams.core.checkpoint.load") + shard_a = tmp_path / "model-00001-of-00002.safetensors" + shard_b = tmp_path / "model-00002-of-00002.safetensors" + index_path = tmp_path / "model.safetensors.index.json" + expected_weight = torch.arange(6, dtype=torch.float32).view(2, 3) + expected_bias = torch.tensor([3.0, 4.0], dtype=torch.float32) + save_safetensors_file({"weight": expected_weight}, shard_a) + save_safetensors_file({"bias": expected_bias}, shard_b) + index_path.write_text( + json.dumps( + { + "metadata": {"total_size": 0}, + "weight_map": { + "weight": shard_a.name, + "bias": shard_b.name, + }, + } + ), + encoding="utf-8", + ) + model = torch.nn.Linear(3, 2) + + def reject_merge(*_args, **_kwargs) -> None: + pytest.fail("sharded model loads must not materialize a merged state dict") + + monkeypatch.setattr( + checkpoint_load, + "_load_sharded_safetensors_index_checkpoint", + reject_merge, + ) + + actual = checkpoint_load.load_checkpoint(str(index_path), model=model) + + assert actual is model + torch.testing.assert_close(model.weight, expected_weight) + torch.testing.assert_close(model.bias, expected_bias) diff --git a/integrations/lingbot/lingbot/config.py b/integrations/lingbot/lingbot/config.py index 72c647f56..ef8530d40 100644 --- a/integrations/lingbot/lingbot/config.py +++ b/integrations/lingbot/lingbot/config.py @@ -76,6 +76,7 @@ in_dim=16 + 4 + 16, ), checkpoint_path=LINGBOT_WORLD_V1_CHECKPOINT_PATH, + stream_checkpoint=True, # Single-rollout layout: tensors flow through the stack as # ``[T, C, H, W]`` (or ``[T, ...]``) with no leading batch/view dim. batch_shape=(), @@ -136,7 +137,8 @@ ) # LingBot-World v2 uses the same architecture and runtime as v1. The -# transformer checkpoint is the only model-level substitution. +# transformer checkpoint is the only model-level substitution; it inherits +# the bounded checkpoint loader from the v1 base config. PIPELINE_LINGBOT_WORLD_V2_14B_CAUSAL_FAST = derive_config( PIPELINE_LINGBOT_WORLD_FAST, name="lingbot-world-v2-14b-causal-fast", diff --git a/integrations/lingbot/lingbot/webrtc/server.py b/integrations/lingbot/lingbot/webrtc/server.py index 4714ed85b..b0c54b248 100644 --- a/integrations/lingbot/lingbot/webrtc/server.py +++ b/integrations/lingbot/lingbot/webrtc/server.py @@ -89,12 +89,14 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--host", type=str, default="0.0.0.0") parser.add_argument("--port", type=int, default=8080) parser.add_argument( + "--config-name", "--config_name", type=str, - default="lingbot-world-fast", + default=LingbotRuntimeConfig().config_name, help="LingBot-World config preset from PIPELINE_CONFIGS.", ) parser.add_argument( + "--no-compile", "--no_compile", action="store_true", help="Disable torch.compile when building the Lingbot pipeline.", @@ -106,12 +108,14 @@ def parse_args() -> argparse.Namespace: help="Torch device used for the Lingbot runtime.", ) parser.add_argument( + "--warmup-chunks", "--warmup_chunks", type=int, default=10, help="Number of synthetic startup chunks to generate for kernel autotuning.", ) parser.add_argument( + "--warmup-timeout-s", "--warmup_timeout_s", type=float, default=600.0, diff --git a/integrations/lingbot/lingbot/webrtc/session.py b/integrations/lingbot/lingbot/webrtc/session.py index c0b59680b..dc7873a96 100644 --- a/integrations/lingbot/lingbot/webrtc/session.py +++ b/integrations/lingbot/lingbot/webrtc/session.py @@ -800,7 +800,10 @@ def _initialize_sync(self) -> None: enable_sync_and_profile=True, diffusion_model=dict( seed=rollout_seed, - transformer=dict(compile_network=self.config.compile_network), + transformer=dict( + compile_network=self.config.compile_network, + init_device=str(self._device), + ), ), ) self._pipeline = pipeline_config.setup().to(device=self._device) diff --git a/integrations/lingbot/tests/test_distributed_server_main.py b/integrations/lingbot/tests/test_distributed_server_main.py index 7f11a776b..b0bbe3be2 100644 --- a/integrations/lingbot/tests/test_distributed_server_main.py +++ b/integrations/lingbot/tests/test_distributed_server_main.py @@ -15,6 +15,7 @@ from __future__ import annotations +import sys from argparse import Namespace import pytest @@ -53,6 +54,27 @@ def _args(device: str = "cuda:0") -> Namespace: ) +def test_parse_args_defaults_to_webrtc_preset_and_accepts_kebab_options( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + sys, + "argv", + [ + "lingbot-webrtc", + "--warmup-chunks", + "0", + "--no-compile", + ], + ) + + args = server.parse_args() + + assert args.config_name == "lingbot-world-fast-taehv-window15-sink3" + assert args.warmup_chunks == 0 + assert args.no_compile is True + + def test_initialize_distributed_single_process_honors_default_device( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/integrations/lingbot/tests/test_smoke.py b/integrations/lingbot/tests/test_smoke.py index bf9c1ab1a..b6bf596a0 100644 --- a/integrations/lingbot/tests/test_smoke.py +++ b/integrations/lingbot/tests/test_smoke.py @@ -195,6 +195,14 @@ def test_lingbot_configs_carry_documented_checkpoint_disk_requirement() -> None: ) +def test_lingbot_configs_enable_streaming_checkpoint_load() -> None: + """Use bounded checkpoint loading for every LingBot model preset.""" + for cfg in RUNNER_CONFIGS.values(): + transformer = cfg.pipeline.diffusion_model.transformer + assert isinstance(transformer, LingbotWorldTransformerConfig) + assert transformer.stream_checkpoint + + def test_v2_only_replaces_the_v1_checkpoint() -> None: """Derive the v2 model by replacing only the v1 checkpoint and slug.""" expected = derive_config(