Optimizing CPU image

server/Dockerfile.cpu

@@ -1,4 +1,4 @@
-FROM pytorch/pytorch:2.1.0-cuda11.8-cudnn8-devel
+FROM python:3.11.7
 ARG DEBIAN_FRONTEND=noninteractive
 
 RUN apt-get update && \
@@ -16,6 +16,5 @@ 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

@@ -2,25 +2,21 @@
 import io
 import os
 import tempfile
-from typing import List, Literal
 import wave
-
-import numpy as np
 import torch
-from fastapi import (
- FastAPI,
- UploadFile,
- Body,
-)
+import numpy as np
+from typing import List
 from pydantic import BaseModel
+
+from fastapi import FastAPI, UploadFile, Body
 from fastapi.responses import StreamingResponse
 
 from TTS.tts.configs.xtts_config import XttsConfig
 from TTS.tts.models.xtts import Xtts
 from TTS.utils.generic_utils import get_user_data_dir
 from TTS.utils.manage import ModelManager
 
-torch.set_num_threads(int(os.environ.get("NUM_THREADS", "2")))
+torch.set_num_threads(int(os.environ.get("NUM_THREADS", os.cpu_count())))
 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.")