main.py 6.05 KB
Newer Older
1
import os
2
import logging
Timothy J. Baek's avatar
Timothy J. Baek committed
3
4
5
6
7
8
9
10
11
12
from fastapi import (
    FastAPI,
    Request,
    Depends,
    HTTPException,
    status,
    UploadFile,
    File,
    Form,
)
Timothy J. Baek's avatar
Timothy J. Baek committed
13
14
15

from fastapi.responses import StreamingResponse, JSONResponse, FileResponse

Timothy J. Baek's avatar
Timothy J. Baek committed
16
17
from fastapi.middleware.cors import CORSMiddleware
from faster_whisper import WhisperModel
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
18
19
from pydantic import BaseModel

Timothy J. Baek's avatar
Timothy J. Baek committed
20

Timothy J. Baek's avatar
Timothy J. Baek committed
21
22
23
24
25
26
import requests
import hashlib
from pathlib import Path
import json


Timothy J. Baek's avatar
Timothy J. Baek committed
27
28
29
30
31
32
33
34
35
from constants import ERROR_MESSAGES
from utils.utils import (
    decode_token,
    get_current_user,
    get_verified_user,
    get_admin_user,
)
from utils.misc import calculate_sha256

Timothy J. Baek's avatar
Timothy J. Baek committed
36
37
38
39
40
41
from config import (
    SRC_LOG_LEVELS,
    CACHE_DIR,
    UPLOAD_DIR,
    WHISPER_MODEL,
    WHISPER_MODEL_DIR,
42
    WHISPER_MODEL_AUTO_UPDATE,
Jannik Streidl's avatar
Jannik Streidl committed
43
    DEVICE_TYPE,
Timothy J. Baek's avatar
Timothy J. Baek committed
44
45
    AUDIO_OPENAI_API_BASE_URL,
    AUDIO_OPENAI_API_KEY,
46
    AUDIO_OPENAI_API_MODEL,
Yanyutin753's avatar
Yanyutin753 committed
47
    AUDIO_OPENAI_API_VOICE,
Timothy J. Baek's avatar
Timothy J. Baek committed
48
)
49
50
51

log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["AUDIO"])
Timothy J. Baek's avatar
Timothy J. Baek committed
52
53
54
55
56
57
58
59
60
61

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

Timothy J. Baek's avatar
Timothy J. Baek committed
62

Timothy J. Baek's avatar
Timothy J. Baek committed
63
64
app.state.OPENAI_API_BASE_URL = AUDIO_OPENAI_API_BASE_URL
app.state.OPENAI_API_KEY = AUDIO_OPENAI_API_KEY
65
app.state.OPENAI_API_MODEL = AUDIO_OPENAI_API_MODEL
Yanyutin753's avatar
Yanyutin753 committed
66
app.state.OPENAI_API_VOICE = AUDIO_OPENAI_API_VOICE
Timothy J. Baek's avatar
Timothy J. Baek committed
67

Jannik Streidl's avatar
Jannik Streidl committed
68
69
70
71
# setting device type for whisper model
whisper_device_type = DEVICE_TYPE if DEVICE_TYPE and DEVICE_TYPE == "cuda" else "cpu"
log.info(f"whisper_device_type: {whisper_device_type}")

Timothy J. Baek's avatar
Timothy J. Baek committed
72
73
74
75
SPEECH_CACHE_DIR = Path(CACHE_DIR).joinpath("./audio/speech/")
SPEECH_CACHE_DIR.mkdir(parents=True, exist_ok=True)


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
76
77
78
class OpenAIConfigUpdateForm(BaseModel):
    url: str
    key: str
79
80
    model: str
    speaker: str
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
81
82
83
84
85
86
87


@app.get("/config")
async def get_openai_config(user=Depends(get_admin_user)):
    return {
        "OPENAI_API_BASE_URL": app.state.OPENAI_API_BASE_URL,
        "OPENAI_API_KEY": app.state.OPENAI_API_KEY,
88
        "OPENAI_API_MODEL": app.state.OPENAI_API_MODEL,
Yanyutin753's avatar
Yanyutin753 committed
89
        "OPENAI_API_VOICE": app.state.OPENAI_API_VOICE,
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
90
91
92
93
94
95
96
97
98
99
100
101
    }


@app.post("/config/update")
async def update_openai_config(
    form_data: OpenAIConfigUpdateForm, user=Depends(get_admin_user)
):
    if form_data.key == "":
        raise HTTPException(status_code=400, detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)

    app.state.OPENAI_API_BASE_URL = form_data.url
    app.state.OPENAI_API_KEY = form_data.key
102
    app.state.OPENAI_API_MODEL = form_data.model
Yanyutin753's avatar
Yanyutin753 committed
103
    app.state.OPENAI_API_VOICE = form_data.speaker
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
104
105
106
107
108

    return {
        "status": True,
        "OPENAI_API_BASE_URL": app.state.OPENAI_API_BASE_URL,
        "OPENAI_API_KEY": app.state.OPENAI_API_KEY,
109
        "OPENAI_API_MODEL": app.state.OPENAI_API_MODEL,
Yanyutin753's avatar
Yanyutin753 committed
110
        "OPENAI_API_VOICE": app.state.OPENAI_API_VOICE,
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
111
112
113
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
114
115
@app.post("/speech")
async def speech(request: Request, user=Depends(get_verified_user)):
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
    body = await request.body()
    name = hashlib.sha256(body).hexdigest()

    file_path = SPEECH_CACHE_DIR.joinpath(f"{name}.mp3")
    file_body_path = SPEECH_CACHE_DIR.joinpath(f"{name}.json")

    # Check if the file already exists in the cache
    if file_path.is_file():
        return FileResponse(file_path)

    headers = {}
    headers["Authorization"] = f"Bearer {app.state.OPENAI_API_KEY}"
    headers["Content-Type"] = "application/json"

    r = None
Timothy J. Baek's avatar
Timothy J. Baek committed
131
    try:
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
        r = requests.post(
            url=f"{app.state.OPENAI_API_BASE_URL}/audio/speech",
            data=body,
            headers=headers,
            stream=True,
        )

        r.raise_for_status()

        # Save the streaming content to a file
        with open(file_path, "wb") as f:
            for chunk in r.iter_content(chunk_size=8192):
                f.write(chunk)

        with open(file_body_path, "w") as f:
            json.dump(json.loads(body.decode("utf-8")), f)

        # Return the saved file
        return FileResponse(file_path)

    except Exception as e:
        log.exception(e)
        error_detail = "Open WebUI: Server Connection Error"
        if r is not None:
            try:
                res = r.json()
                if "error" in res:
                    error_detail = f"External: {res['error']['message']}"
            except:
                error_detail = f"External: {e}"

        raise HTTPException(
            status_code=r.status_code if r != None else 500,
            detail=error_detail,
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
167

Timothy J. Baek's avatar
Timothy J. Baek committed
168

Timothy J. Baek's avatar
Timothy J. Baek committed
169
@app.post("/transcriptions")
Timothy J. Baek's avatar
Timothy J. Baek committed
170
171
172
173
def transcribe(
    file: UploadFile = File(...),
    user=Depends(get_current_user),
):
174
    log.info(f"file.content_type: {file.content_type}")
Timothy J. Baek's avatar
Timothy J. Baek committed
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189

    if file.content_type not in ["audio/mpeg", "audio/wav"]:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.FILE_NOT_SUPPORTED,
        )

    try:
        filename = file.filename
        file_path = f"{UPLOAD_DIR}/{filename}"
        contents = file.file.read()
        with open(file_path, "wb") as f:
            f.write(contents)
            f.close()

190
191
192
193
194
195
196
197
198
199
        whisper_kwargs = {
            "model_size_or_path": WHISPER_MODEL,
            "device": whisper_device_type,
            "compute_type": "int8",
            "download_root": WHISPER_MODEL_DIR,
            "local_files_only": not WHISPER_MODEL_AUTO_UPDATE,
        }

        log.debug(f"whisper_kwargs: {whisper_kwargs}")

Self Denial's avatar
Self Denial committed
200
        try:
201
202
            model = WhisperModel(**whisper_kwargs)
        except:
203
            log.warning(
Self Denial's avatar
Self Denial committed
204
205
                "WhisperModel initialization failed, attempting download with local_files_only=False"
            )
206
207
            whisper_kwargs["local_files_only"] = False
            model = WhisperModel(**whisper_kwargs)
Timothy J. Baek's avatar
Timothy J. Baek committed
208
209

        segments, info = model.transcribe(file_path, beam_size=5)
210
        log.info(
Timothy J. Baek's avatar
Timothy J. Baek committed
211
212
213
214
215
216
            "Detected language '%s' with probability %f"
            % (info.language, info.language_probability)
        )

        transcript = "".join([segment.text for segment in list(segments)])

217
        return {"text": transcript.strip()}
Timothy J. Baek's avatar
Timothy J. Baek committed
218
219

    except Exception as e:
220
        log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
221
222
223
224
225

        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )