update: review config for gpu machine

.dockerignore

@@ -1 +0,0 @@
-venv

Dockerfile

@@ -1,12 +1,16 @@
-FROM python:3.10-slim
-# FROM pytorch/pytorch:2.4.1-cuda12.4-cudnn9-runtime
+FROM pytorch/pytorch:2.6.0-cuda12.4-cudnn9-runtime
 
 WORKDIR /app
 
-RUN apt-get update && apt-get install -y git ffmpeg
-# RUN sed -i 's|http://archive.ubuntu.com/ubuntu|http://fr.archive.ubuntu.com/ubuntu|g' /etc/apt/sources.list && \
-#     apt-get update && apt-get install -y --no-install-recommends git && \
-#     rm -rf /var/lib/apt/lists/*
+ENV PYTHONUNBUFFERED=1 \
+	HF_HOME=/tmp/huggingface \
+	TRANSFORMERS_CACHE=/tmp/huggingface/transformers \
+	HF_DATASETS_CACHE=/tmp/huggingface/datasets \
+	NVIDIA_VISIBLE_DEVICES=all \
+	NVIDIA_DRIVER_CAPABILITIES=compute,utility
+
+RUN apt-get update && apt-get install -y --no-install-recommends git ffmpeg && \
+	rm -rf /var/lib/apt/lists/*
 
 COPY . .
 
@@ -14,7 +18,7 @@ RUN pip install --no-cache-dir --upgrade pip
 RUN pip install --no-cache-dir -r requirements.txt
 
 # CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
-CMD ["python", "-m", "uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
+CMD ["python", "-m", "uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860", "--workers", "1"]
 
 # docker buildx build --platform linux/amd64 --no-cache -t ai-api .
 # docker tag ai-api kora3/ai-api:latest

README.md

@@ -10,3 +10,7 @@ license: apache-2.0
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+## GPU
+
+L'API utilise CUDA automatiquement quand un GPU NVIDIA est disponible. Les modèles de traduction, MMS et TTS sont alors chargés sur le GPU ; `faster-whisper` utilise CUDA en `float16`.

functions/speech_to_text.py

@@ -1,7 +1,12 @@
 from faster_whisper import WhisperModel
 import base64, tempfile, os
+import torch
 
-model = WhisperModel("base")
+model = WhisperModel(
+    "base",
+    device="cuda" if torch.cuda.is_available() else "cpu",
+    compute_type="float16" if torch.cuda.is_available() else "int8",
+)
 
 
 def speechToText(audioBase64: str, sourceLang: str) -> dict:

functions/translation.py

@@ -1,17 +1,19 @@
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
+from functions.utils import DEVICE, TORCH_DTYPE
 
 MODEL_NAME = "facebook/nllb-200-distilled-600M" # facebook/nllb-200-3.3B
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).to(device)
+model = AutoModelForSeq2SeqLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=TORCH_DTYPE,
+).to(DEVICE)
 model.eval()
 
 
 def translateText(text: str, sourceLang: str, targetLang: str) -> str:
     tokenizer.src_lang = sourceLang
-    inputs = tokenizer(text, return_tensors="pt").to(device)
+    inputs = tokenizer(text, return_tensors="pt").to(DEVICE)
 
     with torch.no_grad():
         tokens = model.generate(

functions/utils.py

@@ -1,4 +1,14 @@
 import soundfile as sf
+import torch
+
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+CUDA_AVAILABLE = DEVICE.type == "cuda"
+TORCH_DTYPE = torch.float16 if CUDA_AVAILABLE else torch.float32
+
+print(f"Using device: {DEVICE}")
+if CUDA_AVAILABLE:
+    print(f"CUDA device: {torch.cuda.get_device_name(DEVICE)}")
 
 
 def getAudioDuration(filePath: str) -> float:

language/dioula/dyu_stt.py

@@ -1,11 +1,14 @@
-import base64, tempfile, os, torch
+import base64, tempfile, os
 from transformers import pipeline
-from functions.utils import getAudioDuration
+from functions.utils import CUDA_AVAILABLE, TORCH_DTYPE, getAudioDuration
 
 MODEL_NAME = "facebook/mms-1b-all"
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-pipe = pipeline("automatic-speech-recognition", model=MODEL_NAME, model_kwargs={"target_lang": "dyu"}, device=device)
+pipe = pipeline(
+    "automatic-speech-recognition",
+    model=MODEL_NAME,
+    model_kwargs={"target_lang": "dyu", "torch_dtype": TORCH_DTYPE},
+    device=0 if CUDA_AVAILABLE else -1,
+)
 
 
 def dioulaSTT(audioBase64: str) -> dict:

language/dioula/dyu_tts.py

@@ -1,14 +1,18 @@
 import torch, base64, tempfile, os
 import scipy.io.wavfile as wavfile
 from transformers import VitsModel, VitsTokenizer
+from functions.utils import DEVICE, TORCH_DTYPE
 
 MODEL_NAME = "facebook/mms-tts-dyu"
 tokenizer = VitsTokenizer.from_pretrained(MODEL_NAME)
-model = VitsModel.from_pretrained(MODEL_NAME)
+model = VitsModel.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=TORCH_DTYPE,
+).to(DEVICE)
 
 
 def dioulaTTS(text: str) -> str:
-    inputs = tokenizer(text, return_tensors="pt")
+    inputs = tokenizer(text, return_tensors="pt").to(DEVICE)
 
     with torch.no_grad():
         output = model(**inputs)

language/fr_mos.py

@@ -1,18 +1,21 @@
 import torch
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from functions.utils import DEVICE, TORCH_DTYPE
 
 MODEL_NAME = "code-li/nllb-moore"  
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).to(device)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForSeq2SeqLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=TORCH_DTYPE,
+).to(DEVICE)
 model.eval()
 
 
 def translateFRMOS(text: str, sourceLang: str, targetLang: str):
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, src_lang=sourceLang)
+    tokenizer.src_lang = sourceLang
     TGT_LANG_ID = tokenizer.convert_tokens_to_ids(targetLang)
 
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128).to(device)
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128).to(DEVICE)
     with torch.no_grad():
         tokens = model.generate(
             **inputs,

language/moore/mos_stt.py

@@ -1,11 +1,14 @@
-import base64, tempfile, os, torch
+import base64, tempfile, os
 from transformers import pipeline
-from functions.utils import getAudioDuration
+from functions.utils import CUDA_AVAILABLE, TORCH_DTYPE, getAudioDuration
 
 MODEL_NAME = "facebook/mms-1b-all"
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-pipe = pipeline("automatic-speech-recognition", model=MODEL_NAME, model_kwargs={"target_lang": "mos"}, device=device)
+pipe = pipeline(
+    "automatic-speech-recognition",
+    model=MODEL_NAME,
+    model_kwargs={"target_lang": "mos", "torch_dtype": TORCH_DTYPE},
+    device=0 if CUDA_AVAILABLE else -1,
+)
 
 # MODEL_NAME = "burkimbia/BIA-WHISPER-LARGE-SACHI_V3"
 # pipe = pipeline("automatic-speech-recognition", model=MODEL_NAME)

language/moore/mos_tts.py

@@ -1,14 +1,18 @@
 import torch, base64, tempfile, os
 import scipy.io.wavfile as wavfile
 from transformers import VitsModel, VitsTokenizer
+from functions.utils import DEVICE, TORCH_DTYPE
 
 MODEL_NAME = "facebook/mms-tts-mos"
 tokenizer = VitsTokenizer.from_pretrained(MODEL_NAME)
-model = VitsModel.from_pretrained(MODEL_NAME)
+model = VitsModel.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=TORCH_DTYPE,
+).to(DEVICE)
 
 
 def mooreTTS(text: str) -> str:
-    inputs = tokenizer(text, return_tensors="pt")
+    inputs = tokenizer(text, return_tensors="pt").to(DEVICE)
 
     with torch.no_grad():
         output = model(**inputs)

main.py

@@ -13,7 +13,8 @@ from language.fr_mos import translateFRMOS
 
 import os
 from huggingface_hub import login
-login(token=os.environ["HF_TOKEN"])
+if token := os.getenv("HF_TOKEN"):
+    login(token=token)
 
 
 

requirements.txt

@@ -1,3 +1,4 @@
+--extra-index-url https://download.pytorch.org/whl/cu124
 absl-py==2.4.0
 aiohappyeyeballs==2.6.2
 aiohttp==3.14.1
@@ -128,15 +129,15 @@ srsly==2.5.3
 starlette==1.2.1
 SudachiDict-core==20260428
 SudachiPy==0.6.11
-sympy==1.14.0
+sympy==1.13.1
+torch==2.6.0
+torchaudio==2.6.0
 tabulate==0.10.0
 tensorboard==2.20.0
 tensorboard-data-server==0.7.2
 thinc==8.3.13
 threadpoolctl==3.6.0
 tokenizers==0.22.2
-torch==2.12.0
-torchaudio==2.11.0
 tqdm==4.68.1
 trainer==0.0.36
 transformers==5.12.1