Compare commits

...
Author SHA1 Message Date
tzhouamandCursor d2c149f25b fix: harden omni weight snapshot downloads
Serialize full-repo Hugging Face snapshot materialization with a repo-wide
node lock and retry incomplete snapshots so CI does not race on shared caches.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-06-04 07:51:54 +00:00
Gao Han f0758f5cbe Merge branch 'main' into feat/diffusion-prefetch-protection 2026-06-03 14:59:14 +08:00
Gao Han 69fd466fba Merge branch 'main' into feat/diffusion-prefetch-protection 2026-06-02 23:22:08 +08:00
tzhouam cb08af976b fix: correct logging message format in hub_prefetch.py for better clarity during prefetch attempts 2026-06-02 09:01:47 +00:00
tzhouam 617c6ca55f feat: enhance model loading with prefetching and retry logic
- Introduced `from_pretrained_with_prefetch` to handle racy cache scenarios by re-prefetching and retrying on failures.
- Updated various model pipelines to utilize the new prefetching mechanism, ensuring robust loading of model components.
- Prefetch logic added to multiple models, including HunyuanVideo, LongCatImage, and QwenImage, to mitigate issues with incomplete cache states.
2026-06-02 09:01:47 +00:00
21 changed files with 844 additions and 171 deletions

No files matched your search

+128 -14
View File
@@ -65,10 +65,22 @@ import errno
import logging import logging
import os import os
import time import time
from collections.abc import Iterable, Iterator from collections.abc import Callable, Iterable, Iterator
from typing import Any
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# A racy / partially-evicted HF cache (the exact failure this module defends
# against) is transient: re-running ``snapshot_download`` blocks on the peer
# writer's per-blob ``.lock`` and then returns a complete tree. So a bounded
# retry with linear backoff is what actually closes the window that a single
# best-effort attempt left open (Buildkite vllm-omni-rebase #1858: both the
# ``cuda_ti2v_hsdp`` missing-shard ``OSError`` and the ``wan_2_1_vace`` default
# ``UMT5Config`` size-mismatch were a swallowed prefetch followed by a
# ``from_pretrained`` against the half-written cache).
_PREFETCH_MAX_ATTEMPTS = 3
_PREFETCH_BACKOFF_BASE_S = 1.0
def _node_lock_dir() -> str: def _node_lock_dir() -> str:
"""Return a node-local directory suitable for the prefetch lock files. """Return a node-local directory suitable for the prefetch lock files.
@@ -276,7 +288,11 @@ def prefetch_subfolders(
if local_files_only or not model or os.path.isdir(model): if local_files_only or not model or os.path.isdir(model):
return return
logger.info("Prefetching %s subfolders: %s", model, list(subfolders)) # Materialise ``subfolders`` up-front: it may be a one-shot generator and
# we reference it again in the retry / logging paths below.
subfolders = list(subfolders)
logger.info("Prefetching %s subfolders: %s", model, subfolders)
try: try:
from huggingface_hub import snapshot_download from huggingface_hub import snapshot_download
@@ -306,6 +322,14 @@ def prefetch_subfolders(
# snapshot is now warm. This is what makes the prefetch race-free # snapshot is now warm. This is what makes the prefetch race-free
# even when many ``DiffusionWorker`` subprocesses (or multiple # even when many ``DiffusionWorker`` subprocesses (or multiple
# OmniServer instances on the same node) hit this code in parallel. # OmniServer instances on the same node) hit this code in parallel.
# A single best-effort attempt is not enough: when several diffusion
# workers race a cold cache (and the node-wide lock fails to serialise
# them, as observed for HSDP / ring launches), ``snapshot_download`` can
# raise on a half-written tree. Swallowing that and proceeding straight
# to ``from_pretrained`` is exactly what turned a recoverable prefetch
# hiccup into a hard server crash. Retry with backoff so the snapshot
# actually completes before any loader reads the cache.
for attempt in range(1, _PREFETCH_MAX_ATTEMPTS + 1):
try: try:
with _repo_prefetch_lock(model): with _repo_prefetch_lock(model):
snapshot_download( snapshot_download(
@@ -313,15 +337,12 @@ def prefetch_subfolders(
allow_patterns=allow_patterns, allow_patterns=allow_patterns,
) )
logger.info("Prefetch complete for %s", model) logger.info("Prefetch complete for %s", model)
return
except Exception as exc: except Exception as exc:
# Best-effort: propagate only via logging. The subsequent # Auth / gating never heals on retry - escalate immediately with
# ``from_pretrained`` call will raise a clearer, call-site-specific # an explicit hint so readers of CI logs don't have to correlate
# error (auth, 404, disk full, ...) that we'd rather surface - EXCEPT # the generic "OSError: <repo> does not appear to have a file
# for auth/gating, which we escalate here with an explicit hint so # named ..." that ``from_pretrained`` would otherwise emit later.
# readers of CI logs don't have to correlate the generic "OSError:
# <repo> does not appear to have a file named ..." that
# ``from_pretrained`` would otherwise emit much later with an
# unrelated-looking message.
if _looks_like_auth_error(exc): if _looks_like_auth_error(exc):
logger.error( logger.error(
"Hub prefetch for '%s' failed with an authentication / gated " "Hub prefetch for '%s' failed with an authentication / gated "
@@ -333,12 +354,33 @@ def prefetch_subfolders(
type(exc).__name__, type(exc).__name__,
exc, exc,
) )
else: return
if attempt < _PREFETCH_MAX_ATTEMPTS:
backoff = _PREFETCH_BACKOFF_BASE_S * attempt
logger.warning( logger.warning(
"Hub prefetch for repo '%s' subfolders %s failed (%s: %s); " "Hub prefetch for repo '%s' subfolders %s failed on attempt %d/%d (%s: %s); retrying in %.1fs",
"falling back to on-demand download in from_pretrained",
model, model,
list(subfolders), subfolders,
attempt,
_PREFETCH_MAX_ATTEMPTS,
type(exc).__name__,
exc,
backoff,
)
time.sleep(backoff)
continue
# Exhausted retries. Stay best-effort: propagate only via logging
# so the subsequent ``from_pretrained`` call surfaces the real,
# call-site-specific error (and ``from_pretrained_with_prefetch``
# gets a final chance to heal the cache).
logger.warning(
"Hub prefetch for repo '%s' subfolders %s failed after %d attempts "
"(%s: %s); falling back to on-demand download in from_pretrained",
model,
subfolders,
_PREFETCH_MAX_ATTEMPTS,
type(exc).__name__, type(exc).__name__,
exc, exc,
) )
@@ -376,6 +418,78 @@ def _looks_like_auth_error(exc: BaseException) -> bool:
return "401 client error" in msg or "403 client error" in msg or "gatedrepo" in msg return "401 client error" in msg or "403 client error" in msg or "gatedrepo" in msg
def from_pretrained_with_prefetch(
factory: Callable[..., Any],
model: str,
*,
subfolder: str,
prefetch_list: Iterable[str],
local_files_only: bool = False,
max_attempts: int = _PREFETCH_MAX_ATTEMPTS,
**from_pretrained_kwargs: Any,
) -> Any:
"""Call ``factory.from_pretrained`` healing a racy / partial HF cache.
``factory`` is a bound ``SomeModel.from_pretrained`` (or any callable with
the same ``(model, *, subfolder, local_files_only, **kwargs)`` signature).
This is a stronger sibling of :func:`retry_on_missing_shard`: that helper
only retries the missing-shard ``OSError`` and never re-prefetches, so it
cannot recover the second face of the same race. Two shapes of partial
-cache failure crash the diffusion server outright:
* ``OSError: <repo> does not appear to have a file named
text_encoder/model-0000X-of-0000Y.safetensors`` - a shard is still under
its ``.incomplete`` name.
* ``RuntimeError: You set 'ignore_mismatched_sizes' to 'False' ...`` -
``text_encoder/config.json`` was not present yet, so ``transformers`` v5
silently fell back to the default (tiny) config and then could not load
the real checkpoint into it.
Both heal once the cache is complete. So on those errors we re-run a
*verified* prefetch (which blocks on the peer writer and retries the
download) and reload, instead of letting the worker die. Local paths and
``local_files_only`` loads cannot be healed by re-fetching, so they raise
on the first failure exactly as before.
"""
prefetch_list = list(prefetch_list)
can_heal = not local_files_only and bool(model) and not os.path.isdir(model)
last_exc: BaseException | None = None
for attempt in range(1, max_attempts + 1):
try:
return factory(
model,
subfolder=subfolder,
local_files_only=local_files_only,
**from_pretrained_kwargs,
)
except (OSError, RuntimeError, ValueError) as exc:
last_exc = exc
if not can_heal or attempt >= max_attempts:
break
backoff = _PREFETCH_BACKOFF_BASE_S * attempt
logger.warning(
"from_pretrained(%s, subfolder=%s) failed on attempt %d/%d "
"(%s: %s); re-prefetching repo and retrying in %.1fs",
model,
subfolder,
attempt,
max_attempts,
type(exc).__name__,
exc,
backoff,
)
time.sleep(backoff)
# Force a fresh, verified snapshot of every component this pipeline
# needs - not just ``subfolder`` - so a sibling component that was
# also half-written gets repaired in the same pass.
prefetch_subfolders(model, prefetch_list, local_files_only=False)
assert last_exc is not None # loop only exits via return or a caught exc
raise last_exc
def retry_on_missing_shard(load_fn, *, max_retries: int = 3, base_delay: float = 5.0): def retry_on_missing_shard(load_fn, *, max_retries: int = 3, base_delay: float = 5.0):
"""Call *load_fn* with retry on the transformers v5 shard-resolution race. """Call *load_fn* with retry on the transformers v5 shard-resolution race.
@@ -39,7 +39,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import ( from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
Flux2Transformer2DModel, Flux2Transformer2DModel,
) )
@@ -213,9 +213,10 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
# Avoid the transformers v5 multi-worker subfolder race (see # Avoid the transformers v5 multi-worker subfolder race (see
# ``vllm_omni/diffusion/model_loader/hub_prefetch.py`` for the full # ``vllm_omni/diffusion/model_loader/hub_prefetch.py`` for the full
# analysis; L4 build #1043 hit this on FLUX.2-klein-4B's text_encoder). # analysis; L4 build #1043 hit this on FLUX.2-klein-4B's text_encoder).
flux2_subfolders = ["scheduler", "text_encoder", "tokenizer", "vae"]
prefetch_subfolders( prefetch_subfolders(
model, model,
["scheduler", "text_encoder", "tokenizer", "vae"], flux2_subfolders,
local_files_only=local_files_only, local_files_only=local_files_only,
) )
@@ -224,9 +225,15 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
subfolder="scheduler", subfolder="scheduler",
local_files_only=local_files_only, local_files_only=local_files_only,
) )
self.text_encoder = Qwen3ForCausalLM.from_pretrained( # ``from_pretrained_with_prefetch`` re-prefetches and retries if a peer
# worker left the cache half-written (missing-shard ``OSError`` or the
# default-config size-mismatch ``RuntimeError``) instead of crashing
# the worker - FLUX.2-klein-4B's sharded text_encoder hit this on #1043.
self.text_encoder = from_pretrained_with_prefetch(
Qwen3ForCausalLM.from_pretrained,
model, model,
subfolder="text_encoder", subfolder="text_encoder",
prefetch_list=flux2_subfolders,
local_files_only=local_files_only, local_files_only=local_files_only,
).to(self._execution_device) ).to(self._execution_device)
self.tokenizer = Qwen2TokenizerFast.from_pretrained( self.tokenizer = Qwen2TokenizerFast.from_pretrained(
@@ -234,9 +241,11 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
subfolder="tokenizer", subfolder="tokenizer",
local_files_only=local_files_only, local_files_only=local_files_only,
) )
self.vae = AutoencoderKLFlux2.from_pretrained( self.vae = from_pretrained_with_prefetch(
AutoencoderKLFlux2.from_pretrained,
model, model,
subfolder="vae", subfolder="vae",
prefetch_list=flux2_subfolders,
local_files_only=local_files_only, local_files_only=local_files_only,
).to(self._execution_device) ).to(self._execution_device)
@@ -31,6 +31,7 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size
from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.models.hidream_image import HiDreamImageTransformer2DModel from vllm_omni.diffusion.models.hidream_image import HiDreamImageTransformer2DModel
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
@@ -177,24 +178,55 @@ class HiDreamI1ImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfi
# Check if model is a local path # Check if model is a local path
local_files_only = os.path.exists(model) local_files_only = os.path.exists(model)
# See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
# race; prefetch the in-repo component set before any from_pretrained
# (``text_encoder_4`` lives in a separate Llama repo and is unaffected).
hidream_subfolders = [
"scheduler",
"vae",
"text_encoder",
"tokenizer",
"text_encoder_2",
"tokenizer_2",
"text_encoder_3",
"tokenizer_3",
]
prefetch_subfolders(model, hidream_subfolders, local_files_only=local_files_only)
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
model, subfolder="scheduler", local_files_only=local_files_only model, subfolder="scheduler", local_files_only=local_files_only
) )
self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.vae = from_pretrained_with_prefetch(
self.device AutoencoderKL.from_pretrained,
) model,
self.text_encoder = CLIPTextModelWithProjection.from_pretrained( subfolder="vae",
model, subfolder="text_encoder", local_files_only=local_files_only prefetch_list=hidream_subfolders,
local_files_only=local_files_only,
).to(self.device)
self.text_encoder = from_pretrained_with_prefetch(
CLIPTextModelWithProjection.from_pretrained,
model,
subfolder="text_encoder",
prefetch_list=hidream_subfolders,
local_files_only=local_files_only,
) )
self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
self.text_encoder_2 = CLIPTextModelWithProjection.from_pretrained( self.text_encoder_2 = from_pretrained_with_prefetch(
model, subfolder="text_encoder_2", local_files_only=local_files_only CLIPTextModelWithProjection.from_pretrained,
model,
subfolder="text_encoder_2",
prefetch_list=hidream_subfolders,
local_files_only=local_files_only,
) )
self.tokenizer_2 = CLIPTokenizer.from_pretrained( self.tokenizer_2 = CLIPTokenizer.from_pretrained(
model, subfolder="tokenizer_2", local_files_only=local_files_only model, subfolder="tokenizer_2", local_files_only=local_files_only
) )
self.text_encoder_3 = T5EncoderModel.from_pretrained( self.text_encoder_3 = from_pretrained_with_prefetch(
model, subfolder="text_encoder_3", local_files_only=local_files_only T5EncoderModel.from_pretrained,
model,
subfolder="text_encoder_3",
prefetch_list=hidream_subfolders,
local_files_only=local_files_only,
) )
self.tokenizer_3 = T5Tokenizer.from_pretrained( self.tokenizer_3 = T5Tokenizer.from_pretrained(
model, subfolder="tokenizer_3", local_files_only=local_files_only model, subfolder="tokenizer_3", local_files_only=local_files_only
@@ -31,6 +31,7 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size
from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.models.hidream_image import HiDreamImageTransformer2DModel from vllm_omni.diffusion.models.hidream_image import HiDreamImageTransformer2DModel
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
@@ -177,24 +178,55 @@ class HiDreamImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfile
# Check if model is a local path # Check if model is a local path
local_files_only = os.path.exists(model) local_files_only = os.path.exists(model)
# See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
# race; prefetch the in-repo component set before any from_pretrained
# (``text_encoder_4`` lives in a separate Llama repo and is unaffected).
hidream_subfolders = [
"scheduler",
"vae",
"text_encoder",
"tokenizer",
"text_encoder_2",
"tokenizer_2",
"text_encoder_3",
"tokenizer_3",
]
prefetch_subfolders(model, hidream_subfolders, local_files_only=local_files_only)
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
model, subfolder="scheduler", local_files_only=local_files_only model, subfolder="scheduler", local_files_only=local_files_only
) )
self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.vae = from_pretrained_with_prefetch(
self.device AutoencoderKL.from_pretrained,
) model,
self.text_encoder = CLIPTextModelWithProjection.from_pretrained( subfolder="vae",
model, subfolder="text_encoder", local_files_only=local_files_only prefetch_list=hidream_subfolders,
local_files_only=local_files_only,
).to(self.device)
self.text_encoder = from_pretrained_with_prefetch(
CLIPTextModelWithProjection.from_pretrained,
model,
subfolder="text_encoder",
prefetch_list=hidream_subfolders,
local_files_only=local_files_only,
) )
self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
self.text_encoder_2 = CLIPTextModelWithProjection.from_pretrained( self.text_encoder_2 = from_pretrained_with_prefetch(
model, subfolder="text_encoder_2", local_files_only=local_files_only CLIPTextModelWithProjection.from_pretrained,
model,
subfolder="text_encoder_2",
prefetch_list=hidream_subfolders,
local_files_only=local_files_only,
) )
self.tokenizer_2 = CLIPTokenizer.from_pretrained( self.tokenizer_2 = CLIPTokenizer.from_pretrained(
model, subfolder="tokenizer_2", local_files_only=local_files_only model, subfolder="tokenizer_2", local_files_only=local_files_only
) )
self.text_encoder_3 = T5EncoderModel.from_pretrained( self.text_encoder_3 = from_pretrained_with_prefetch(
model, subfolder="text_encoder_3", local_files_only=local_files_only T5EncoderModel.from_pretrained,
model,
subfolder="text_encoder_3",
prefetch_list=hidream_subfolders,
local_files_only=local_files_only,
) )
self.tokenizer_3 = T5Tokenizer.from_pretrained( self.tokenizer_3 = T5Tokenizer.from_pretrained(
model, subfolder="tokenizer_3", local_files_only=local_files_only model, subfolder="tokenizer_3", local_files_only=local_files_only
@@ -25,6 +25,7 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan_video_1
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.models.hunyuan_video.hunyuan_video_15_transformer import HunyuanVideo15Transformer3DModel from vllm_omni.diffusion.models.hunyuan_video.hunyuan_video_15_transformer import HunyuanVideo15Transformer3DModel
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel
@@ -101,9 +102,20 @@ class HunyuanVideo15Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, Diff
model = od_config.model model = od_config.model
local_files_only = os.path.exists(model) local_files_only = os.path.exists(model)
# See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
# race; prefetch the whole component set before any from_pretrained so
# peer workers load from a warm, complete cache.
hv_subfolders = ["tokenizer", "text_encoder", "tokenizer_2", "text_encoder_2", "vae", "scheduler"]
prefetch_subfolders(model, hv_subfolders, local_files_only=local_files_only)
self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
self.text_encoder = Qwen2_5_VLTextModel.from_pretrained( self.text_encoder = from_pretrained_with_prefetch(
model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only Qwen2_5_VLTextModel.from_pretrained,
model,
subfolder="text_encoder",
prefetch_list=hv_subfolders,
local_files_only=local_files_only,
torch_dtype=dtype,
).to(self.device) ).to(self.device)
self.tokenizer_2 = ByT5Tokenizer.from_pretrained( self.tokenizer_2 = ByT5Tokenizer.from_pretrained(
@@ -112,8 +124,13 @@ class HunyuanVideo15Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, Diff
t5_config = AutoConfig.from_pretrained(model, subfolder="text_encoder_2", local_files_only=local_files_only) t5_config = AutoConfig.from_pretrained(model, subfolder="text_encoder_2", local_files_only=local_files_only)
self.text_encoder_2 = T5EncoderModel(t5_config, prefix="text_encoder_2").to(dtype=dtype, device=self.device) self.text_encoder_2 = T5EncoderModel(t5_config, prefix="text_encoder_2").to(dtype=dtype, device=self.device)
self.vae = DistributedAutoencoderKLHunyuanVideo15.from_pretrained( self.vae = from_pretrained_with_prefetch(
model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only DistributedAutoencoderKLHunyuanVideo15.from_pretrained,
model,
subfolder="vae",
prefetch_list=hv_subfolders,
local_files_only=local_files_only,
torch_dtype=torch.float32,
).to(self.device) ).to(self.device)
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
@@ -32,6 +32,7 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan_video_1
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.models.hunyuan_video.hunyuan_video_15_transformer import HunyuanVideo15Transformer3DModel from vllm_omni.diffusion.models.hunyuan_video.hunyuan_video_15_transformer import HunyuanVideo15Transformer3DModel
from vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5 import ( from vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5 import (
extract_glyph_texts, extract_glyph_texts,
@@ -123,9 +124,29 @@ class HunyuanVideo15I2VPipeline(
model = od_config.model model = od_config.model
local_files_only = os.path.exists(model) local_files_only = os.path.exists(model)
# See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
# race; prefetch the whole component set before any from_pretrained so
# peer workers load from a warm, complete cache.
hv_subfolders = [
"tokenizer",
"text_encoder",
"tokenizer_2",
"text_encoder_2",
"image_encoder",
"feature_extractor",
"vae",
"scheduler",
]
prefetch_subfolders(model, hv_subfolders, local_files_only=local_files_only)
self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
self.text_encoder = Qwen2_5_VLTextModel.from_pretrained( self.text_encoder = from_pretrained_with_prefetch(
model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only Qwen2_5_VLTextModel.from_pretrained,
model,
subfolder="text_encoder",
prefetch_list=hv_subfolders,
local_files_only=local_files_only,
torch_dtype=dtype,
).to(self.device) ).to(self.device)
self.tokenizer_2 = ByT5Tokenizer.from_pretrained( self.tokenizer_2 = ByT5Tokenizer.from_pretrained(
@@ -134,15 +155,25 @@ class HunyuanVideo15I2VPipeline(
t5_config = AutoConfig.from_pretrained(model, subfolder="text_encoder_2", local_files_only=local_files_only) t5_config = AutoConfig.from_pretrained(model, subfolder="text_encoder_2", local_files_only=local_files_only)
self.text_encoder_2 = T5EncoderModel(t5_config, prefix="text_encoder_2").to(dtype=dtype, device=self.device) self.text_encoder_2 = T5EncoderModel(t5_config, prefix="text_encoder_2").to(dtype=dtype, device=self.device)
self.image_encoder = SiglipVisionModel.from_pretrained( self.image_encoder = from_pretrained_with_prefetch(
model, subfolder="image_encoder", torch_dtype=dtype, local_files_only=local_files_only SiglipVisionModel.from_pretrained,
model,
subfolder="image_encoder",
prefetch_list=hv_subfolders,
local_files_only=local_files_only,
torch_dtype=dtype,
).to(self.device) ).to(self.device)
self.feature_extractor = SiglipImageProcessor.from_pretrained( self.feature_extractor = SiglipImageProcessor.from_pretrained(
model, subfolder="feature_extractor", local_files_only=local_files_only model, subfolder="feature_extractor", local_files_only=local_files_only
) )
self.vae = DistributedAutoencoderKLHunyuanVideo15.from_pretrained( self.vae = from_pretrained_with_prefetch(
model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only DistributedAutoencoderKLHunyuanVideo15.from_pretrained,
model,
subfolder="vae",
prefetch_list=hv_subfolders,
local_files_only=local_files_only,
torch_dtype=torch.float32,
).to(self.device) ).to(self.device)
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
@@ -27,6 +27,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import LongCatImageTransformer2DModel from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import LongCatImageTransformer2DModel
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -223,19 +224,32 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfile
model = od_config.model model = od_config.model
local_files_only = os.path.exists(model) local_files_only = os.path.exists(model)
# See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
# race; prefetch the whole component set before any from_pretrained.
longcat_subfolders = ["scheduler", "text_encoder", "tokenizer", "vae"]
prefetch_subfolders(model, longcat_subfolders, local_files_only=local_files_only)
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
model, subfolder="scheduler", local_files_only=local_files_only model, subfolder="scheduler", local_files_only=local_files_only
) )
self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( self.text_encoder = from_pretrained_with_prefetch(
model, subfolder="text_encoder", local_files_only=local_files_only Qwen2_5_VLForConditionalGeneration.from_pretrained,
model,
subfolder="text_encoder",
prefetch_list=longcat_subfolders,
local_files_only=local_files_only,
) )
self.text_processor = Qwen2VLProcessor.from_pretrained( self.text_processor = Qwen2VLProcessor.from_pretrained(
model, subfolder="tokenizer", local_files_only=local_files_only model, subfolder="tokenizer", local_files_only=local_files_only
) )
self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.vae = from_pretrained_with_prefetch(
self.device AutoencoderKL.from_pretrained,
) model,
subfolder="vae",
prefetch_list=longcat_subfolders,
local_files_only=local_files_only,
).to(self.device)
self.transformer = LongCatImageTransformer2DModel(od_config=od_config) self.transformer = LongCatImageTransformer2DModel(od_config=od_config)
self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
@@ -29,6 +29,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.models.interface import SupportImageInput from vllm_omni.diffusion.models.interface import SupportImageInput
from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import ( from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import (
LongCatImageTransformer2DModel, LongCatImageTransformer2DModel,
@@ -244,19 +245,32 @@ class LongCatImageEditPipeline(nn.Module, CFGParallelMixin, SupportImageInput, D
model = od_config.model model = od_config.model
local_files_only = os.path.exists(model) local_files_only = os.path.exists(model)
# See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
# race; prefetch the whole component set before any from_pretrained.
longcat_subfolders = ["scheduler", "text_encoder", "text_processor", "tokenizer", "vae"]
prefetch_subfolders(model, longcat_subfolders, local_files_only=local_files_only)
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
model, subfolder="scheduler", local_files_only=local_files_only model, subfolder="scheduler", local_files_only=local_files_only
) )
self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( self.text_encoder = from_pretrained_with_prefetch(
model, subfolder="text_encoder", local_files_only=local_files_only Qwen2_5_VLForConditionalGeneration.from_pretrained,
model,
subfolder="text_encoder",
prefetch_list=longcat_subfolders,
local_files_only=local_files_only,
) )
self.text_processor = Qwen2VLProcessor.from_pretrained( self.text_processor = Qwen2VLProcessor.from_pretrained(
model, subfolder="text_processor", local_files_only=local_files_only model, subfolder="text_processor", local_files_only=local_files_only
) )
self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.vae = from_pretrained_with_prefetch(
self.device AutoencoderKL.from_pretrained,
) model,
subfolder="vae",
prefetch_list=longcat_subfolders,
local_files_only=local_files_only,
).to(self.device)
self.transformer = LongCatImageTransformer2DModel(od_config=od_config) self.transformer = LongCatImageTransformer2DModel(od_config=od_config)
self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
@@ -36,6 +36,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
@@ -185,6 +186,19 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
), ),
] ]
# See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
# race; prefetch the whole component set before any from_pretrained.
ltx2_subfolders = [
"tokenizer",
"text_encoder",
"connectors",
"vae",
"audio_vae",
"vocoder",
"scheduler",
]
prefetch_subfolders(model, ltx2_subfolders, local_files_only=local_files_only)
self.tokenizer = AutoTokenizer.from_pretrained( self.tokenizer = AutoTokenizer.from_pretrained(
model, model,
subfolder="tokenizer", subfolder="tokenizer",
@@ -193,36 +207,46 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
# prefer mmap loading as default device is cuda, and the output of text encoder # prefer mmap loading as default device is cuda, and the output of text encoder
# could be deterministic. # could be deterministic.
with torch.device("cpu"): with torch.device("cpu"):
self.text_encoder = Gemma3ForConditionalGeneration.from_pretrained( self.text_encoder = from_pretrained_with_prefetch(
Gemma3ForConditionalGeneration.from_pretrained,
model, model,
subfolder="text_encoder", subfolder="text_encoder",
torch_dtype=dtype, prefetch_list=ltx2_subfolders,
local_files_only=local_files_only, local_files_only=local_files_only,
torch_dtype=dtype,
).to(self.device) ).to(self.device)
self.connectors = LTX2TextConnectors.from_pretrained( self.connectors = from_pretrained_with_prefetch(
LTX2TextConnectors.from_pretrained,
model, model,
subfolder="connectors", subfolder="connectors",
torch_dtype=dtype, prefetch_list=ltx2_subfolders,
local_files_only=local_files_only, local_files_only=local_files_only,
torch_dtype=dtype,
).to(self.device) ).to(self.device)
self.vae = AutoencoderKLLTX2Video.from_pretrained( self.vae = from_pretrained_with_prefetch(
AutoencoderKLLTX2Video.from_pretrained,
model, model,
subfolder="vae", subfolder="vae",
torch_dtype=dtype, prefetch_list=ltx2_subfolders,
local_files_only=local_files_only, local_files_only=local_files_only,
torch_dtype=dtype,
).to(self.device) ).to(self.device)
self.audio_vae = AutoencoderKLLTX2Audio.from_pretrained( self.audio_vae = from_pretrained_with_prefetch(
AutoencoderKLLTX2Audio.from_pretrained,
model, model,
subfolder="audio_vae", subfolder="audio_vae",
torch_dtype=dtype, prefetch_list=ltx2_subfolders,
local_files_only=local_files_only, local_files_only=local_files_only,
torch_dtype=dtype,
).to(self.device) ).to(self.device)
self.vocoder = LTX2Vocoder.from_pretrained( self.vocoder = from_pretrained_with_prefetch(
LTX2Vocoder.from_pretrained,
model, model,
subfolder="vocoder", subfolder="vocoder",
torch_dtype=dtype, prefetch_list=ltx2_subfolders,
local_files_only=local_files_only, local_files_only=local_files_only,
torch_dtype=dtype,
).to(self.device) ).to(self.device)
transformer_config = load_transformer_config(model, "transformer", local_files_only) transformer_config = load_transformer_config(model, "transformer", local_files_only)
@@ -45,6 +45,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
) )
from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -149,26 +150,59 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
), ),
] ]
# See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
# race; prefetch the whole component set before any from_pretrained.
ltx2_subfolders = [
"tokenizer",
"text_encoder",
"connectors",
"vae",
"audio_vae",
"vocoder",
"scheduler",
]
prefetch_subfolders(model, ltx2_subfolders, local_files_only=local_files_only)
# --- Tokenizer (lightweight, stays wherever) --- # --- Tokenizer (lightweight, stays wherever) ---
self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
# --- Text encoder: load on CPU, move to GPU only during encoding --- # --- Text encoder: load on CPU, move to GPU only during encoding ---
with torch.device("cpu"): with torch.device("cpu"):
self.text_encoder = Gemma3ForConditionalGeneration.from_pretrained( self.text_encoder = from_pretrained_with_prefetch(
model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only Gemma3ForConditionalGeneration.from_pretrained,
model,
subfolder="text_encoder",
prefetch_list=ltx2_subfolders,
local_files_only=local_files_only,
torch_dtype=dtype,
) )
# --- Connectors: CPU (LTX-2.3 connectors include caption projection) --- # --- Connectors: CPU (LTX-2.3 connectors include caption projection) ---
self.connectors = LTX2TextConnectors.from_pretrained( self.connectors = from_pretrained_with_prefetch(
model, subfolder="connectors", torch_dtype=dtype, local_files_only=local_files_only LTX2TextConnectors.from_pretrained,
model,
subfolder="connectors",
prefetch_list=ltx2_subfolders,
local_files_only=local_files_only,
torch_dtype=dtype,
) )
# --- VAE, Audio VAE: CPU --- # --- VAE, Audio VAE: CPU ---
self.vae = AutoencoderKLLTX2Video.from_pretrained( self.vae = from_pretrained_with_prefetch(
model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only AutoencoderKLLTX2Video.from_pretrained,
model,
subfolder="vae",
prefetch_list=ltx2_subfolders,
local_files_only=local_files_only,
torch_dtype=dtype,
) )
self.audio_vae = AutoencoderKLLTX2Audio.from_pretrained( self.audio_vae = from_pretrained_with_prefetch(
model, subfolder="audio_vae", torch_dtype=dtype, local_files_only=local_files_only AutoencoderKLLTX2Audio.from_pretrained,
model,
subfolder="audio_vae",
prefetch_list=ltx2_subfolders,
local_files_only=local_files_only,
torch_dtype=dtype,
) )
# --- Vocoder: prefer BWE vocoder (48kHz) for LTX-2.3 --- # --- Vocoder: prefer BWE vocoder (48kHz) for LTX-2.3 ---
@@ -32,6 +32,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
from vllm_omni.diffusion.models.omnigen2.omnigen2_transformer import ( from vllm_omni.diffusion.models.omnigen2.omnigen2_transformer import (
OmniGen2RotaryPosEmbed, OmniGen2RotaryPosEmbed,
@@ -673,23 +674,40 @@ class OmniGen2Pipeline(CFGParallelMixin, nn.Module, SupportsComponentDiscovery):
) )
] ]
# See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
# race; prefetch the whole component set before any from_pretrained.
omnigen2_subfolders = ["scheduler", "vae", "mllm", "processor"]
prefetch_subfolders(model, omnigen2_subfolders, local_files_only=local_files_only)
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
model, subfolder="scheduler", local_files_only=local_files_only model, subfolder="scheduler", local_files_only=local_files_only
) )
self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.vae = from_pretrained_with_prefetch(
self.device AutoencoderKL.from_pretrained,
) model,
subfolder="vae",
prefetch_list=omnigen2_subfolders,
local_files_only=local_files_only,
).to(self.device)
transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, OmniGen2Transformer2DModel) transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, OmniGen2Transformer2DModel)
self.transformer = OmniGen2Transformer2DModel( self.transformer = OmniGen2Transformer2DModel(
**transformer_kwargs, **transformer_kwargs,
quant_config=od_config.quantization_config, quant_config=od_config.quantization_config,
) )
self.mllm = Qwen2_5_VLForConditionalGeneration.from_pretrained( self.mllm = from_pretrained_with_prefetch(
model, subfolder="mllm", local_files_only=local_files_only Qwen2_5_VLForConditionalGeneration.from_pretrained,
model,
subfolder="mllm",
prefetch_list=omnigen2_subfolders,
local_files_only=local_files_only,
).to(self.device) ).to(self.device)
self.processor = Qwen2_5_VLProcessor.from_pretrained( self.processor = from_pretrained_with_prefetch(
model, subfolder="processor", local_files_only=local_files_only Qwen2_5_VLProcessor.from_pretrained,
model,
subfolder="processor",
prefetch_list=omnigen2_subfolders,
local_files_only=local_files_only,
) )
self.vae_scale_factor = ( self.vae_scale_factor = (
2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8 2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8
@@ -38,6 +38,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.models.ovis_image.ovis_image_transformer import OvisImageTransformer2DModel from vllm_omni.diffusion.models.ovis_image.ovis_image_transformer import OvisImageTransformer2DModel
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -163,20 +164,32 @@ class OvisImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMi
self._execution_device = get_local_device() self._execution_device = get_local_device()
model = od_config.model model = od_config.model
local_files_only = os.path.exists(model) local_files_only = os.path.exists(model)
# See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
# race; prefetch the whole component set before any from_pretrained.
ovis_subfolders = ["scheduler", "text_encoder", "vae", "tokenizer"]
prefetch_subfolders(model, ovis_subfolders, local_files_only=local_files_only)
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
model, subfolder="scheduler", local_files_only=local_files_only model, subfolder="scheduler", local_files_only=local_files_only
) )
self.text_encoder = Qwen3Model.from_pretrained( self.text_encoder = from_pretrained_with_prefetch(
Qwen3Model.from_pretrained,
model, model,
subfolder="text_encoder", subfolder="text_encoder",
prefetch_list=ovis_subfolders,
local_files_only=local_files_only, local_files_only=local_files_only,
torch_dtype=od_config.dtype, torch_dtype=od_config.dtype,
) )
self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.vae = from_pretrained_with_prefetch(
self._execution_device AutoencoderKL.from_pretrained,
) model,
subfolder="vae",
prefetch_list=ovis_subfolders,
local_files_only=local_files_only,
).to(self._execution_device)
self.tokenizer = Qwen2TokenizerFast.from_pretrained( self.tokenizer = Qwen2TokenizerFast.from_pretrained(
model, subfolder="tokenizer", local_files_only=local_files_only model, subfolder="tokenizer", local_files_only=local_files_only
@@ -26,7 +26,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage
from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
from vllm_omni.diffusion.models.qwen_image.cfg_parallel import ( from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
QwenImageCFGParallelMixin, QwenImageCFGParallelMixin,
@@ -279,18 +279,25 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
# See pipeline_qwen_image_edit_plus: guard against transformers v5 # See pipeline_qwen_image_edit_plus: guard against transformers v5
# multi-worker race on partial subfolder shard sets (Buildkite #1043). # multi-worker race on partial subfolder shard sets (Buildkite #1043).
qwen_subfolders = ["scheduler", "text_encoder", "vae", "tokenizer"]
prefetch_subfolders( prefetch_subfolders(
model, model,
["scheduler", "text_encoder", "vae", "tokenizer"], qwen_subfolders,
) )
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
model, subfolder="scheduler", local_files_only=local_files_only model, subfolder="scheduler", local_files_only=local_files_only
) )
self.text_encoder = retry_on_missing_shard( # ``from_pretrained_with_prefetch`` re-prefetches and retries on a
lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained( # half-written cache (missing-shard ``OSError`` *and* the default
model, subfolder="text_encoder", local_files_only=local_files_only # -config size-mismatch ``RuntimeError`` that ``retry_on_missing_shard``
) # could not recover) instead of crashing the worker.
self.text_encoder = from_pretrained_with_prefetch(
Qwen2_5_VLForConditionalGeneration.from_pretrained,
model,
subfolder="text_encoder",
prefetch_list=qwen_subfolders,
local_files_only=local_files_only,
) )
# Qwen2.5-VL ships a vision tower that text-to-image does not use. # Qwen2.5-VL ships a vision tower that text-to-image does not use.
# Drop it while the model is still on CPU, before moving to GPU, so # Drop it while the model is still on CPU, before moving to GPU, so
@@ -307,8 +314,12 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
else: else:
logger.warning("Qwen-Image: vision tower not found on text encoder; skipping drop") logger.warning("Qwen-Image: vision tower not found on text encoder; skipping drop")
self.text_encoder = self.text_encoder.to(self.device) self.text_encoder = self.text_encoder.to(self.device)
self.vae = DistributedAutoencoderKLQwenImage.from_pretrained( self.vae = from_pretrained_with_prefetch(
model, subfolder="vae", local_files_only=local_files_only DistributedAutoencoderKLQwenImage.from_pretrained,
model,
subfolder="vae",
prefetch_list=qwen_subfolders,
local_files_only=local_files_only,
).to(self.device) ).to(self.device)
transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel) transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel)
self.transformer = QwenImageTransformer2DModel( self.transformer = QwenImageTransformer2DModel(
@@ -27,7 +27,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.models.interface import SupportImageInput from vllm_omni.diffusion.models.interface import SupportImageInput
from vllm_omni.diffusion.models.qwen_image.cfg_parallel import ( from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
QwenImageCFGParallelMixin, QwenImageCFGParallelMixin,
@@ -247,28 +247,43 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
# See pipeline_qwen_image_edit_plus: guard against transformers v5 # See pipeline_qwen_image_edit_plus: guard against transformers v5
# multi-worker race on partial subfolder shard sets (Buildkite #1043). # multi-worker race on partial subfolder shard sets (Buildkite #1043).
qwen_subfolders = ["scheduler", "text_encoder", "vae", "tokenizer", "processor"]
prefetch_subfolders( prefetch_subfolders(
model, model,
["scheduler", "text_encoder", "vae", "tokenizer", "processor"], qwen_subfolders,
) )
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
model, subfolder="scheduler", local_files_only=local_files_only model, subfolder="scheduler", local_files_only=local_files_only
) )
self.text_encoder = retry_on_missing_shard( # ``from_pretrained_with_prefetch`` re-prefetches and retries on a
lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained( # half-written cache (missing-shard ``OSError`` *and* the default
model, subfolder="text_encoder", local_files_only=local_files_only # -config size-mismatch ``RuntimeError`` that ``retry_on_missing_shard``
# could not recover) instead of crashing the worker.
self.text_encoder = from_pretrained_with_prefetch(
Qwen2_5_VLForConditionalGeneration.from_pretrained,
model,
subfolder="text_encoder",
prefetch_list=qwen_subfolders,
local_files_only=local_files_only,
).to(self.device) ).to(self.device)
)
self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.vae = from_pretrained_with_prefetch(
self.device AutoencoderKLQwenImage.from_pretrained,
) model,
subfolder="vae",
prefetch_list=qwen_subfolders,
local_files_only=local_files_only,
).to(self.device)
transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel) transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel)
self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs) self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs)
self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
self.processor = retry_on_missing_shard( self.processor = from_pretrained_with_prefetch(
lambda: Qwen2VLProcessor.from_pretrained(model, subfolder="processor", local_files_only=local_files_only) Qwen2VLProcessor.from_pretrained,
model,
subfolder="processor",
prefetch_list=qwen_subfolders,
local_files_only=local_files_only,
) )
self.stage = None self.stage = None
@@ -25,7 +25,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
from vllm_omni.diffusion.models.interface import SupportImageInput from vllm_omni.diffusion.models.interface import SupportImageInput
from vllm_omni.diffusion.models.qwen_image.cfg_parallel import ( from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
@@ -216,29 +216,44 @@ class QwenImageEditPlusPipeline(
# #1043 for Qwen/Qwen-Image-Edit-2509). snapshot_download takes a # #1043 for Qwen/Qwen-Image-Edit-2509). snapshot_download takes a
# per-repo file lock so the first worker downloads and the rest # per-repo file lock so the first worker downloads and the rest
# wait on a warm cache before loading. # wait on a warm cache before loading.
qwen_subfolders = ["scheduler", "text_encoder", "vae", "tokenizer", "processor"]
prefetch_subfolders( prefetch_subfolders(
model, model,
["scheduler", "text_encoder", "vae", "tokenizer", "processor"], qwen_subfolders,
) )
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
model, subfolder="scheduler", local_files_only=local_files_only model, subfolder="scheduler", local_files_only=local_files_only
) )
self.text_encoder = retry_on_missing_shard( # ``from_pretrained_with_prefetch`` re-prefetches and retries on a
lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained( # half-written cache (missing-shard ``OSError`` *and* the default
model, subfolder="text_encoder", local_files_only=local_files_only # -config size-mismatch ``RuntimeError`` that ``retry_on_missing_shard``
# could not recover) instead of crashing the worker.
self.text_encoder = from_pretrained_with_prefetch(
Qwen2_5_VLForConditionalGeneration.from_pretrained,
model,
subfolder="text_encoder",
prefetch_list=qwen_subfolders,
local_files_only=local_files_only,
).to(self.device) ).to(self.device)
)
self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.vae = from_pretrained_with_prefetch(
self.device AutoencoderKLQwenImage.from_pretrained,
) model,
subfolder="vae",
prefetch_list=qwen_subfolders,
local_files_only=local_files_only,
).to(self.device)
transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel) transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel)
self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs) self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs)
self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
self.processor = retry_on_missing_shard( self.processor = from_pretrained_with_prefetch(
lambda: Qwen2VLProcessor.from_pretrained(model, subfolder="processor", local_files_only=local_files_only) Qwen2VLProcessor.from_pretrained,
model,
subfolder="processor",
prefetch_list=qwen_subfolders,
local_files_only=local_files_only,
) )
self.stage = None self.stage = None
@@ -24,7 +24,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.models.interface import SupportImageInput from vllm_omni.diffusion.models.interface import SupportImageInput
from vllm_omni.diffusion.models.qwen_image.autoencoder_kl_qwenimage import ( from vllm_omni.diffusion.models.qwen_image.autoencoder_kl_qwenimage import (
AutoencoderKLQwenImage, AutoencoderKLQwenImage,
@@ -220,26 +220,41 @@ class QwenImageLayeredPipeline(nn.Module, SupportImageInput, QwenImageCFGParalle
# See pipeline_qwen_image_edit_plus: guard against transformers v5 # See pipeline_qwen_image_edit_plus: guard against transformers v5
# multi-worker race on partial subfolder shard sets (Buildkite #1043). # multi-worker race on partial subfolder shard sets (Buildkite #1043).
qwen_subfolders = ["scheduler", "text_encoder", "vae", "tokenizer", "processor"]
prefetch_subfolders( prefetch_subfolders(
model, model,
["scheduler", "text_encoder", "vae", "tokenizer", "processor"], qwen_subfolders,
) )
# modules keep same as transformers & diffusers # modules keep same as transformers & diffusers
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
model, subfolder="scheduler", local_files_only=local_files_only model, subfolder="scheduler", local_files_only=local_files_only
) )
self.text_encoder = retry_on_missing_shard( # ``from_pretrained_with_prefetch`` re-prefetches and retries on a
lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained( # half-written cache (missing-shard ``OSError`` *and* the default
model, subfolder="text_encoder", local_files_only=local_files_only # -config size-mismatch ``RuntimeError`` that ``retry_on_missing_shard``
# could not recover) instead of crashing the worker.
self.text_encoder = from_pretrained_with_prefetch(
Qwen2_5_VLForConditionalGeneration.from_pretrained,
model,
subfolder="text_encoder",
prefetch_list=qwen_subfolders,
local_files_only=local_files_only,
).to(self.device)
self.vae = from_pretrained_with_prefetch(
AutoencoderKLQwenImage.from_pretrained,
model,
subfolder="vae",
prefetch_list=qwen_subfolders,
local_files_only=local_files_only,
).to(self.device) ).to(self.device)
)
self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
self.device
)
self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
self.processor = retry_on_missing_shard( self.processor = from_pretrained_with_prefetch(
lambda: Qwen2VLProcessor.from_pretrained(model, subfolder="processor", local_files_only=local_files_only) Qwen2VLProcessor.from_pretrained,
model,
subfolder="processor",
prefetch_list=qwen_subfolders,
local_files_only=local_files_only,
) )
# modules re-implemented for vLLM-Omni # modules re-implemented for vLLM-Omni
+26 -13
View File
@@ -19,7 +19,7 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl import Distribu
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.models.sd3.sd3_transformer import ( from vllm_omni.diffusion.models.sd3.sd3_transformer import (
SD3Transformer2DModel, SD3Transformer2DModel,
) )
@@ -156,9 +156,7 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
# See ``hub_prefetch.py`` for the transformers v5 subfolder race. # See ``hub_prefetch.py`` for the transformers v5 subfolder race.
# SD3.5 loads six subfolders in a row, each with multi-shard # SD3.5 loads six subfolders in a row, each with multi-shard
# safetensors - it is the worst-case fan-out for the race window. # safetensors - it is the worst-case fan-out for the race window.
prefetch_subfolders( sd3_subfolders = [
model,
[
"scheduler", "scheduler",
"tokenizer", "tokenizer",
"tokenizer_2", "tokenizer_2",
@@ -167,7 +165,10 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
"text_encoder_2", "text_encoder_2",
"text_encoder_3", "text_encoder_3",
"vae", "vae",
], ]
prefetch_subfolders(
model,
sd3_subfolders,
local_files_only=local_files_only, local_files_only=local_files_only,
) )
@@ -198,31 +199,43 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
self.tokenizer_3 = T5Tokenizer.from_pretrained( self.tokenizer_3 = T5Tokenizer.from_pretrained(
model, subfolder="tokenizer_3", local_files_only=local_files_only model, subfolder="tokenizer_3", local_files_only=local_files_only
) )
self.text_encoder = CLIPTextModelWithProjection.from_pretrained( # ``from_pretrained_with_prefetch`` re-prefetches and retries if a
# peer worker left the cache half-written (missing-shard ``OSError`` or
# the default-config size-mismatch ``RuntimeError``) instead of
# crashing the worker - critical here given the six-subfolder fan-out.
self.text_encoder = from_pretrained_with_prefetch(
CLIPTextModelWithProjection.from_pretrained,
model, model,
subfolder="text_encoder", subfolder="text_encoder",
torch_dtype=dtype, prefetch_list=sd3_subfolders,
local_files_only=local_files_only, local_files_only=local_files_only,
torch_dtype=dtype,
) )
self.text_encoder_2 = CLIPTextModelWithProjection.from_pretrained( self.text_encoder_2 = from_pretrained_with_prefetch(
CLIPTextModelWithProjection.from_pretrained,
model, model,
subfolder="text_encoder_2", subfolder="text_encoder_2",
torch_dtype=dtype, prefetch_list=sd3_subfolders,
local_files_only=local_files_only, local_files_only=local_files_only,
torch_dtype=dtype,
) )
self.text_encoder_3 = T5EncoderModel.from_pretrained( self.text_encoder_3 = from_pretrained_with_prefetch(
T5EncoderModel.from_pretrained,
model, model,
subfolder="text_encoder_3", subfolder="text_encoder_3",
torch_dtype=dtype, prefetch_list=sd3_subfolders,
local_files_only=local_files_only, local_files_only=local_files_only,
torch_dtype=dtype,
) )
self.transformer = SD3Transformer2DModel(od_config=od_config) self.transformer = SD3Transformer2DModel(od_config=od_config)
self.vae = DistributedAutoencoderKL.from_pretrained( self.vae = from_pretrained_with_prefetch(
DistributedAutoencoderKL.from_pretrained,
model, model,
subfolder="vae", subfolder="vae",
torch_dtype=dtype, prefetch_list=sd3_subfolders,
local_files_only=local_files_only, local_files_only=local_files_only,
torch_dtype=dtype,
).to(self.device) ).to(self.device)
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
@@ -28,6 +28,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.models.interface import SupportAudioOutput from vllm_omni.diffusion.models.interface import SupportAudioOutput
from vllm_omni.diffusion.models.stable_audio.stable_audio_transformer import ( from vllm_omni.diffusion.models.stable_audio.stable_audio_transformer import (
StableAudioDiTModel, StableAudioDiTModel,
@@ -109,6 +110,11 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfil
), ),
] ]
# See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
# race; prefetch the whole component set before any from_pretrained.
sa_subfolders = ["tokenizer", "text_encoder", "vae", "projection_model", "scheduler"]
prefetch_subfolders(model, sa_subfolders, local_files_only=local_files_only)
# Load tokenizer # Load tokenizer
self.tokenizer = T5TokenizerFast.from_pretrained( self.tokenizer = T5TokenizerFast.from_pretrained(
model, model,
@@ -117,27 +123,33 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfil
) )
# Load text encoder # Load text encoder
self.text_encoder = T5EncoderModel.from_pretrained( self.text_encoder = from_pretrained_with_prefetch(
T5EncoderModel.from_pretrained,
model, model,
subfolder="text_encoder", subfolder="text_encoder",
torch_dtype=dtype, prefetch_list=sa_subfolders,
local_files_only=local_files_only, local_files_only=local_files_only,
torch_dtype=dtype,
).to(self.device) ).to(self.device)
# Load VAE (AutoencoderOobleck for audio) # Load VAE (AutoencoderOobleck for audio)
self.vae = AutoencoderOobleck.from_pretrained( self.vae = from_pretrained_with_prefetch(
AutoencoderOobleck.from_pretrained,
model, model,
subfolder="vae", subfolder="vae",
torch_dtype=torch.float32, prefetch_list=sa_subfolders,
local_files_only=local_files_only, local_files_only=local_files_only,
torch_dtype=torch.float32,
).to(self.device) ).to(self.device)
# Load projection model (using diffusers implementation) # Load projection model (using diffusers implementation)
self.projection_model = StableAudioProjectionModel.from_pretrained( self.projection_model = from_pretrained_with_prefetch(
StableAudioProjectionModel.from_pretrained,
model, model,
subfolder="projection_model", subfolder="projection_model",
torch_dtype=dtype, prefetch_list=sa_subfolders,
local_files_only=local_files_only, local_files_only=local_files_only,
torch_dtype=dtype,
).to(self.device) ).to(self.device)
# Initialize transformer from HF config to keep architecture aligned with checkpoint. # Initialize transformer from HF config to keep architecture aligned with checkpoint.
@@ -26,7 +26,7 @@ from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, Pipe
from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
@@ -336,18 +336,39 @@ class Wan22Pipeline(
) )
# See ``hub_prefetch.py`` for the transformers v5 subfolder race. # See ``hub_prefetch.py`` for the transformers v5 subfolder race.
component_subfolders = ["tokenizer", "text_encoder", "vae"]
prefetch_subfolders( prefetch_subfolders(
model, model,
["tokenizer", "text_encoder", "vae"], component_subfolders,
local_files_only=local_files_only, local_files_only=local_files_only,
) )
self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) # ``from_pretrained_with_prefetch`` re-prefetches and retries if the
self.text_encoder = UMT5EncoderModel.from_pretrained( # cache is still half-written (the missing-shard ``OSError`` and the
model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only # default-``UMT5Config`` size-mismatch ``RuntimeError`` seen on multi
# -worker HSDP / ring launches), instead of crashing the worker.
self.tokenizer = from_pretrained_with_prefetch(
AutoTokenizer.from_pretrained,
model,
subfolder="tokenizer",
prefetch_list=component_subfolders,
local_files_only=local_files_only,
)
self.text_encoder = from_pretrained_with_prefetch(
UMT5EncoderModel.from_pretrained,
model,
subfolder="text_encoder",
prefetch_list=component_subfolders,
local_files_only=local_files_only,
torch_dtype=dtype,
).to(self.device) ).to(self.device)
self.vae = DistributedAutoencoderKLWan.from_pretrained( self.vae = from_pretrained_with_prefetch(
model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only DistributedAutoencoderKLWan.from_pretrained,
model,
subfolder="vae",
prefetch_list=component_subfolders,
local_files_only=local_files_only,
torch_dtype=dtype,
).to(self.device) ).to(self.device)
# Initialize transformers with correct config (weights loaded via load_weights) # Initialize transformers with correct config (weights loaded via load_weights)
@@ -26,7 +26,7 @@ from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, Pipe
from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
from vllm_omni.diffusion.models.interface import SupportImageInput from vllm_omni.diffusion.models.interface import SupportImageInput
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
@@ -197,35 +197,64 @@ class Wan22I2VPipeline(
) )
) )
# Text encoder
self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
self.text_encoder = UMT5EncoderModel.from_pretrained(
model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
).to(self.device)
# Image encoder (CLIP) - optional, for Wan2.1-style I2V # Image encoder (CLIP) - optional, for Wan2.1-style I2V
self.has_image_encoder = "image_encoder" in model_index and model_index["image_encoder"][0] is not None self.has_image_encoder = "image_encoder" in model_index and model_index["image_encoder"][0] is not None
# See ``hub_prefetch.py`` for the transformers v5 subfolder race. # See ``hub_prefetch.py`` for the transformers v5 subfolder race. The
# prefetch MUST run before any ``from_pretrained`` below - previously
# the tokenizer / text_encoder were loaded ahead of this call, so they
# never benefited from the prefetch and could hit the half-written
# cache directly.
subfolders = ["tokenizer", "text_encoder", "vae"] subfolders = ["tokenizer", "text_encoder", "vae"]
if self.has_image_encoder: if self.has_image_encoder:
subfolders.extend(["image_processor", "image_encoder"]) subfolders.extend(["image_processor", "image_encoder"])
prefetch_subfolders(model, subfolders, local_files_only=local_files_only) prefetch_subfolders(model, subfolders, local_files_only=local_files_only)
if self.has_image_encoder: # Text encoder
self.image_processor = CLIPImageProcessor.from_pretrained( self.tokenizer = from_pretrained_with_prefetch(
model, subfolder="image_processor", local_files_only=local_files_only AutoTokenizer.from_pretrained,
model,
subfolder="tokenizer",
prefetch_list=subfolders,
local_files_only=local_files_only,
) )
self.image_encoder = CLIPVisionModel.from_pretrained( self.text_encoder = from_pretrained_with_prefetch(
model, subfolder="image_encoder", torch_dtype=dtype, local_files_only=local_files_only UMT5EncoderModel.from_pretrained,
model,
subfolder="text_encoder",
prefetch_list=subfolders,
local_files_only=local_files_only,
torch_dtype=dtype,
).to(self.device)
if self.has_image_encoder:
self.image_processor = from_pretrained_with_prefetch(
CLIPImageProcessor.from_pretrained,
model,
subfolder="image_processor",
prefetch_list=subfolders,
local_files_only=local_files_only,
)
self.image_encoder = from_pretrained_with_prefetch(
CLIPVisionModel.from_pretrained,
model,
subfolder="image_encoder",
prefetch_list=subfolders,
local_files_only=local_files_only,
torch_dtype=dtype,
).to(self.device) ).to(self.device)
else: else:
self.image_processor = None self.image_processor = None
self.image_encoder = None self.image_encoder = None
# VAE # VAE
self.vae = DistributedAutoencoderKLWan.from_pretrained( self.vae = from_pretrained_with_prefetch(
model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only DistributedAutoencoderKLWan.from_pretrained,
model,
subfolder="vae",
prefetch_list=subfolders,
local_files_only=local_files_only,
torch_dtype=dtype,
).to(self.device) ).to(self.device)
# Transformers (weights loaded via load_weights) # Transformers (weights loaded via load_weights)
@@ -1,4 +1,8 @@
import contextlib
import errno
import os
import time import time
from collections.abc import Iterator
from pathlib import Path from pathlib import Path
import huggingface_hub import huggingface_hub
@@ -13,6 +17,159 @@ else:
logger = init_logger(__name__) logger = init_logger(__name__)
_DOWNLOAD_MAX_ATTEMPTS = 3
_DOWNLOAD_BACKOFF_BASE_S = 1.0
_FULL_SNAPSHOT_METADATA_MARKERS = (
"config.json",
"preprocessor_config.json",
"processor_config.json",
"tokenizer_config.json",
"model_index.json",
)
def _node_lock_dir() -> str:
candidates: list[str] = []
hf_home = os.environ.get("HF_HOME")
if hf_home:
candidates.append(os.path.join(hf_home, "locks", "vllm-omni-weight-prefetch"))
xdg_cache = os.environ.get("XDG_CACHE_HOME")
if xdg_cache:
candidates.append(os.path.join(xdg_cache, "huggingface", "locks", "vllm-omni-weight-prefetch"))
candidates.append(
os.path.join(os.path.expanduser("~"), ".cache", "huggingface", "locks", "vllm-omni-weight-prefetch")
)
candidates.append(os.path.join("/tmp", "vllm-omni-weight-prefetch-locks"))
for cand in candidates:
try:
os.makedirs(cand, exist_ok=True)
probe = os.path.join(cand, ".write_check")
with open(probe, "a"):
pass
with contextlib.suppress(OSError):
os.remove(probe)
return cand
except OSError:
continue
fallback = os.path.join("/tmp", "vllm-omni-weight-prefetch-locks")
os.makedirs(fallback, exist_ok=True)
return fallback
def _safe_repo_filename(model: str) -> str:
return model.replace("/", "__").replace(os.sep, "__") + ".lock"
def _dotfile_lock_acquire(lock_dir: str, model: str, timeout: float = 300.0, poll_interval: float = 0.5) -> str | None:
lock_path = os.path.join(lock_dir, _safe_repo_filename(model) + ".dir")
deadline = time.monotonic() + timeout
while True:
try:
os.makedirs(lock_path, exist_ok=False)
logger.info("Acquired dotfile weight prefetch lock for %s at %s", model, lock_path)
return lock_path
except FileExistsError:
if time.monotonic() >= deadline:
logger.warning(
"Timed out waiting for dotfile weight prefetch lock %s after %.0fs; proceeding unlocked",
lock_path,
timeout,
)
return None
time.sleep(poll_interval)
@contextlib.contextmanager
def _repo_download_lock(model: str) -> Iterator[None]:
"""Serialize full-repo snapshot materialization across Omni processes."""
lock_dir = None
dotfile_held = None
fd = None
flock_held = False
try:
import fcntl # type: ignore[import-not-found]
except ImportError: # pragma: no cover - non-POSIX
fcntl = None
if fcntl is not None:
try:
lock_dir = _node_lock_dir()
lock_path = os.path.join(lock_dir, _safe_repo_filename(model))
fd = os.open(lock_path, os.O_RDWR | os.O_CREAT, 0o644)
fcntl.flock(fd, fcntl.LOCK_EX)
flock_held = True
logger.info("Acquired flock weight prefetch lock for %s at %s", model, lock_path)
except OSError as exc:
if fd is not None:
with contextlib.suppress(OSError):
os.close(fd)
fd = None
if exc.errno not in (errno.ENOLCK, errno.EOPNOTSUPP, errno.EACCES, errno.EINVAL):
raise
logger.warning("fcntl.flock unavailable for weight prefetch of %s (%s); using dotfile lock", model, exc)
if not flock_held:
try:
lock_dir = lock_dir or _node_lock_dir()
dotfile_held = _dotfile_lock_acquire(lock_dir, model)
except OSError as exc:
logger.warning("Could not allocate weight prefetch lock dir for %s (%s); proceeding unlocked", model, exc)
try:
yield
finally:
if flock_held and fd is not None:
with contextlib.suppress(OSError):
fcntl.flock(fd, fcntl.LOCK_UN)
if fd is not None:
with contextlib.suppress(OSError):
os.close(fd)
if dotfile_held is not None:
with contextlib.suppress(OSError):
os.rmdir(dotfile_held)
def _looks_like_auth_error(exc: BaseException) -> bool:
try:
from huggingface_hub.errors import GatedRepoError, RepositoryNotFoundError # type: ignore[import-not-found]
if isinstance(exc, GatedRepoError | RepositoryNotFoundError):
return True
except ImportError: # pragma: no cover - very old huggingface_hub
pass
status = getattr(getattr(exc, "response", None), "status_code", None)
if status in (401, 403):
return True
msg = str(exc).lower()
return "401 client error" in msg or "403 client error" in msg or "gatedrepo" in msg
def _is_full_snapshot_request(allow_patterns: list[str]) -> bool:
return any(pattern.strip() in {"*", "**", "./*", "**/*"} for pattern in allow_patterns)
def _verify_snapshot_materialized(hf_folder: str, allow_patterns: list[str]) -> None:
root = Path(hf_folder)
if not root.exists():
raise OSError(f"Downloaded snapshot folder does not exist: {hf_folder}")
if not any(root.glob(pattern) for pattern in allow_patterns):
raise OSError(f"Downloaded snapshot {hf_folder} has no files matching {allow_patterns}")
# Full-repo downloads are used before transformers loads tokenizers,
# processors, and feature extractors. A half-materialized shared cache can
# contain weights while missing these small metadata files, which later
# surfaces as "Can't load feature extractor ... preprocessor_config.json".
if _is_full_snapshot_request(allow_patterns) and not any(
(root / marker).exists() for marker in _FULL_SNAPSHOT_METADATA_MARKERS
):
raise OSError(
f"Downloaded full snapshot {hf_folder} is missing expected metadata files {_FULL_SNAPSHOT_METADATA_MARKERS}"
)
def download_weights_from_hf_specific( def download_weights_from_hf_specific(
model_name_or_path: str, model_name_or_path: str,
@@ -44,14 +201,20 @@ def download_weights_from_hf_specific(
str: The path to the downloaded model weights. str: The path to the downloaded model weights.
""" """
assert len(allow_patterns) > 0 assert len(allow_patterns) > 0
allow_patterns = list(allow_patterns)
local_only = huggingface_hub.constants.HF_HUB_OFFLINE local_only = huggingface_hub.constants.HF_HUB_OFFLINE
download_kwargs = {"tqdm_class": DisabledTqdm} if not envs.VLLM_USE_MODELSCOPE else {} download_kwargs = {"tqdm_class": DisabledTqdm} if not envs.VLLM_USE_MODELSCOPE else {}
logger.info("Using model weights format %s", allow_patterns) logger.info("Using model weights format %s", allow_patterns)
# Use file lock to prevent multiple processes from
# downloading the same model weights at the same time.
with get_lock(model_name_or_path, cache_dir):
start_time = time.perf_counter() start_time = time.perf_counter()
last_exc: BaseException | None = None
for attempt in range(1, _DOWNLOAD_MAX_ATTEMPTS + 1):
try:
# Use both vLLM's cache lock and a repo-wide node lock. The latter
# protects the transformers v5 eager metadata resolver from seeing
# a peer process' half-written shared HF cache.
with get_lock(model_name_or_path, cache_dir), _repo_download_lock(model_name_or_path):
if require_all: if require_all:
hf_folder = snapshot_download( hf_folder = snapshot_download(
model_name_or_path, model_name_or_path,
@@ -62,7 +225,10 @@ def download_weights_from_hf_specific(
local_files_only=local_only, local_files_only=local_only,
**download_kwargs, **download_kwargs,
) )
else: _verify_snapshot_materialized(hf_folder, allow_patterns)
break
hf_folder = None
for allow_pattern in allow_patterns: for allow_pattern in allow_patterns:
hf_folder = snapshot_download( hf_folder = snapshot_download(
model_name_or_path, model_name_or_path,
@@ -73,10 +239,34 @@ def download_weights_from_hf_specific(
local_files_only=local_only, local_files_only=local_only,
**download_kwargs, **download_kwargs,
) )
_verify_snapshot_materialized(hf_folder, [allow_pattern])
# If we have downloaded weights for this allow_pattern, # If we have downloaded weights for this allow_pattern,
# we don't need to check the rest, unless require_all is set. # we don't need to check the rest, unless require_all is set.
if any(Path(hf_folder).glob(allow_pattern)): if any(Path(hf_folder).glob(allow_pattern)):
break break
if hf_folder is None:
raise OSError(f"No snapshot downloaded for patterns {allow_patterns}")
break
except Exception as exc:
last_exc = exc
if _looks_like_auth_error(exc) or attempt == _DOWNLOAD_MAX_ATTEMPTS or local_only:
raise
backoff = _DOWNLOAD_BACKOFF_BASE_S * attempt
logger.warning(
"Downloading weights for %s with patterns %s failed on attempt %d/%d (%s: %s); retrying in %.1fs",
model_name_or_path,
allow_patterns,
attempt,
_DOWNLOAD_MAX_ATTEMPTS,
type(exc).__name__,
exc,
backoff,
)
time.sleep(backoff)
if last_exc is not None and "hf_folder" not in locals():
raise last_exc
time_taken = time.perf_counter() - start_time time_taken = time.perf_counter() - start_time
if time_taken > 0.5: if time_taken > 0.5:
logger.info( logger.info(