From 89fe7c064b085304fd786063d0f0d83680ea3a44 Mon Sep 17 00:00:00 2001 From: MuShan Date: Fri, 27 Mar 2026 22:14:03 +0800 Subject: [PATCH] compatible ARM. --- api/main.py | 5 ++++- api/models.py | 3 ++- api/requirements.txt | 1 + api/service.py | 4 +++- cli/podcast.py | 6 ++++++ cli/tts.py | 6 ++++++ requirements.txt | 8 +++++--- soulxpodcast/engine/llm_engine.py | 15 ++++++++++---- soulxpodcast/models/soulxpodcast.py | 29 ++++++++++++++------------ soulxpodcast/utils/commons.py | 3 ++- soulxpodcast/utils/device.py | 32 +++++++++++++++++++++++++++++ 11 files changed, 88 insertions(+), 24 deletions(-) create mode 100644 soulxpodcast/utils/device.py diff --git a/api/main.py b/api/main.py index 139b06c..f152163 100644 --- a/api/main.py +++ b/api/main.py @@ -31,6 +31,7 @@ validate_dialogue_format, cleanup_old_files, ) +from soulxpodcast.utils.device import get_best_device # 配置日志 logging.basicConfig( @@ -127,11 +128,13 @@ async def health_check(): """健康检查""" service = get_service() task_manager = get_task_manager() + device = get_best_device() return HealthResponse( status="healthy", model_loaded=service.is_loaded(), - gpu_available=torch.cuda.is_available(), + gpu_available=device.type in {"cuda", "mps"}, + accelerator=device.type, llm_engine=config.llm_engine, active_tasks=task_manager.get_active_task_count(), version="1.0.0" diff --git a/api/models.py b/api/models.py index 10d8e3f..01fe722 100644 --- a/api/models.py +++ b/api/models.py @@ -103,7 +103,8 @@ class HealthResponse(BaseModel): """健康检查响应""" status: str = Field(default="healthy", description="服务状态") model_loaded: bool = Field(..., description="模型是否已加载") - gpu_available: bool = Field(..., description="GPU是否可用") + gpu_available: bool = Field(..., description="是否存在可用的硬件加速后端（CUDA/MPS）") + accelerator: str = Field(..., description="当前使用的加速后端（cuda/mps/cpu）") llm_engine: str = Field(..., description="当前使用的LLM引擎 (hf/vllm)") active_tasks: int = Field(default=0, description="正在处理的任务数") version: str = Field(default="1.0.0", description="API版本") diff --git a/api/requirements.txt b/api/requirements.txt index 92306a4..d27431c 100644 --- a/api/requirements.txt +++ b/api/requirements.txt @@ -11,3 +11,4 @@ scipy>=1.11.0 # Logging and utilities aiofiles>=23.2.0 +requests>=2.31.0 diff --git a/api/service.py b/api/service.py index 13c5f63..f5519c9 100644 --- a/api/service.py +++ b/api/service.py @@ -75,7 +75,9 @@ def _load_model(self): ) self.config = model_config - logger.info(f"Model loaded successfully with {api_config.llm_engine} engine!") + logger.info( + f"Model loaded successfully with {api_config.llm_engine} engine on {self.model.device}!" + ) except Exception as e: logger.error(f"Failed to load model: {e}") diff --git a/cli/podcast.py b/cli/podcast.py index d61dfd5..27f5515 100644 --- a/cli/podcast.py +++ b/cli/podcast.py @@ -2,10 +2,16 @@ import json import torch import argparse +import sys +from pathlib import Path import s3tokenizer import soundfile as sf +PROJECT_ROOT = Path(__file__).resolve().parents[1] +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + from soulxpodcast.config import SamplingParams from soulxpodcast.utils.parser import podcast_format_parser from soulxpodcast.utils.infer_utils import initiate_model, process_single_input diff --git a/cli/tts.py b/cli/tts.py index 25b172c..67d68f3 100644 --- a/cli/tts.py +++ b/cli/tts.py @@ -2,10 +2,16 @@ import json import torch import argparse +import sys +from pathlib import Path import s3tokenizer import soundfile as sf +PROJECT_ROOT = Path(__file__).resolve().parents[1] +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + from soulxpodcast.config import SamplingParams from soulxpodcast.utils.parser import podcast_format_parser from soulxpodcast.utils.infer_utils import initiate_model, process_single_input diff --git a/requirements.txt b/requirements.txt index da8131c..ab0cf2d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,10 +5,12 @@ s3tokenizer diffusers torch==2.7.1 torchaudio==2.7.1 -triton>=3.0.0 +triton>=3.0.0; platform_system == "Linux" transformers==4.57.1 accelerate==1.10.1 onnxruntime -onnxruntime-gpu +onnxruntime-gpu; platform_system == "Linux" einops -gradio \ No newline at end of file +gradio +soundfile +tqdm diff --git a/soulxpodcast/engine/llm_engine.py b/soulxpodcast/engine/llm_engine.py index 6694974..42ea8d4 100644 --- a/soulxpodcast/engine/llm_engine.py +++ b/soulxpodcast/engine/llm_engine.py @@ -19,6 +19,7 @@ from soulxpodcast.config import Config, SamplingParams from soulxpodcast.models.modules.sampler import _ras_sample_hf_engine +from soulxpodcast.utils.device import get_best_device, get_llm_dtype class HFLLMEngine: @@ -29,8 +30,12 @@ def __init__(self, model, **kwargs): self.tokenizer = AutoTokenizer.from_pretrained(model, use_fast=True) config.eos = config.hf_config.eos_token_id # speech eos token; - self.device = "cuda:0" if torch.cuda.is_available() else "cpu" - self.model = AutoModelForCausalLM.from_pretrained(model, torch_dtype=torch.bfloat16, device_map=self.device) + self.device = get_best_device() + self.model = AutoModelForCausalLM.from_pretrained( + model, + torch_dtype=get_llm_dtype(self.device), + ).to(self.device) + self.model.eval() self.config = config self.pad_token_id = self.tokenizer.pad_token_id @@ -85,7 +90,9 @@ def __init__(self, model, **kwargs): self.tokenizer = AutoTokenizer.from_pretrained(config.model, use_fast=True) config.eos = config.hf_config.eos_token_id # speech eos token; - self.device = "cuda:0" if torch.cuda.is_available() else "cpu" + self.device = get_best_device() + if self.device.type != "cuda": + raise RuntimeError("vLLM currently requires a CUDA device.") os.environ["VLLM_USE_V1"] = "0" if SUPPORT_VLLM: self.model = LLM(model=model, enforce_eager=True, dtype="bfloat16", max_model_len=8192, enable_prefix_caching=True,) @@ -111,4 +118,4 @@ def generate( "text": self.tokenizer.decode(generated_ids), "token_ids": list(generated_ids), } - return output \ No newline at end of file + return output diff --git a/soulxpodcast/models/soulxpodcast.py b/soulxpodcast/models/soulxpodcast.py index 05b2a01..6e96e61 100644 --- a/soulxpodcast/models/soulxpodcast.py +++ b/soulxpodcast/models/soulxpodcast.py @@ -16,13 +16,15 @@ ) from soulxpodcast.models.modules.flow import CausalMaskedDiffWithXvec from soulxpodcast.models.modules.hifigan import HiFTGenerator +from soulxpodcast.utils.device import get_amp_context, get_best_device class SoulXPodcast(torch.nn.Module): def __init__(self, config: Config = None): super().__init__() self.config = Config() if config is None else config + self.device = get_best_device() - self.audio_tokenizer = s3tokenizer.load_model("speech_tokenizer_v2_25hz").cuda().eval() + self.audio_tokenizer = s3tokenizer.load_model("speech_tokenizer_v2_25hz").to(self.device).eval() if self.config.llm_engine == "hf": self.llm = HFLLMEngine(**self.config.__dict__) elif self.config.llm_engine == "vllm": @@ -38,12 +40,12 @@ def __init__(self, config: Config = None): tqdm.write(f"[{timestamp}] - [INFO] - Casting flow to fp16") self.flow.half() self.flow.load_state_dict(torch.load(f"{self.config.model}/flow.pt", map_location="cpu", weights_only=True), strict=True) - self.flow.cuda().eval() + self.flow.to(self.device).eval() self.hift = HiFTGenerator() hift_state_dict = {k.replace('generator.', ''): v for k, v in torch.load(f"{self.config.model}/hift.pt", map_location="cpu", weights_only=True).items()} self.hift.load_state_dict(hift_state_dict, strict=True) - self.hift.cuda().eval() + self.hift.to(self.device).eval() @torch.inference_mode() @@ -66,7 +68,7 @@ def forward_longform( # Audio tokenization prompt_speech_tokens_ori, prompt_speech_tokens_lens_ori = self.audio_tokenizer.quantize( - prompt_mels_for_llm.cuda(), prompt_mels_lens_for_llm.cuda() + prompt_mels_for_llm.to(self.device), prompt_mels_lens_for_llm.to(self.device) ) # align speech token with speech feat as to reduce @@ -81,10 +83,11 @@ def forward_longform( prompt_mel_len = prompt_mel.shape[0] if prompt_speech_token_len * 2 > prompt_mel_len: prompt_speech_token = prompt_speech_token[:int(prompt_mel_len/2)] - prompt_mel_len = torch.tensor([prompt_mel_len]).cuda() + prompt_mel = prompt_mel.detach().clone().to(self.device) + prompt_mel_len = torch.tensor([prompt_mel_len], device=self.device) else: - prompt_mel = prompt_mel.detach().clone()[:prompt_speech_token_len * 2].cuda() - prompt_mel_len = torch.tensor([prompt_speech_token_len * 2]).cuda() + prompt_mel = prompt_mel.detach().clone()[:prompt_speech_token_len * 2].to(self.device) + prompt_mel_len = torch.tensor([prompt_speech_token_len * 2], device=self.device) prompt_speech_tokens.append(prompt_speech_token) prompt_mels_for_flow.append(prompt_mel) prompt_mels_lens_for_flow.append(prompt_mel_len) @@ -141,8 +144,8 @@ def forward_longform( turn_spk = spk_ids[i] generated_speech_tokens = [token - self.config.hf_config.speech_token_offset for token in llm_outputs['token_ids'][:-1]] # ignore last eos prompt_speech_token = prompt_speech_tokens[turn_spk].tolist() - flow_input = torch.tensor([prompt_speech_token + generated_speech_tokens]) - flow_inputs_len = torch.tensor([len(prompt_speech_token) + len(generated_speech_tokens)]) + flow_input = torch.tensor([prompt_speech_token + generated_speech_tokens], device=self.device) + flow_inputs_len = torch.tensor([len(prompt_speech_token) + len(generated_speech_tokens)], device=self.device) # Flow generation and HiFi-GAN generation start_idx = spk_ids[i] @@ -151,10 +154,10 @@ def forward_longform( spk_emb = spk_emb_for_flow[start_idx:start_idx+1] # Flow generation - with torch.amp.autocast("cuda", dtype=torch.float16 if self.config.hf_config.fp16_flow else torch.float32): + with get_amp_context(self.device, enabled=self.config.hf_config.fp16_flow): generated_mels, generated_mels_lens = self.flow( - flow_input.cuda(), flow_inputs_len.cuda(), - prompt_mels, prompt_mels_lens, spk_emb.cuda(), + flow_input, flow_inputs_len, + prompt_mels, prompt_mels_lens, spk_emb.to(self.device), streaming=False, finalize=True ) @@ -165,4 +168,4 @@ def forward_longform( # Save the generated wav; results_dict['generated_wavs'] = generated_wavs - return results_dict \ No newline at end of file + return results_dict diff --git a/soulxpodcast/utils/commons.py b/soulxpodcast/utils/commons.py index 4ff87d0..2f57674 100644 --- a/soulxpodcast/utils/commons.py +++ b/soulxpodcast/utils/commons.py @@ -7,4 +7,5 @@ def set_all_random_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) \ No newline at end of file + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) diff --git a/soulxpodcast/utils/device.py b/soulxpodcast/utils/device.py new file mode 100644 index 0000000..3c72d0c --- /dev/null +++ b/soulxpodcast/utils/device.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +from contextlib import nullcontext + +import torch + + +def get_best_device() -> torch.device: + if torch.cuda.is_available(): + return torch.device("cuda") + mps_backend = getattr(torch.backends, "mps", None) + if mps_backend is not None and torch.backends.mps.is_available(): + return torch.device("mps") + return torch.device("cpu") + + +def get_llm_dtype(device: torch.device) -> torch.dtype: + if device.type == "cuda": + return torch.bfloat16 + if device.type == "mps": + return torch.float16 + return torch.float32 + + +def get_amp_context(device: torch.device, enabled: bool): + if not enabled: + return nullcontext() + if device.type == "cuda": + return torch.amp.autocast("cuda", dtype=torch.float16) + if device.type == "mps": + return torch.amp.autocast("mps", dtype=torch.float16) + return nullcontext()