[Refactor][Misc] Remove unused code in diffusion hooks/model_loader (#5270)

Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>

vllm_omni/diffusion/hooks/__init__.py

@@ -3,7 +3,6 @@
 """Hook mechanism for model forward interception."""
 
 from vllm_omni.diffusion.hooks.base import (
-    BaseState,
     HookRegistry,
     ModelHook,
     StateManager,
@@ -12,14 +11,11 @@ from vllm_omni.diffusion.hooks.sequence_parallel import (
     SequenceParallelGatherHook,
     SequenceParallelSplitHook,
     apply_sequence_parallel,
-    disable_sequence_parallel_for_model,
-    enable_sequence_parallel_for_model,
     remove_sequence_parallel,
 )
 
 __all__ = [
     # Base hooks
-    "BaseState",
     "StateManager",
     "ModelHook",
     "HookRegistry",
@@ -28,6 +24,4 @@ __all__ = [
     "SequenceParallelGatherHook",
     "apply_sequence_parallel",
     "remove_sequence_parallel",
-    "enable_sequence_parallel_for_model",
-    "disable_sequence_parallel_for_model",
 ]

vllm_omni/diffusion/hooks/base.py

@@ -20,21 +20,18 @@ from vllm.logger import init_logger
 logger = init_logger(__name__)
 
 
-class BaseState:
-    """Base class for hook state containers."""
-
-    def reset(self) -> None:  # pragma: no cover - default is no-op
-        pass
-
-
 class StateManager:
-    """Manage per-context hook state instances."""
+    """Manage per-context hook state instances.
 
-    def __init__(self, state_cls: Callable[[], BaseState], init_args: tuple = (), init_kwargs: dict | None = None):
+    State classes (e.g. TeaCacheState, MagCacheState) are plain classes;
+    they only need to be constructible with the given init args.
+    """
+
+    def __init__(self, state_cls: Callable[..., Any], init_args: tuple = (), init_kwargs: dict | None = None):
         self._state_cls = state_cls
         self._init_args = init_args
         self._init_kwargs = init_kwargs or {}
-        self._states: dict[str, BaseState] = {}
+        self._states: dict[str, Any] = {}
         self._context: str = "default"
 
     @property
@@ -45,7 +42,7 @@ class StateManager:
     def set_context(self, name: str) -> None:
         self._context = name or "default"
 
-    def get_state(self) -> BaseState:
+    def get_state(self) -> Any:
         if self._context not in self._states:
             self._states[self._context] = self._state_cls(*self._init_args, **self._init_kwargs)
         return self._states[self._context]
@@ -332,17 +329,3 @@ class HookRegistry:
         hook = self._hooks.get(name)
         if hook is not None:
             hook.reset_state(self.module)
-
-    def reset(self) -> None:
-        """Reset all hooks and clear the registry.
-
-        This removes all hooks from the registry and resets each hook's state.
-        Also restores module.forward to its original implementation.
-        """
-        for name, hook in list(self._hooks.items()):
-            hook.reset_state(self.module)
-        self._hooks.clear()
-
-        if hasattr(self.module, "_omni_original_forward"):
-            self.module.forward = self.module._omni_original_forward  # type: ignore[attr-defined]
-            delattr(self.module, "_omni_original_forward")

vllm_omni/diffusion/hooks/sequence_parallel.py

@@ -679,87 +679,3 @@ def remove_sequence_parallel(
                 continue
 
             registry.remove_hook(hook_name)
-
-
-def enable_sequence_parallel_for_model(
-    model: nn.Module,
-    config: SequenceParallelConfig | None = None,
-) -> None:
-    """Enable sequence parallelism for a model using its _sp_plan.
-
-    This is a convenience function that reads the model's _sp_plan attribute
-    and applies sequence parallelism automatically.
-
-    Note: This corresponds to `enable_context_parallel_for_model` in diffusers,
-    but uses vLLM-Omni's _sp_plan instead of diffusers' _cp_plan.
-
-    The function performs two main tasks:
-    1. Applies _sp_plan hooks to shard inputs and gather outputs
-    2. Ensures Attention layers are configured for the correct parallel mode
-       (handled automatically by vLLM-Omni's forward_context mechanism)
-
-    Args:
-        model: The model to enable SP for. Must have a _sp_plan attribute.
-        config: Optional config. If None, uses default based on current
-            parallel state.
-
-    Raises:
-        ValueError: If model has no _sp_plan defined.
-
-    Note:
-        vLLM-Omni supports Ulysses + Ring hybrid parallelism:
-        - ulysses_degree > 1: Uses All-to-All communication over Q/K/V heads
-        - ring_degree > 1: Uses Ring attention with K/V passing
-        - Both > 1: Hybrid mode (Ulysses handles head redistribution,
-          Ring handles K/V circulation)
-    """
-    from vllm_omni.diffusion.distributed.parallel_state import (
-        get_ring_parallel_world_size,
-        get_ulysses_parallel_world_size,
-    )
-    from vllm_omni.diffusion.distributed.sp_plan import get_sp_plan_from_model
-
-    plan = get_sp_plan_from_model(model)
-    if plan is None:
-        raise ValueError(
-            f"Model {model.__class__.__name__} has no _sp_plan defined. "
-            f"Define _sp_plan as a class attribute or pass a plan explicitly."
-        )
-
-    if config is None:
-        # Create config from current parallel state
-        ulysses_degree = get_ulysses_parallel_world_size()
-        ring_degree = get_ring_parallel_world_size()
-        config = SequenceParallelConfig(
-            ulysses_degree=ulysses_degree,
-            ring_degree=ring_degree,
-        )
-        if ulysses_degree > 1 and ring_degree > 1:
-            mode = "hybrid"
-        elif ulysses_degree > 1:
-            mode = "ulysses"
-        else:
-            mode = "ring"
-        logger.info(
-            f"Created SP config from parallel state: "
-            f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, "
-            f"mode={mode}"
-        )
-
-    apply_sequence_parallel(model, config, plan)
-    logger.info(f"Enabled sequence parallelism for {model.__class__.__name__}")
-
-
-def disable_sequence_parallel_for_model(model: nn.Module) -> None:
-    """Disable sequence parallelism for a model.
-
-    Note: This corresponds to `disable_context_parallel_for_model` in diffusers.
-
-    Args:
-        model: The model to disable SP for.
-    """
-    from vllm_omni.diffusion.distributed.sp_plan import get_sp_plan_from_model
-
-    plan = get_sp_plan_from_model(model)
-    if plan is not None:
-        remove_sequence_parallel(model, plan)

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -12,7 +12,6 @@ from typing import cast
 
 import torch
 from torch import nn
-from vllm.config import ModelConfig
 from vllm.config.load import LoadConfig
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
@@ -81,7 +80,6 @@ def _natural_sort_key(filepath: str) -> list:
     return [int(s) if s.isdigit() else s for s in re.split(r"(\d+)", os.path.basename(filepath))]
 
 
-MODEL_INDEX = "model_index.json"
 DIFFUSION_MODEL_WEIGHTS_INDEX = "diffusion_pytorch_model.safetensors.index.json"
 TRANSFORMER_WEIGHTS_INDEX = "model.safetensors.index.json"
 INDEX_FILES = [DIFFUSION_MODEL_WEIGHTS_INDEX, TRANSFORMER_WEIGHTS_INDEX]
@@ -326,15 +324,6 @@ class DiffusersPipelineLoader:
             return all_parameter_names
         return {name for name in all_parameter_names if name.startswith(source_prefixes)}
 
-    def download_model(self, model_config: ModelConfig) -> None:
-        self._prepare_weights(
-            model_name_or_path=model_config.model,
-            subfolder=None,
-            revision=model_config.revision,
-            fall_back_to_pt=True,
-            allow_patterns_overrides=None,
-        )
-
     def load_model(
         self,
         load_device: str,

vllm_omni/diffusion/model_loader/hub_prefetch.py

@@ -433,10 +433,7 @@ def from_pretrained_with_prefetch(
     ``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:
+    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
@@ -488,31 +485,3 @@ def from_pretrained_with_prefetch(
 
     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):
-    """Call *load_fn* with retry on the transformers v5 shard-resolution race.
-
-    When the prefetch lock cannot be acquired (e.g. flock unsupported on
-    the filesystem and dotfile lock times out), ``from_pretrained`` may
-    still hit the ``cached_files`` race. This wrapper retries with
-    exponential backoff when the OSError message matches the specific
-    "does not appear to have a file named" pattern.
-    """
-    for attempt in range(max_retries):
-        try:
-            return load_fn()
-        except OSError as exc:
-            if "does not appear to have a file" not in str(exc):
-                raise
-            if attempt == max_retries - 1:
-                raise
-            delay = base_delay * (attempt + 1)
-            logger.warning(
-                "from_pretrained failed with shard-resolution race (%s); retrying in %.1fs (attempt %d/%d)",
-                exc,
-                delay,
-                attempt + 1,
-                max_retries,
-            )
-            time.sleep(delay)

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -309,9 +309,8 @@ class QwenImagePipeline(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
         # ``from_pretrained_with_prefetch`` re-prefetches and retries on a
-        # half-written cache (missing-shard ``OSError`` *and* the default
-        # -config size-mismatch ``RuntimeError`` that ``retry_on_missing_shard``
-        # could not recover) instead of crashing the worker.
+        # half-written cache (missing-shard ``OSError`` and the default
+        # -config size-mismatch ``RuntimeError``) instead of crashing the worker.
         self.text_encoder = from_pretrained_with_prefetch(
             Qwen2_5_VLForConditionalGeneration.from_pretrained,
             model,

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py

@@ -263,9 +263,8 @@ class QwenImageEditPipeline(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
         # ``from_pretrained_with_prefetch`` re-prefetches and retries on a
-        # half-written cache (missing-shard ``OSError`` *and* the default
-        # -config size-mismatch ``RuntimeError`` that ``retry_on_missing_shard``
-        # could not recover) instead of crashing the worker.
+        # half-written cache (missing-shard ``OSError`` and the default
+        # -config size-mismatch ``RuntimeError``) instead of crashing the worker.
         self.text_encoder = from_pretrained_with_prefetch(
             Qwen2_5_VLForConditionalGeneration.from_pretrained,
             model,

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py

@@ -230,9 +230,8 @@ class QwenImageEditPlusPipeline(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
         # ``from_pretrained_with_prefetch`` re-prefetches and retries on a
-        # half-written cache (missing-shard ``OSError`` *and* the default
-        # -config size-mismatch ``RuntimeError`` that ``retry_on_missing_shard``
-        # could not recover) instead of crashing the worker.
+        # half-written cache (missing-shard ``OSError`` and the default
+        # -config size-mismatch ``RuntimeError``) instead of crashing the worker.
         self.text_encoder = from_pretrained_with_prefetch(
             Qwen2_5_VLForConditionalGeneration.from_pretrained,
             model,

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py

@@ -237,9 +237,8 @@ class QwenImageLayeredPipeline(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
         # ``from_pretrained_with_prefetch`` re-prefetches and retries on a
-        # half-written cache (missing-shard ``OSError`` *and* the default
-        # -config size-mismatch ``RuntimeError`` that ``retry_on_missing_shard``
-        # could not recover) instead of crashing the worker.
+        # half-written cache (missing-shard ``OSError`` and the default
+        # -config size-mismatch ``RuntimeError``) instead of crashing the worker.
         self.text_encoder = from_pretrained_with_prefetch(
             Qwen2_5_VLForConditionalGeneration.from_pretrained,
             model,