From 898a5f2e9b026e1e7761d64b38fb6cbb20903da4 Mon Sep 17 00:00:00 2001
From: WeberJulian <julian.weber@hotmail.fr>
Date: Wed, 13 Dec 2023 16:03:18 +0100
Subject: [PATCH] Optimizing CPU image

---
 server/Dockerfile.cpu |  3 +--
 server/main.py        | 14 +++++---------
 2 files changed, 6 insertions(+), 11 deletions(-)

diff --git a/server/Dockerfile.cpu b/server/Dockerfile.cpu
index 9ef9dac..57dc3a7 100644
--- a/server/Dockerfile.cpu
+++ b/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"]
diff --git a/server/main.py b/server/main.py
index 6dbe082..9a6aeb7 100644
--- a/server/main.py
+++ b/server/main.py
@@ -2,17 +2,13 @@
 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
@@ -20,7 +16,7 @@
 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.")