From 60260031d2939e44202d84013cec3d7189189304 Mon Sep 17 00:00:00 2001
From: ultrasev <i@ultrasev.com>
Date: Thu, 30 May 2024 23:46:02 +0800
Subject: [PATCH] feat: add docker build

---
 .github/workflows/docker-build.yml | 27 ++++++++++++++
 Dockerfile                         | 13 +++++++
 docker-compose.yml                 | 11 ++++++
 main.py                            | 56 ++++++++++++++++++++++++++++++
 requirements.txt                   |  4 +++
 5 files changed, 111 insertions(+)
 create mode 100644 .github/workflows/docker-build.yml
 create mode 100644 Dockerfile
 create mode 100644 docker-compose.yml
 create mode 100644 main.py

diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml
new file mode 100644
index 000000000..a4fdb9bef
--- /dev/null
+++ b/.github/workflows/docker-build.yml
@@ -0,0 +1,27 @@
+name: Build Docker chattts image
+
+on:
+  push:
+    branches:
+      - master
+
+jobs:
+  build-and-push:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v2
+
+    - name: Log in to GitHub Container Registry
+      uses: docker/login-action@v1
+      with:
+        registry: ghcr.io
+        username: ${{ github.actor }}
+        password: ${{ secrets.CR_PAT }}
+
+    - name: Build and push Docker image
+      uses: docker/build-push-action@v2
+      with:
+        context: .
+        file: ./Dockerfile
+        push: true
+        tags: ghcr.io/${{ github.repository_owner }}/chattts:latest
\ No newline at end of file
diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 000000000..0fa161e74
--- /dev/null
+++ b/Dockerfile
@@ -0,0 +1,13 @@
+FROM python:3.8-slim
+
+RUN apt-get update && apt-get install -y gcc libgomp1 && apt-get clean
+
+WORKDIR /app
+
+COPY requirements.txt .
+
+RUN pip3 install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
\ No newline at end of file
diff --git a/docker-compose.yml b/docker-compose.yml
new file mode 100644
index 000000000..069bf2474
--- /dev/null
+++ b/docker-compose.yml
@@ -0,0 +1,11 @@
+version: '3.8'
+
+services:
+  chattts:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    ports:
+      - "8000:8000"
+    environment:
+      - ENV=production
diff --git a/main.py b/main.py
new file mode 100644
index 000000000..77c7504b2
--- /dev/null
+++ b/main.py
@@ -0,0 +1,56 @@
+#!/usr/bin/env python3
+import wave
+import numpy as np
+from fastapi import FastAPI, HTTPException, Depends
+import pydantic
+import ChatTTS
+
+app = FastAPI()
+
+
+class TTSInput(pydantic.BaseModel):
+    text: str
+    output_path: str
+    seed: int = 42
+
+
+def get_chat_model() -> ChatTTS.Chat:
+    chat = ChatTTS.Chat()
+    chat.load_models()
+    return chat
+
+
+@app.post("/tts")
+def tts(input: TTSInput, chat: ChatTTS.Chat = Depends(get_chat_model)):
+    try:
+        texts = [input.text]
+        r = chat.sample_random_speaker(seed=input.seed)
+
+        params_infer_code = {
+            'spk_emb': r,  # add sampled speaker
+            'temperature': .3,  # using customtemperature
+            'top_P': 0.7,  # top P decode
+            'top_K': 20,  # top K decode
+        }
+
+        params_refine_text = {
+            'prompt': '[oral_2][laugh_0][break_6]'
+        }
+
+        wavs = chat.infer(texts,
+                          params_infer_code=params_infer_code,
+                          params_refine_text=params_refine_text, use_decoder=True)
+
+        audio_data = np.array(wavs[0], dtype=np.float32)
+        sample_rate = 24000
+        audio_data = (audio_data * 32767).astype(np.int16)
+
+        with wave.open(input.output_path, "w") as wf:
+            wf.setnchannels(1)
+            wf.setsampwidth(2)
+            wf.setframerate(sample_rate)
+            wf.writeframes(audio_data.tobytes())
+        return {"output_path": input.output_path}
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
diff --git a/requirements.txt b/requirements.txt
index 87868ec9c..32977d3f7 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,3 +5,7 @@ einops
 vector_quantize_pytorch
 transformers~=4.41.1
 vocos
+fastapi
+uvicorn
+starlette
+pydantic