Merge pull request #507 from transformerlab/add/audio-stt

Add/audio stt

Author: mina-parham

transformerlab/fastchat_openai_api.py

@@ -93,7 +93,7 @@ class ChatCompletionRequest(BaseModel):
     tools: Optional[List[Dict[str, Any]]] = None
 
 
-class AudioRequest(BaseModel):
+class AudioGenerationRequest(BaseModel):
     experiment_id: str
     model: str
     adaptor: Optional[str] = ""
@@ -106,6 +106,14 @@ class AudioRequest(BaseModel):
     voice: Optional[str] = None
     audio_path: Optional[str] = None
 
+class AudioTranscriptionsRequest(BaseModel):
+    experiment_id: str
+    model: str
+    adaptor: Optional[str] = ""
+    audio_path: str
+    # format: str
+    # output_path: str note: probably we set this by ourself
+
 
 class VisualizationRequest(PydanticBaseModel):
     model: str
@@ -506,7 +514,7 @@ async def show_available_models():
 
 
 @router.post("/v1/audio/speech", tags=["audio"])
-async def create_audio_tts(request: AudioRequest):
+async def create_audio_tts(request: AudioGenerationRequest):
     error_check_ret = await check_model(request)
     if error_check_ret is not None:
         if isinstance(error_check_ret, JSONResponse):
@@ -532,6 +540,8 @@ async def create_audio_tts(request: AudioRequest):
         "top_p": request.top_p,
         "audio_path": request.audio_path,
     }
+    gen_params["task"] = "tts"
+
 
     # Add voice parameter if provided
     if request.voice:
@@ -564,6 +574,34 @@ async def upload_audio_reference(experimentId: str, audio: UploadFile = File(...
 
     return JSONResponse({"audioPath": file_path})
 
+@router.post("/v1/audio/transcriptions", tags=["audio"])
+async def create_text_stt(request: AudioTranscriptionsRequest):
+    error_check_ret = await check_model(request)
+    if error_check_ret is not None:
+        if isinstance(error_check_ret, JSONResponse):
+            return error_check_ret
+        elif isinstance(error_check_ret, dict) and "model_name" in error_check_ret.keys():
+            request.model = error_check_ret["model_name"]
+
+    exp_obj = Experiment.get(request.experiment_id)
+    experiment_dir = exp_obj.get_dir()
+    transcription_dir = os.path.join(experiment_dir, "transcriptions")
+    os.makedirs(transcription_dir, exist_ok=True)
+
+    gen_params = {
+        "model": request.model,
+        "audio_path": request.audio_path,
+        "output_path": transcription_dir,
+        #"format": request.format,
+    }
+    gen_params["task"] = "stt"
+    try:
+        content = await generate_completion(gen_params)
+        return content
+    except Exception as e:
+        return create_error_response(ErrorCode.INTERNAL_ERROR, str(e))
+
+
 
 @router.post("/v1/chat/completions", dependencies=[Depends(check_api_key)], tags=["chat"])
 async def create_openapi_chat_completion(request: ChatCompletionRequest):

transformerlab/plugins/mlx_audio_server/index.json

@@ -1,12 +1,20 @@
 {
   "name": "Apple Audio MLX Server",
   "uniqueId": "mlx_audio_server",
-  "description": "A text-to-speech (TTS) library built on Apple's MLX framework, providing efficient speech synthesis on Apple Silicon.",
+  "description": "A text-to-speech (TTS), speech-to-text(STT) library built on Apple's MLX framework, providing efficient speech synthesis on Apple Silicon.",
   "plugin-format": "python",
   "type": "loader",
-  "version": "0.0.7",
-  "supports": ["Text-to-Speech", "Audio"],
-  "model_architectures": ["MLXTextToSpeech", "StyleTTS2"],
+  "version": "0.1.0",
+  "supports": [
+    "Text-to-Speech",
+    "Audio",
+    "Speech-to-Text"
+  ],
+  "model_architectures": [
+    "MLXTextToSpeech",
+    "StyleTTS2",
+    "MLXSpeechToText"
+  ],
   "supported_hardware_architectures": ["mlx"],
   "files": ["main.py", "setup.sh"],
   "setup-script": "setup.sh"

transformerlab/plugins/mlx_audio_server/main.py

@@ -16,11 +16,14 @@
 from fastapi.responses import JSONResponse
 
 from fastchat.serve.model_worker import logger
-from transformerlab.plugin import WORKSPACE_DIR
 
 from mlx_audio.tts.generate import generate_audio
+from mlx_audio.stt.generate import generate
 from datetime import datetime
 
+from lab.dirs import get_experiments_dir, get_workspace_dir
+from werkzeug.utils import secure_filename
+
 worker_id = str(uuid.uuid4())[:8]
 
 from fastchat.serve.base_model_worker import BaseModelWorker  # noqa
@@ -70,78 +73,134 @@ def __init__(
     async def generate(self, params):
         self.call_ct += 1
 
-        text = params.get("text", "")
-        model = params.get("model", None)
-        speed = params.get("speed", 1.0)
-        # file_prefix = params.get("file_prefix", "audio")
-        audio_format = params.get("audio_format", "wav")
-        sample_rate = params.get("sample_rate", 24000)
-        temperature = params.get("temperature", 0.0)
-        top_p = params.get("top_p", 1.0)
-        stream = params.get("stream", False)
-        voice = params.get("voice", None)
-        lang_code = params.get("lang_code", None)
-
-        audio_dir = params.get("audio_dir", None)
-        if not audio_dir:
-            audio_dir = os.path.join(WORKSPACE_DIR, "audio")
-        os.makedirs(name=audio_dir, exist_ok=True)
-
-        # Generate a UUID for this file name:
-        file_prefix = str(uuid.uuid4())
-
-        try:
-            kwargs = {
-                "text": text,
-                "model_path": model,
-                "speed": speed,
-                "file_prefix": os.path.join(audio_dir, file_prefix),
-                "sample_rate": sample_rate,
-                "join_audio": True,
-                "verbose": True,
-                "temperature": temperature,
-                "top_p": top_p,
-                "stream": stream,
-                "voice": voice,
-            }
-            if lang_code:
-                kwargs["lang_code"] = lang_code
-
-            generate_audio(**kwargs)
-
-            # Also save the parameters and metadata used to generate the audio
-            metadata = {
-                "type": "audio",
-                "text": text,
-                "voice": voice,
-                "filename": f"{file_prefix}.{audio_format}",
-                "model": model,
-                "speed": speed,
-                "audio_format": audio_format,
-                "sample_rate": sample_rate,
-                "temperature": temperature,
-                "top_p": top_p,
-                "date": datetime.now().isoformat(),  # Store the real date and time
-            }
-
-            metadata_file = os.path.join(audio_dir, f"{file_prefix}.json")
-            with open(metadata_file, "w") as f:
-                json.dump(metadata, f)
-
-            logger.info(f"Audio successfully generated: {audio_dir}/{file_prefix}.{audio_format}")
-
-            return {
-                "status": "success",
-                "message": f"{audio_dir}/{file_prefix}.{audio_format}",
-            }
-        except Exception:
-            logger.error(f"Error generating audio: {audio_dir}/{file_prefix}.{audio_format}")
+        task = params.get("task")
+        if task == "tts":
+
+            text = params.get("text", "")
+            model = params.get("model", None)
+            speed = params.get("speed", 1.0)
+            file_prefix = secure_filename(params.get("file_prefix", "audio"))
+            audio_format = params.get("audio_format", "wav")
+            sample_rate = params.get("sample_rate", 24000)
+            temperature = params.get("temperature", 0.0)
+            top_p = params.get("top_p", 1.0)
+            stream = params.get("stream", False)
+            voice = params.get("voice", None)
+            lang_code = params.get("lang_code", None)
+            stream = params.get("stream", False)
+            
+            experiment_dir = get_experiments_dir()
+            audio_dir_name = secure_filename(params.get("audio_dir", "audio"))
+            audio_dir = os.path.join(experiment_dir, audio_dir_name)
+            os.makedirs(name=audio_dir, exist_ok=True)
+
+            try:
+                kwargs = {
+                    "text": text,
+                    "model_path": model,
+                    "speed": speed,
+                    "file_prefix": os.path.join(audio_dir, file_prefix),
+                    "sample_rate": sample_rate,
+                    "join_audio": True,
+                    "verbose": True,
+                    "temperature": temperature,
+                    "top_p": top_p,
+                    "stream": stream,
+                    "voice": voice,
+                }
+                if lang_code:
+                    kwargs["lang_code"] = lang_code
+
+                generate_audio(**kwargs)
+
+                # Also save the parameters and metadata used to generate the audio
+                metadata = {
+                    "type": "audio",
+                    "text": text,
+                    "voice": voice,
+                    "filename": f"{file_prefix}.{audio_format}",
+                    "model": model,
+                    "speed": speed,
+                    "audio_format": audio_format,
+                    "sample_rate": sample_rate,
+                    "temperature": temperature,
+                    "top_p": top_p,
+                    "date": datetime.now().isoformat(),  # Store the real date and time
+                }
+
+                metadata_file = os.path.join(audio_dir, f"{file_prefix}.json")
+                with open(metadata_file, "w") as f:
+                    json.dump(metadata, f)
+
+                logger.info(f"Audio successfully generated: {audio_dir}/{file_prefix}.{audio_format}")
+
+                return {
+                    "status": "success",
+                    "message": f"{audio_dir}/{file_prefix}.{audio_format}",
+                }
+            except Exception:
+                logger.error(f"Error generating audio: {audio_dir}/{file_prefix}.{audio_format}")
+                return {
+                    "status": "error",
+                    "message": f"Error generating audio: {audio_dir}/{file_prefix}.{audio_format}",
+                }
+
+        elif task == "stt":
+            audio_path = params.get("audio_path", "")
+            model = params.get("model", None)
+            format = params.get("format", "txt")
+            output_path_name = secure_filename(params.get("output_path", "transcriptions"))
+            transcriptions_dir = os.path.join(get_workspace_dir(), output_path_name)
+            os.makedirs(name=transcriptions_dir, exist_ok=True)
+
+            # Generate a UUID for this file name:
+            file_prefix = str(uuid.uuid4())
+
+            try:
+                generate(
+                    audio_path=audio_path,
+                    model_path=model,
+                    format=format,
+                    output_path=os.path.join(transcriptions_dir, file_prefix),
+                    verbose=True,  # Set to False to disable print messages
+                )
+
+                # Also save the parameters and metadata used to generate the audio
+                metadata = {
+                    "type": "text",
+                    "audio_folder": "uploaded_audio",
+                    "audio_path": audio_path.split("/").pop(),
+                    "filename": f"{file_prefix}.{format}",
+                    "model": model,
+                    "text_format": format,
+                    "date": datetime.now().isoformat(),  # Store the real date and time
+                }
+                metadata_file = os.path.join(transcriptions_dir, f"{file_prefix}.json")
+                with open(metadata_file, "w") as f:
+                    json.dump(metadata, f)
+
+                logger.info(f"Transcription successfully generated: {transcriptions_dir}/{file_prefix}.{format}")
+
+                return {
+                    "status": "success",
+                    "message": f"{transcriptions_dir}/{file_prefix}.{format}",
+                }
+            except Exception:
+                logger.error(f"Error generating transcription: {transcriptions_dir}/{file_prefix}.{format}")
+                return {
+                    "status": "error",
+                    "message": f"Error generating transcription: {transcriptions_dir}/{file_prefix}.{format}",
+                }
+      
+        else:
+            logger.error(f"Unknown task type: {task}")
             return {
                 "status": "error",
-                "message": f"Error generating audio: {audio_dir}/{file_prefix}.{audio_format}",
+                "message": f"Unknown task type: {task}",
             }
 
 
+
 def release_worker_semaphore():
     worker.semaphore.release()
 

transformerlab/routers/experiment/conversations.py

@@ -130,10 +130,13 @@ async def list_audio(experimentId: str):
 
 
 @router.get(path="/download_audio")
-async def download_audio(experimentId: str, filename: str):
+async def download_audio(experimentId: str, filename: str, audioFolder: str = "audio"):
     exp_obj = Experiment.get(experimentId)
     experiment_dir = exp_obj.get_dir()
-    audio_dir = os.path.join(experiment_dir, "audio")
+    
+    # Use the provided audioFolder parameter, defaulting to "audio"
+    audioFolder = secure_filename(audioFolder)
+    audio_dir = os.path.join(experiment_dir, audioFolder)
 
     # now download the audio file
     filename = secure_filename(filename)
@@ -193,3 +196,51 @@ async def delete_audio(experimentId: str, id: str):
         os.remove(audio_path)
 
     return {"message": f"Audio file {id} deleted from experiment {experimentId}"}
+
+@router.get("/list_transcription")
+async def list_transcription(experimentId: str):
+    # Get experiment object and directory
+    exp_obj = Experiment.get(experimentId)
+    experiment_dir = exp_obj.get_dir()
+    transcription_dir = os.path.join(experiment_dir, "transcriptions")
+    os.makedirs(transcription_dir, exist_ok=True)
+
+    # List all .json files in the transcription directory
+    transcription_files_metadata = []
+    for filename in os.listdir(transcription_dir):
+        if filename.endswith(".json"):
+            file_path = os.path.join(transcription_dir, filename)
+            with open(file_path, "r") as f:
+                try:
+                    data = json.load(f)
+                    # Add the file modification time for sorting
+                    data["id"] = filename[:-5]  # Remove .json from the filename
+                    data["file_date"] = os.path.getmtime(file_path)
+                    transcription_files_metadata.append(data)
+                except Exception:
+                    continue
+    transcription_files_metadata.sort(key=lambda x: x["file_date"], reverse=True)
+    return transcription_files_metadata
+
+@router.get("/download_transcription")
+async def download_transcription(experimentId: str, filename: str):
+    exp_obj = Experiment.get(experimentId)
+    experiment_dir = exp_obj.get_dir()
+    text_dir = os.path.join(experiment_dir, "transcriptions")
+    filename = secure_filename(filename)
+    file_path = os.path.join(text_dir, filename)
+    if not os.path.exists(file_path):
+        return {"message": f"Text file {filename} does not exist in experiment {experimentId}"}
+    return FileResponse(path=file_path, filename=filename, media_type="text/plain")
+
+@router.delete("/delete_transcription")
+async def delete_transcription(experimentId: str, id: str):
+    exp_obj = Experiment.get(experimentId)
+    experiment_dir = exp_obj.get_dir()
+    text_dir = os.path.join(experiment_dir, "transcriptions")
+    id = secure_filename(id)
+    text_path = os.path.join(text_dir, id + ".json")
+    if not os.path.exists(text_path):
+        return {"message": f"Text file {id} does not exist in experiment {experimentId}"}
+    os.remove(text_path)
+    return {"message": f"Text file {id} deleted from experiment {experimentId}"}