Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d2c149f25b | ||
|
|
f0758f5cbe | ||
|
|
69fd466fba | ||
|
|
cb08af976b | ||
|
|
617c6ca55f |
No files matched your search
@@ -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
|
||||||
|
|||||||
@@ -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(
|
||||||
|
|||||||
Reference in new issue
Block a user