diff --git a/.gitignore b/.gitignore index 7fb302410..6a004ce25 100644 --- a/.gitignore +++ b/.gitignore @@ -20,6 +20,7 @@ parts/ sdist/ var/ wheels/ +vllm/ share/python-wheels/ *.egg-info/ .installed.cfg diff --git a/vllm_omni/__init__.py b/vllm_omni/__init__.py index 7d8be2a9f..afeb09c1b 100644 --- a/vllm_omni/__init__.py +++ b/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 diff --git a/vllm_omni/model_executor/model_loader/__init__.py b/vllm_omni/model_executor/model_loader/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/model_executor/model_loader/weight_utils.py b/vllm_omni/model_executor/model_loader/weight_utils.py new file mode 100644 index 000000000..ee5e5a6d0 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni.py index e4ff6cbaa..894cd6084 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni.py +++ b/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"],