Add CPU image

server/Dockerfile.cpu

@@ -0,0 +1,21 @@
+FROM pytorch/pytorch:2.1.0-cuda11.8-cudnn8-devel
+ARG DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update && \
+    apt-get install --no-install-recommends -y sox libsox-fmt-all curl wget gcc git git-lfs build-essential libaio-dev libsndfile1 ssh ffmpeg && \
+    apt-get clean && apt-get -y autoremove
+
+WORKDIR /app
+COPY requirements_cpu.txt .
+RUN python -m pip install --use-deprecated=legacy-resolver -r requirements_cpu.txt \
+    && python -m pip cache purge
+
+RUN python -m unidic download
+RUN mkdir -p /app/tts_models
+
+COPY main.py .
+ENV USE_CPU=1
+
+ENV NUM_THREADS=2
+EXPOSE 80
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "80"]

server/main.py

@@ -21,7 +21,9 @@
 from TTS.utils.manage import ModelManager
 
 torch.set_num_threads(int(os.environ.get("NUM_THREADS", "2")))
-device = torch.device("cuda")
+device = torch.device("cuda" if os.environ.get("USE_CPU", "0") == "0" else "cpu")
+if not torch.cuda.is_available() and device == "cuda":
+    raise RuntimeError("CUDA device unavailable, please use Dockerfile.cpu instead.") 
 
 custom_model_path = os.environ.get("CUSTOM_MODEL_PATH", "/app/tts_models")
 
@@ -40,7 +42,7 @@
 config = XttsConfig()
 config.load_json(os.path.join(model_path, "config.json"))
 model = Xtts.init_from_config(config)
-model.load_checkpoint(config, checkpoint_dir=model_path, eval=True, use_deepspeed=True)
+model.load_checkpoint(config, checkpoint_dir=model_path, eval=True, use_deepspeed=True if device == "cuda" else False)
 model.to(device)
 print("XTTS Loaded.", flush=True)
 

server/requirements_cpu.txt

@@ -0,0 +1,11 @@
+TTS @ git+https://github.com/coqui-ai/TTS@fa28f99f1508b5b5366539b2149963edcb80ba62
+uvicorn[standard]==0.23.2
+fastapi==0.95.2
+pydantic==1.10.13
+python-multipart==0.0.6
+typing-extensions>=4.8.0
+numpy==1.24.3
+cutlet
+mecab-python3==1.0.6
+unidic-lite==1.0.8
+unidic==1.1.0