Merge pull request #47 from Gaohan123/refractor

[Bugfix] fix huggingface download problem for spk_dict.pt

.gitignore

@@ -20,6 +20,7 @@ parts/
 sdist/
 var/
 wheels/
+vllm/
 share/python-wheels/
 *.egg-info/
 .installed.cfg

vllm_omni/__init__.py

@@ -17,7 +17,6 @@ from .config import OmniModelConfig
 
 # Main entry points
 from .entrypoints.omni_llm import OmniLLM
-from .entrypoints.omni_llm import OmniLLM
 
 __all__ = [
     # Main components

vllm_omni/model_executor/model_loader/weight_utils.py

@@ -0,0 +1,59 @@
+import time
+from typing import Optional, Union
+
+import huggingface_hub
+from huggingface_hub import snapshot_download
+from vllm.logger import init_logger
+from vllm.model_executor.model_loader.weight_utils import DisabledTqdm, get_lock
+
+logger = init_logger(__name__)
+
+
+def download_weights_from_hf_specific(
+    model_name_or_path: str,
+    cache_dir: Optional[str],
+    allow_patterns: list[str],
+    revision: Optional[str] = None,
+    ignore_patterns: Optional[Union[str, list[str]]] = None,
+) -> str:
+    """Download model weights from Hugging Face Hub.
+
+    Args:
+        model_name_or_path (str): The model name or path.
+        cache_dir (Optional[str]): The cache directory to store the model
+            weights. If None, will use HF defaults.
+        allow_patterns (list[str]): The allowed patterns for the
+            weight files. Files matched by any of the patterns will be
+            downloaded.
+        revision (Optional[str]): The revision of the model.
+        ignore_patterns (Optional[Union[str, list[str]]]): The patterns to
+            filter out the weight files. Files matched by any of the patterns
+            will be ignored.
+
+    Returns:
+        str: The path to the downloaded model weights.
+    """
+    local_only = huggingface_hub.constants.HF_HUB_OFFLINE
+
+    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()
+        hf_folder = snapshot_download(
+            model_name_or_path,
+            allow_patterns=allow_patterns,
+            ignore_patterns=ignore_patterns,
+            cache_dir=cache_dir,
+            tqdm_class=DisabledTqdm,
+            revision=revision,
+            local_files_only=local_only,
+        )
+        time_taken = time.perf_counter() - start_time
+        if time_taken > 0.5:
+            logger.info(
+                "Time spent downloading weights for %s: %.6f seconds",
+                model_name_or_path,
+                time_taken,
+            )
+    return hf_folder

vllm_omni/model_executor/models/qwen2_5_omni.py

@@ -14,7 +14,6 @@ from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
-from vllm.model_executor.model_loader.weight_utils import download_weights_from_hf
 from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
@@ -29,6 +28,9 @@ from vllm.model_executor.sampling_metadata import SamplingMetadata
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.sequence import IntermediateTensors
 
+from vllm_omni.model_executor.model_loader.weight_utils import (
+    download_weights_from_hf_specific,
+)
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
 
 TALKER_CODEC_EOS_TOKEN_ID = 8294
@@ -817,7 +819,7 @@ class Qwen2_5OmniForConditionalGeneration(
 
         # Load token2wav weights (if any)
         if token2wav_weights and self.token2wav is not None:
-            hf_model_folder = download_weights_from_hf(
+            hf_model_folder = download_weights_from_hf_specific(
                 self.vllm_config.model_config.model,
                 self.vllm_config.load_config.download_dir,
                 allow_patterns=["*.safetensors", "*.bin", "*.pt"],