fix(vllm): No module named 'ChatTTS.model'

Author: fumiama

ChatTTS/model/velocity/block_manager.py

@@ -4,7 +4,7 @@
 from typing import Dict, List, Optional, Set, Tuple
 
 from vllm.block import PhysicalTokenBlock
-from ChatTTS.model.velocity.sequence import Sequence, SequenceGroup, SequenceStatus
+from .sequence import Sequence, SequenceGroup, SequenceStatus
 from vllm.utils import Device
 
 # Mapping: logical block number -> physical block.

ChatTTS/model/velocity/llm.py

@@ -3,10 +3,10 @@
 from tqdm import tqdm
 from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
 
-from ChatTTS.model.velocity.configs import EngineArgs
-from ChatTTS.model.velocity.llm_engine import LLMEngine
-from ChatTTS.model.velocity.output import RequestOutput
-from ChatTTS.model.velocity.sampling_params import SamplingParams
+from .configs import EngineArgs
+from .llm_engine import LLMEngine
+from .output import RequestOutput
+from .sampling_params import SamplingParams
 from vllm.utils import Counter
 
 

ChatTTS/model/velocity/llm_engine.py

@@ -5,14 +5,14 @@
 from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union
 
 from vllm.config import CacheConfig, ModelConfig, ParallelConfig, SchedulerConfig
-from ChatTTS.model.velocity.scheduler import Scheduler, SchedulerOutputs
-from ChatTTS.model.velocity.configs import EngineArgs
+from .scheduler import Scheduler, SchedulerOutputs
+from .configs import EngineArgs
 from vllm.engine.metrics import record_metrics
 from vllm.engine.ray_utils import RayWorkerVllm, initialize_cluster, ray
 from vllm.logger import init_logger
-from ChatTTS.model.velocity.output import RequestOutput
-from ChatTTS.model.velocity.sampling_params import SamplingParams
-from ChatTTS.model.velocity.sequence import (
+from .output import RequestOutput
+from .sampling_params import SamplingParams
+from .sequence import (
     SamplerOutput,
     Sequence,
     SequenceGroup,
@@ -127,7 +127,7 @@ def __init__(
     def _init_workers(self):
         # Lazy import the Worker to avoid importing torch.cuda/xformers
         # before CUDA_VISIBLE_DEVICES is set in the Worker
-        from ChatTTS.model.velocity.worker import Worker
+        from .worker import Worker
 
         assert (
             self.parallel_config.world_size == 1

ChatTTS/model/velocity/model_loader.py

@@ -24,7 +24,7 @@ def _set_default_torch_dtype(dtype: torch.dtype):
 
 def _get_model_architecture(config: PretrainedConfig) -> Type[nn.Module]:
     model_cls = getattr(
-        importlib.import_module("ChatTTS.model.velocity.llama"), "LlamaModel", None
+        importlib.import_module(".llama"), "LlamaModel", None
     )
     return model_cls
 

ChatTTS/model/velocity/model_runner.py

@@ -5,24 +5,24 @@
 import torch
 import torch.nn as nn
 
-from ChatTTS.model.velocity.configs import ModelConfig, ParallelConfig, SchedulerConfig
+from .configs import ModelConfig, ParallelConfig, SchedulerConfig
 from vllm.logger import init_logger
-from ChatTTS.model.velocity.model_loader import get_model
+from .model_loader import get_model
 from vllm.model_executor import InputMetadata, SamplingMetadata
 from vllm.model_executor.parallel_utils.communication_op import (
     broadcast,
     broadcast_object_list,
 )
-from ChatTTS.model.velocity.sampling_params import SamplingParams, SamplingType
-from ChatTTS.model.velocity.sequence import (
+from .sampling_params import SamplingParams, SamplingType
+from .sequence import (
     SamplerOutput,
     SequenceData,
     SequenceGroupMetadata,
     SequenceGroupOutput,
     SequenceOutput,
 )
 from vllm.utils import in_wsl
-from ChatTTS.model.velocity.post_model import PostModel, Sampler
+from .post_model import PostModel, Sampler
 from safetensors.torch import safe_open
 
 logger = init_logger(__name__)

ChatTTS/model/velocity/output.py

@@ -1,7 +1,7 @@
 from typing import List, Optional
 import torch
 
-from ChatTTS.model.velocity.sequence import (
+from .sequence import (
     PromptLogprobs,
     SampleLogprobs,
     SequenceGroup,

ChatTTS/model/velocity/scheduler.py

@@ -3,10 +3,10 @@
 from typing import Dict, Iterable, List, Optional, Tuple, Union
 
 from vllm.config import CacheConfig, SchedulerConfig
-from ChatTTS.model.velocity.block_manager import AllocStatus, BlockSpaceManager
+from .block_manager import AllocStatus, BlockSpaceManager
 from vllm.core.policy import PolicyFactory
 from vllm.logger import init_logger
-from ChatTTS.model.velocity.sequence import (
+from .sequence import (
     Sequence,
     SequenceData,
     SequenceGroup,

ChatTTS/model/velocity/sequence.py

@@ -5,7 +5,7 @@
 from typing import Dict, List, Optional, Union
 import torch
 from vllm.block import LogicalTokenBlock
-from ChatTTS.model.velocity.sampling_params import SamplingParams
+from .sampling_params import SamplingParams
 
 PromptLogprobs = List[Optional[Dict[int, float]]]
 SampleLogprobs = List[Dict[int, float]]

ChatTTS/model/velocity/worker.py

@@ -12,7 +12,7 @@
 from vllm.model_executor.parallel_utils.parallel_state import initialize_model_parallel
 from vllm.sequence import SamplerOutput, SequenceGroupMetadata
 from vllm.worker.cache_engine import CacheEngine
-from ChatTTS.model.velocity.model_runner import ModelRunner
+from .model_runner import ModelRunner
 
 
 class Worker:

requirements.txt

@@ -15,3 +15,4 @@ nemo_text_processing; sys_platform == 'linux'
 av
 pydub
 safetensors
+vllm>=0.2.7; sys_platform == 'linux'