[Diffusion] Refactor diffusion models weights loading (#157)

Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -0,0 +1,228 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import dataclasses
+import glob
+import os
+import time
+from collections.abc import Generator, Iterable
+from pathlib import Path
+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.model_loader.utils import set_default_torch_dtype
+from vllm.model_executor.model_loader.weight_utils import (
+    download_safetensors_index_file_from_hf,
+    download_weights_from_hf,
+    filter_duplicate_safetensors_files,
+    filter_files_not_needed_for_inference,
+    maybe_download_from_modelscope,
+    safetensors_weights_iterator,
+)
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.registry import initialize_model
+
+logger = init_logger(__name__)
+
+
+MODEL_INDEX = "model_index.json"
+DIFFUSION_MODEL_WEIGHTS_INDEX = "diffusion_pytorch_model.safetensors.index.json"
+
+
+class DiffusersPipelineLoader:
+    """Model loader that can load diffusers pipeline components from disk."""
+
+    # default number of thread when enable multithread weight loading
+    DEFAULT_NUM_THREADS = 8
+
+    @dataclasses.dataclass
+    class ComponentSource:
+        """A source for weights."""
+
+        model_or_path: str
+        """The model ID or path."""
+
+        subfolder: str | None
+        """The subfolder inside the model repo."""
+
+        revision: str | None
+        """The optional model revision."""
+
+        prefix: str = ""
+        """A prefix to prepend to all weights."""
+
+        fall_back_to_pt: bool = True
+        """Whether .pt weights can be used."""
+
+        allow_patterns_overrides: list[str] | None = None
+        """If defined, weights will load exclusively using these patterns."""
+
+    counter_before_loading_weights: float = 0.0
+    counter_after_loading_weights: float = 0.0
+
+    def __init__(self, load_config: LoadConfig):
+        self.load_config = load_config
+
+        # TODO(Isotr0py): Enable multithreaded weight loading
+        # extra_config = load_config.model_loader_extra_config
+        # allowed_keys = {"enable_multithread_load", "num_threads"}
+        # unexpected_keys = set(extra_config.keys()) - allowed_keys
+
+        # if unexpected_keys:
+        #     raise ValueError(
+        #         f"Unexpected extra config keys for load format {load_config.load_format}: {unexpected_keys}"
+        #     )
+
+    def _prepare_weights(
+        self,
+        model_name_or_path: Path,
+        subfolder: str | None,
+        revision: str | None,
+        fall_back_to_pt: bool,
+        allow_patterns_overrides: list[str] | None,
+    ) -> tuple[str, list[str], bool]:
+        """Prepare weights for the model.
+
+        If the model is not local, it will be downloaded."""
+        model_name_or_path = maybe_download_from_modelscope(model_name_or_path, revision) or model_name_or_path
+
+        is_local = os.path.isdir(model_name_or_path)
+        load_format = self.load_config.load_format
+        use_safetensors = False
+        index_file = DIFFUSION_MODEL_WEIGHTS_INDEX
+
+        # only hf is supported currently
+        if load_format == "auto":
+            load_format = "hf"
+
+        # Some quantized models use .pt files for storing the weights.
+        if load_format == "hf":
+            allow_patterns = ["*.safetensors", "*.bin"]
+        else:
+            raise ValueError(f"Unknown load_format: {load_format}")
+
+        if fall_back_to_pt:
+            allow_patterns += ["*.pt"]
+
+        if allow_patterns_overrides is not None:
+            allow_patterns = allow_patterns_overrides
+
+        if subfolder is not None:
+            allow_patterns = [f"{subfolder}/{pattern}" for pattern in allow_patterns]
+
+        if not is_local:
+            hf_folder = download_weights_from_hf(
+                model_name_or_path,
+                self.load_config.download_dir,
+                allow_patterns,
+                revision,
+                ignore_patterns=self.load_config.ignore_patterns,
+            )
+        else:
+            hf_folder = model_name_or_path
+
+        hf_weights_files: list[str] = []
+        for pattern in allow_patterns:
+            hf_weights_files += glob.glob(os.path.join(hf_folder, pattern))
+            if len(hf_weights_files) > 0:
+                if pattern == "*.safetensors":
+                    use_safetensors = True
+                break
+
+        if use_safetensors:
+            # For models like Mistral-7B-Instruct-v0.3
+            # there are both sharded safetensors files and a consolidated
+            # safetensors file. Using both breaks.
+            # Here, we download the `model.safetensors.index.json` and filter
+            # any files not found in the index.
+            if not is_local:
+                download_safetensors_index_file_from_hf(
+                    model_name_or_path,
+                    index_file,
+                    self.load_config.download_dir,
+                    revision,
+                )
+            hf_weights_files = filter_duplicate_safetensors_files(hf_weights_files, hf_folder, index_file)
+        else:
+            hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files)
+
+        if len(hf_weights_files) == 0:
+            raise RuntimeError(f"Cannot find any model weights with `{model_name_or_path}`")
+
+        return hf_folder, hf_weights_files, use_safetensors
+
+    def _get_weights_iterator(self, source: "ComponentSource") -> Generator[tuple[str, torch.Tensor], None, None]:
+        """Get an iterator for the model weights based on the load format."""
+        hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
+            source.model_or_path,
+            source.subfolder,
+            source.revision,
+            source.fall_back_to_pt,
+            source.allow_patterns_overrides,
+        )
+        weights_iterator = safetensors_weights_iterator(
+            hf_weights_files,
+            self.load_config.use_tqdm_on_load,
+            self.load_config.safetensors_load_strategy,
+        )
+
+        if self.counter_before_loading_weights == 0.0:
+            self.counter_before_loading_weights = time.perf_counter()
+        # Apply the prefix.
+        return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator)
+
+    def get_all_weights(
+        self,
+        model: nn.Module,
+    ) -> Generator[tuple[str, torch.Tensor], None, None]:
+        sources = cast(
+            Iterable[DiffusersPipelineLoader.ComponentSource],
+            getattr(model, "weights_sources", ()),
+        )
+        for source in sources:
+            yield from self._get_weights_iterator(source)
+
+    def download_model(self, model_config: ModelConfig) -> None:
+        self._prepare_weights(
+            model_config.model,
+            model_config.revision,
+            fall_back_to_pt=True,
+            allow_patterns_overrides=None,
+        )
+
+    def load_model(self, od_config: OmniDiffusionConfig, load_device: str) -> nn.Module:
+        """Load a model with the given configurations."""
+        target_device = torch.device(load_device)
+        with set_default_torch_dtype(od_config.dtype):
+            with target_device:
+                model = initialize_model(od_config)
+
+            logger.debug("Loading weights on %s ...", load_device)
+            # Quantization does not happen in `load_weights` but after it
+            self.load_weights(model)
+        return model.eval()
+
+    def load_weights(self, model: nn.Module) -> None:
+        weights_to_load = {name for name, _ in model.named_parameters()}
+        loaded_weights = model.load_weights(self.get_all_weights(model))
+
+        self.counter_after_loading_weights = time.perf_counter()
+        logger.info_once(
+            "Loading weights took %.2f seconds",
+            self.counter_after_loading_weights - self.counter_before_loading_weights,
+        )
+        # TODO(Isotr0py): Enable weights loading check after decoupling
+        # all components' weights loading (AutoModel.from_pretrained etc).
+        # We only enable strict check for non-quantized models
+        # that have loaded weights tracking currently.
+        if loaded_weights is not None:
+            _ = weights_to_load - loaded_weights
+        #     if weights_not_loaded:
+        #         raise ValueError(
+        #             "Following weights were not initialized from "
+        #             f"checkpoint: {weights_not_loaded}"
+        #         )

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -6,6 +6,7 @@ import json
 import logging
 import math
 import os
+from collections.abc import Iterable
 from typing import Any, Optional, Union
 
 import numpy as np
@@ -20,9 +21,11 @@ from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
 from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer
+from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
     QwenImageTransformer2DModel,
 )
@@ -240,9 +243,18 @@ class QwenImagePipeline(
     ):
         super().__init__()
         self.od_config = od_config
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="transformer",
+                revision=None,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            )
+        ]
+
         self.device = get_local_device()
         model = od_config.model
-
         # Check if model is a local path
         local_files_only = os.path.exists(model)
 
@@ -731,52 +743,6 @@ class QwenImagePipeline(
 
         return DiffusionOutput(output=image)
 
-    def load_weights(self):
-        self.load_transformer()
-
-    # TODO: find a better way to load
-    def load_transformer(self):
-        import glob
-        import os
-
-        # Define the weight iterator
-        def weight_iterator(transformer_path):
-            if not os.path.exists(transformer_path):
-                logger.warning(f"Path {transformer_path} does not exist.")
-                return
-
-            # Look for safetensors first
-            safetensors_files = glob.glob(os.path.join(transformer_path, "*.safetensors"))
-            if safetensors_files:
-                try:
-                    from safetensors.torch import load_file
-                except ImportError:
-                    logger.warning("safetensors not installed, cannot load .safetensors files.")
-                    return
-
-                for file_path in safetensors_files:
-                    state_dict = load_file(file_path)
-                    for name, tensor in state_dict.items():
-                        yield name, tensor
-            else:
-                # Fallback to bin
-                bin_files = glob.glob(os.path.join(transformer_path, "*.bin"))
-                for file_path in bin_files:
-                    state_dict = torch.load(file_path)
-                    for name, tensor in state_dict.items():
-                        yield name, tensor
-
-        try:
-            # Get model path from config or download from HF
-            model_name = self.od_config.model if hasattr(self, "od_config") else "Qwen/Qwen-Image"
-            if os.path.exists(model_name):
-                model_path = model_name
-            else:
-                model_path = download_weights_from_hf_specific(model_name, None, ["*"])
-
-            transformer_path = os.path.join(model_path, "transformer")
-            self.transformer.load_weights(weight_iterator(transformer_path))
-
-        except Exception as e:
-            print(f"An error occurred: {e}")
-            raise e
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -16,6 +16,7 @@ from diffusers.models.normalization import AdaLayerNormContinuous
 from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
 
@@ -603,79 +604,39 @@ class QwenImageTransformer2DModel(nn.Module):
 
         return Transformer2DModelOutput(sample=output)
 
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
-        # Collect incoming weights into a dict for flexible mapping and grouping
-        weights_map: dict[str, torch.Tensor] = {name: tensor for name, tensor in weights}
-
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         stacked_params_mapping = [
-            (
-                "to_qkv",
-                [
-                    {"aliases": ["q_proj", "to_q", "q", "query"], "order": 0},
-                    {"aliases": ["k_proj", "to_k", "k", "key"], "order": 1},
-                    {"aliases": ["v_proj", "to_v", "v", "value"], "order": 2},
-                ],
-            ),
-            (
-                "add_kv_proj",
-                [
-                    {"aliases": ["add_q_proj", "add_q"], "order": 0},
-                    {"aliases": ["add_k_proj", "add_k"], "order": 1},
-                    {"aliases": ["add_v_proj", "add_v"], "order": 2},
-                ],
-            ),
+            # (param_name, shard_name, shard_id)
+            # self-attn
+            (".to_qkv", ".to_q", "q"),
+            (".to_qkv", ".to_k", "k"),
+            (".to_qkv", ".to_v", "v"),
+            # cross-attn
+            (".add_kv_proj", ".add_q_proj", "q"),
+            (".add_kv_proj", ".add_k_proj", "k"),
+            (".add_kv_proj", ".add_v_proj", "v"),
         ]
 
-        # Helper that attempts to gather shards per base key using aliases and merges them
-        for target, shards in stacked_params_mapping:
-            base_candidates: dict[str, dict[str, dict[int, tuple[str, torch.Tensor]]]] = {}
-            for key in list(weights_map.keys()):
-                for shard in shards:
-                    for alias in shard["aliases"]:
-                        weight_suffix = f".{alias}.weight"
-                        bias_suffix = f".{alias}.bias"
-                        if key.endswith(weight_suffix):
-                            base = key[: -len(weight_suffix)]
-                            base_entry = base_candidates.setdefault(base, {"weight": {}, "bias": {}})
-                            base_entry["weight"][shard["order"]] = (
-                                key,
-                                weights_map[key],
-                            )
-                            break
-                        if key.endswith(bias_suffix):
-                            base = key[: -len(bias_suffix)]
-                            base_entry = base_candidates.setdefault(base, {"weight": {}, "bias": {}})
-                            base_entry["bias"][shard["order"]] = (key, weights_map[key])
-                            break
-
-            for base, entry in base_candidates.items():
-                if len(entry["weight"]) == len(shards):
-                    ordered = [entry["weight"][order] for order in sorted(entry["weight"])]
-                    combined = torch.cat([tensor for _, tensor in ordered], dim=0)
-                    for key, _ in ordered:
-                        weights_map.pop(key, None)
-                    weights_map[base + f".{target}.weight"] = combined
-                if len(entry["bias"]) == len(shards):
-                    ordered_bias = [entry["bias"][order] for order in sorted(entry["bias"])]
-                    bcombined = torch.cat([tensor for _, tensor in ordered_bias], dim=0)
-                    for key, _ in ordered_bias:
-                        weights_map.pop(key, None)
-                    weights_map[base + f".{target}.bias"] = bcombined
-
-        # Now perform the parameter copy using the possibly-updated weights_map
         params_dict = dict(self.named_parameters())
-        for name, loaded_weight in weights_map.items():
-            # Handle potential name mismatches for FeedForward layers
-            if name.endswith(".proj.weight") and name not in params_dict:
-                mapped_name = name.replace(".proj.weight", ".weight")
-            elif name.endswith(".proj.bias") and name not in params_dict:
-                mapped_name = name.replace(".proj.bias", ".bias")
-            else:
-                mapped_name = name
 
-            if mapped_name in params_dict:
-                param = params_dict[mapped_name]
-                with torch.no_grad():
-                    param.copy_(loaded_weight)
+        # we need to load the buffers for beta and eps (XIELU)
+        for name, buffer in self.named_buffers():
+            if name.endswith(".beta") or name.endswith(".eps"):
+                params_dict[name] = buffer
+
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
             else:
-                logger.warning(f"Weight {name} (mapped to {mapped_name}) not found in model params.")
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params

vllm_omni/diffusion/models/z_image/pipeline_z_image.py

@@ -18,6 +18,7 @@
 import inspect
 import json
 import os
+from collections.abc import Iterable
 from typing import Any, Callable
 
 import torch
@@ -28,9 +29,11 @@ from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
 from diffusers.utils import logging
 from diffusers.utils.torch_utils import randn_tensor
 from transformers import AutoModel, AutoTokenizer
+from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.z_image.z_image_transformer import (
     ZImageTransformer2DModel,
 )
@@ -148,6 +151,15 @@ class ZImagePipeline(nn.Module):
     ):
         super().__init__()
         self.od_config = od_config
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="transformer",
+                revision=None,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            )
+        ]
         self._execution_device = get_local_device()
         model = od_config.model
         local_files_only = os.path.exists(model)
@@ -598,51 +610,6 @@ class ZImagePipeline(nn.Module):
 
         return DiffusionOutput(output=image)
 
-    def load_weights(self):
-        self.load_transformer()
-
-    def load_transformer(self):
-        import glob
-        import os
-
-        # Define the weight iterator
-        def weight_iterator(transformer_path):
-            if not os.path.exists(transformer_path):
-                logger.warning(f"Path {transformer_path} does not exist.")
-                return
-
-            # Look for safetensors first
-            safetensors_files = glob.glob(os.path.join(transformer_path, "*.safetensors"))
-            if safetensors_files:
-                try:
-                    from safetensors.torch import load_file
-                except ImportError:
-                    logger.warning("safetensors not installed, cannot load .safetensors files.")
-                    return
-
-                for file_path in safetensors_files:
-                    state_dict = load_file(file_path)
-                    for name, tensor in state_dict.items():
-                        yield name, tensor
-            else:
-                # Fallback to bin
-                bin_files = glob.glob(os.path.join(transformer_path, "*.bin"))
-                for file_path in bin_files:
-                    state_dict = torch.load(file_path)
-                    for name, tensor in state_dict.items():
-                        yield name, tensor
-
-        try:
-            # Get model path from config or download from HF
-            model_name = self.od_config.model
-            if os.path.exists(model_name):
-                model_path = model_name
-            else:
-                model_path = download_weights_from_hf_specific(model_name, None, ["*"])
-
-            transformer_path = os.path.join(model_path, "transformer")
-            self.transformer.load_weights(weight_iterator(transformer_path))
-
-        except Exception as e:
-            print(f"An error occurred: {e}")
-            raise e
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)

vllm_omni/diffusion/models/z_image/z_image_transformer.py

@@ -24,6 +24,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch.nn.utils.rnn import pad_sequence
 from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
 
@@ -642,33 +643,13 @@ class ZImageTransformer2DModel(nn.Module):
 
         return x, {}
 
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
-        weights_map: dict[str, torch.Tensor] = {name: tensor for name, tensor in weights}
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         params_dict = dict(self.named_parameters())
 
-        for name, loaded_weight in weights_map.items():
-            # In ZImageAttention, to_q, to_k, to_v are separate.
-            # The weight names are like '...attention.to_q.weight'.
-            # The logic from qwen_image_transformer for stacking is not needed here.
-            # We can do a direct mapping.
-
-            # The names in the state dict seem to match the model structure.
-            # Let's try a direct load.
-            if name in params_dict:
-                param = params_dict[name]
-                with torch.no_grad():
-                    param.copy_(loaded_weight)
-            else:
-                # The provided weight list might use a different naming convention.
-                # For example, `layers.0.attention.to_out.0.weight` vs `...`
-                # Let's try to handle simple cases.
-
-                # Replace dots with hyphens for ModuleDict keys
-                # e.g. all_x_embedder.2-1.weight -> all_x_embedder.2-1.weight
-                # The names seem to match already.
-
-                # If there are any specific mismatches, they can be handled here.
-                # For now, we assume a direct mapping and warn if a weight is not found.
-                import logging
-
-                logging.warning(f"Weight {name} not found in model parameters.")
+        loaded_params = set[str]()
+        for name, loaded_weight in weights:
+            param = params_dict[name]
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params

vllm_omni/diffusion/worker/gpu_worker.py

@@ -2,20 +2,21 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import multiprocessing as mp
 import os
+import time
 
 import torch
 import zmq
-from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.config import LoadConfig, VllmConfig, set_current_vllm_config
 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
 from vllm.distributed.parallel_state import (
     init_distributed_environment,
     initialize_model_parallel,
 )
 from vllm.logger import init_logger
-from vllm.model_executor.model_loader.utils import set_default_torch_dtype
+from vllm.utils import DeviceMemoryProfiler, GiB_bytes
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.registry import initialize_model
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 logger = init_logger(__name__)
@@ -60,13 +61,23 @@ class GPUWorker:
 
         init_distributed_environment(world_size=world_size, rank=rank)
         initialize_model_parallel(tensor_model_parallel_size=world_size)
+        logger.info(f"Worker {self.rank}: Initialized device and distributed environment.")
 
-        with device:
-            with set_default_torch_dtype(self.od_config.dtype):
-                self.pipeline = initialize_model(self.od_config)
-                self.pipeline.load_weights()
-                self.pipeline.eval()
-        logger.info(f"Worker {self.rank}: Initialized device, model, and distributed environment.")
+        load_config = LoadConfig()
+        model_loader = DiffusersPipelineLoader(load_config)
+        time_before_load = time.perf_counter()
+        with DeviceMemoryProfiler() as m:
+            self.pipeline = model_loader.load_model(
+                od_config=self.od_config,
+                load_device=f"cuda:{rank}",
+            )
+        time_after_load = time.perf_counter()
+
+        logger.info(
+            "Model loading took %.4f GiB and %.6f seconds",
+            m.consumed_memory / GiB_bytes,
+            time_after_load - time_before_load,
+        )
         logger.info(f"Worker {self.rank}: Model loaded successfully.")
 
     @torch.inference_mode()

vllm_omni/diffusion/worker/npu/npu_worker.py

@@ -2,21 +2,22 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import multiprocessing as mp
 import os
+import time
 
 import torch
 import zmq
 from transformers import PretrainedConfig
-from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
+from vllm.config import LoadConfig, ModelConfig, VllmConfig, set_current_vllm_config
 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
 from vllm.distributed.parallel_state import (
     init_distributed_environment,
     initialize_model_parallel,
 )
 from vllm.logger import init_logger
-from vllm.model_executor.model_loader.utils import set_default_torch_dtype
+from vllm.utils import DeviceMemoryProfiler, GiB_bytes
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.registry import initialize_model
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 logger = init_logger(__name__)
@@ -60,19 +61,27 @@ class NPUWorker:
             def __getattr__(self, name):
                 return "fake"
 
-        vllm_config = VllmConfig(model_config=ModelConfig(hf_config=_FakePretrainedConfig()))
+        vllm_config = VllmConfig(model_config=ModelConfig(hf_config=_FakePretrainedConfig()), load_config=LoadConfig())
         vllm_config.parallel_config.tensor_parallel_size = self.od_config.num_gpus
         set_current_vllm_config(vllm_config)
 
         init_distributed_environment(world_size=world_size, rank=rank)
         initialize_model_parallel(tensor_model_parallel_size=world_size)
 
-        with device:
-            with set_default_torch_dtype(self.od_config.dtype):
-                self.pipeline = initialize_model(self.od_config)
-                self.pipeline.load_weights()
-                self.pipeline.eval()
-        logger.info(f"Worker {self.rank}: Initialized device, model, and distributed environment.")
+        model_loader = DiffusersPipelineLoader(vllm_config.load_config)
+        time_before_load = time.perf_counter()
+        with DeviceMemoryProfiler() as m:
+            self.pipeline = model_loader.load_model(
+                od_config=self.od_config,
+                load_device=f"npu:{rank}",
+            )
+        time_after_load = time.perf_counter()
+
+        logger.info(
+            "Model loading took %.4f GiB and %.6f seconds",
+            m.consumed_memory / GiB_bytes,
+            time_after_load - time_before_load,
+        )
         logger.info(f"Worker {self.rank}: Model loaded successfully.")
 
     @torch.inference_mode()

vllm_omni/model_executor/model_loader/weight_utils.py

@@ -3,10 +3,15 @@ from pathlib import Path
 from typing import Optional, Union
 
 import huggingface_hub
-from huggingface_hub import snapshot_download
+import vllm.envs as envs
 from vllm.logger import init_logger
 from vllm.model_executor.model_loader.weight_utils import DisabledTqdm, get_lock
 
+if envs.VLLM_USE_MODELSCOPE:
+    from modelscope.hub.snapshot_download import snapshot_download
+else:
+    from huggingface_hub import snapshot_download
+
 logger = init_logger(__name__)
 
 
@@ -37,6 +42,7 @@ def download_weights_from_hf_specific(
     """
     assert len(allow_patterns) > 0
     local_only = huggingface_hub.constants.HF_HUB_OFFLINE
+    download_kwargs = {"tqdm_class": DisabledTqdm} if not envs.VLLM_USE_MODELSCOPE else {}
 
     logger.info("Using model weights format %s", allow_patterns)
     # Use file lock to prevent multiple processes from
@@ -49,9 +55,9 @@ def download_weights_from_hf_specific(
                 allow_patterns=allow_pattern,
                 ignore_patterns=ignore_patterns,
                 cache_dir=cache_dir,
-                tqdm_class=DisabledTqdm,
                 revision=revision,
                 local_files_only=local_only,
+                **download_kwargs,
             )
             # If we have downloaded weights for this allow_pattern,
             # we don't need to check the rest.