Compatible ARM.

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"

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版本")

api/requirements.txt

@@ -11,3 +11,4 @@ scipy>=1.11.0
 
 # Logging and utilities
 aiofiles>=23.2.0
+requests>=2.31.0

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}")

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

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

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
+gradio
+soundfile
+tqdm

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
+        return output

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
+        return results_dict

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)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)

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()