Repository: huggingface/transformers
Pull request: [serve] Add speech to text (`/v1/audio/transcriptions`)

# What does this PR do?

This PR:
- Adds `/v1/audio/transcriptions`, the speech to text (transcription) OAI API
- Adds `TimedModel`: models get deleted from memory after a certain amount of time (e.g try launching the server as `transformers serve --enable-cors --log-level debug --model-timeout 10`)

Changed files:
filename: docs/source/en/serving.md
patch: @@ -70,8 +70,13 @@ vllm serve Qwen/Qwen2.5-1.5B-Instruct \
 > [!WARNING]
 > This section is experimental and subject to change in future versions
 
-<!-- TODO: LLMs -> models, after we add audio/image input/output support -->
-You can serve LLMs supported by `transformers` with the `transformers serve` CLI. It spawns a local server that offers a Chat Completion API or a Response API compatible with the OpenAI SDK, which are the _de facto_ standard for LLM conversations. This way, you can use the server from many third party applications, or test it using the `transformers chat` CLI ([docs](conversations.md#chat-cli)).
+You can serve models of diverse modalities supported by `transformers` with the `transformers serve` CLI. It spawns a local server that offers compatibility with the OpenAI SDK, which is the _de facto_ standard for LLM conversations and other related tasks. This way, you can use the server from many third party applications, or test it using the `transformers chat` CLI ([docs](conversations.md#chat-cli)).
+
+The server supports the following REST APIs:
+- `/v1/chat/completions`
+- `/v1/responses`
+- `/v1/audio/transcriptions`
+- `/v1/models`
 
 To launch a server, simply use the `transformers serve` CLI command:
 
@@ -109,7 +114,7 @@ The server is also an MCP client, so it can interact with MCP tools in agentic u
 <!-- TODO: example with a minimal python example, and explain that it is possible to pass a full generation config in the request -->
 
 
-### Usage example 1: apps with local requests (feat. Jan)
+### Usage example 1: chat with local requests (feat. Jan)
 
 This example shows how to use `transformers serve` as a local LLM provider for the [Jan](https://jan.ai/) app. Jan is a ChatGPT-alternative graphical interface, fully running on your machine. The requests to `transformers serve` come directly from the local app -- while this section focuses on Jan, you can extrapolate some instructions to other apps that make local requests.
 
@@ -139,17 +144,17 @@ ssh -N -f -L 8000:localhost:8000 your_server_account@your_server_IP -p port_to_s
 Port forwarding is not Jan-specific: you can use it to connect `transformers serve` running in a different machine with an app of your choice.
 
 
-### Usage example 2: apps with external requests (feat. Cursor)
+### Usage example 2: chat with external requests (feat. Cursor)
 
 This example shows how to use `transformers serve` as a local LLM provider for [Cursor](https://cursor.com/), the popular IDE. Unlike in the previous example, requests to `transformers serve` will come from an external IP (Cursor's server IPs), which requires some additional setup. Furthermore, some of Cursor's requests require [CORS](https://developer.mozilla.org/en-US/docs/Web/HTTP/Guides/CORS), which is disabled by default for security reasons.
 
-To launch our server with CORS enabled, run
+To launch a server with CORS enabled, run
 
 ```shell
 transformers serve --enable-cors
 ```
 
-We'll also need to expose our server to external IPs. A potential solution is to use [`ngrok`](https://ngrok.com/), which has a permissive free tier. After setting up your `ngrok` account and authenticating on your server machine, you run
+You'll also need to expose your server to external IPs. A potential solution is to use [`ngrok`](https://ngrok.com/), which has a permissive free tier. After setting up your `ngrok` account and authenticating on your server machine, you run
 
 ```shell
 ngrok http [port]
@@ -161,7 +166,7 @@ where `port` is the port used by `transformers serve` (`8000` by default). On th
     <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/transformers_serve_ngrok.png"/>
 </h3>
 
-We're now ready to set things up on the app side! In Cursor, while we can't set a new provider, we can change the endpoint for OpenAI requests in the model selection settings. First, navigate to "Settings" > "Cursor Settings", "Models" tab, and expand the "API Keys" collapsible. To set our `transformers serve` endpoint, follow this order:
+You're now ready to set things up on the app side! In Cursor, while you can't set a new provider, you can change the endpoint for OpenAI requests in the model selection settings. First, navigate to "Settings" > "Cursor Settings", "Models" tab, and expand the "API Keys" collapsible. To set your `transformers serve` endpoint, follow this order:
 1. Unselect ALL models in the list above (e.g. `gpt4`, ...);
 2. Add and select the model you want to use (e.g. `Qwen/Qwen3-4B`)
 3. Add some random text to OpenAI API Key. This field won't be used, but it can’t be empty;
@@ -225,3 +230,26 @@ Image URL: https://evalstate-flux1-schnell.hf.space/gradio_api/file=/tmp/gradio/
 
 I have generated an image of a cat on the moon using the Flux 1 Schnell Image Generator. The image is 1024x1024 pixels and was created with 4 inference steps. Let me know if you would like to make any changes or need further assistance!
 ```
+
+### Usage example 4: speech to text transcription (feat. Open WebUI)
+
+This guide shows how to do audio transcription for chat purposes, using `transformers serve` and [Open WebUI](https://openwebui.com/). This guide assumes you have Open WebUI installed on your machine and ready to run. Please refer to the examples above to use the text functionalities of `transformer serve` with Open WebUI -- the instructions are the same.
+
+To start, let's launch the server. Some of Open WebUI's requests require [CORS](https://developer.mozilla.org/en-US/docs/Web/HTTP/Guides/CORS), which is disabled by default for security reasons, so you need to enable it:
+
+```shell
+transformers serve --enable-cors
+```
+
+Before you can speak into Open WebUI, you need to update its settings to use your server for speech to text (STT) tasks. Launch Open WebUI, and navigate to the audio tab inside the admin settings. If you're using Open WebUI with the default ports, [this link (default)](http://localhost:3000/admin/settings/audio) or [this link (python deployment)](http://localhost:8080/admin/settings/audio) will take you there. Do the following changes there:
+1. Change the type of "Speech-to-Text Engine" to "OpenAI";
+2. Update the address to your server's address -- `http://localhost:8000/v1` by default;
+3. Type your model of choice into the "STT Model" field, e.g. `openai/whisper-large-v3` ([available models](https://huggingface.co/models?pipeline_tag=automatic-speech-recognition&sort=trending)).
+
+If you've done everything correctly, the audio tab should look like this
+
+<h3 align="center">
+    <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/transformers_openwebui_stt_settings.png"/>
+</h3>
+
+You're now ready to speak! Open a new chat, utter a few words after hitting the microphone button, and you should see the corresponding text on the chat input after the model transcribes it.

filename: src/transformers/commands/serving.py
patch: @@ -14,24 +14,28 @@
 
 import copy
 import functools
+import gc
+import io
 import json
 import re
+import threading
 import time
 from argparse import ArgumentParser, Namespace
 from dataclasses import dataclass, field
 from threading import Thread
-from typing import Generator, Optional
+from typing import Generator, Optional, Union
 
 from huggingface_hub import ModelInfo, model_info
 
 from transformers.utils.import_utils import (
     is_fastapi_available,
+    is_librosa_available,
     is_openai_available,
     is_pydantic_available,
     is_uvicorn_available,
 )
 
-from .. import LogitsProcessorList, PreTrainedTokenizerFast, TextIteratorStreamer
+from .. import LogitsProcessorList, PreTrainedTokenizerFast, ProcessorMixin, TextIteratorStreamer
 from ..generation.continuous_batching import ContinuousBatchingManager, RequestStatus
 from ..utils import is_torch_available, logging
 from . import BaseTransformersCLICommand
@@ -42,12 +46,16 @@
 
     from transformers import (
         AutoModelForCausalLM,
-        AutoTokenizer,
+        AutoModelForSpeechSeq2Seq,
+        AutoProcessor,
         BitsAndBytesConfig,
         GenerationConfig,
         PreTrainedModel,
     )
 
+if is_librosa_available():
+    import librosa
+
 serve_dependencies_available = (
     is_pydantic_available() and is_fastapi_available() and is_uvicorn_available() and is_openai_available()
 )
@@ -56,6 +64,8 @@
     from fastapi import FastAPI, HTTPException
     from fastapi.middleware.cors import CORSMiddleware
     from fastapi.responses import JSONResponse, StreamingResponse
+    from openai.types.audio.transcription import Transcription
+    from openai.types.audio.transcription_create_params import TranscriptionCreateParamsBase
     from openai.types.chat.chat_completion_chunk import (
         ChatCompletionChunk,
         Choice,
@@ -90,20 +100,28 @@ class TransformersResponseCreateParamsStreaming(ResponseCreateParamsStreaming, t
         OpenAI's ResponseCreateParamsStreaming with an additional field for the generation config (as a json string).
         """
 
-        generation_config: Optional[str]
+        generation_config: str
 
     class TransformersCompletionCreateParamsStreaming(CompletionCreateParamsStreaming, total=False):
         """
-        OpenAI's CompletionCreateParamsStreaming with additional fields for the generation config (as a json string)
-        and the request ID to re-use the previous KV cache.
+        OpenAI's CompletionCreateParamsStreaming with an additional field for the generation config (as a json string).
         """
 
-        generation_config: Optional[str]
-        request_id: Optional[str]
+        generation_config: str
+
+    class TransformersTranscriptionCreateParams(TranscriptionCreateParamsBase, total=False):
+        """
+        OpenAI's TranscriptionCreateParamsBase with an additional field for the generation config (as a json string).
+        """
 
-    # Contrarily to OpenAI's output types, input types are `TypedDict`, which don't have validation
+        file: bytes  # Overwritten -- pydantic isn't happy with `typing.IO[bytes]`, present in the original type
+        generation_config: str
+        stream: Optional[bool] = False
+
+    # Contrarily to OpenAI's output types, input types are `TypedDict`, which don't have built-in validation.
     response_validator = TypeAdapter(TransformersResponseCreateParamsStreaming)
     completion_validator = TypeAdapter(TransformersCompletionCreateParamsStreaming)
+    transcription_validator = TypeAdapter(TransformersTranscriptionCreateParams)
 
     # Define request fields that are not yet used in `transformers serve`. Receiving these fields will raise an
     # HTTPException.
@@ -146,6 +164,14 @@ class TransformersCompletionCreateParamsStreaming(CompletionCreateParamsStreamin
         "user",
         "web_search_options",
     }
+    UNUSED_TRANSCRIPTION_FIELDS = {
+        "chunking_strategy",
+        "include",
+        "language",
+        "prompt",
+        "response_format",
+        "timestamp_granularities",
+    }
 
 
 logger = logging.get_logger(__name__)
@@ -226,10 +252,6 @@ def create_generation_config_from_req(
     if req.get("seed") is not None:
         torch.manual_seed(req["seed"])
 
-    # Sets server-specific defaults, if unset
-    if generation_config.max_new_tokens is None:
-        generation_config.max_new_tokens = 1024
-
     return generation_config
 
 
@@ -247,6 +269,53 @@ def reset(self):
         self.buffer = ""
 
 
+class TimedModel:
+    """
+    A class that holds a PreTrainedModel instance and its associated processor (tokenizer, audio processor, etc.).
+    Automatically deletes the instances after a specified timeout.
+    """
+
+    def __init__(
+        self,
+        model: "PreTrainedModel",
+        timeout_seconds: int,
+        processor: Optional[Union["ProcessorMixin", "PreTrainedTokenizerFast"]] = None,
+    ):
+        self.model = model
+        self._name_or_path = str(model.name_or_path)
+        self.processor = processor
+        self.timeout_seconds = timeout_seconds
+        self._timer = threading.Timer(self.timeout_seconds, self._delete_model)
+        self._timer.start()
+
+    def reset_timer(self):
+        """Reset the timer for the deletion of the instances."""
+        self._timer.cancel()
+        self._timer = threading.Timer(self.timeout_seconds, self._delete_model)
+        self._timer.start()
+
+    def _delete_model(self):
+        """Delete the wrapped model and processor and clean up resources."""
+        if hasattr(self, "model") and self.model is not None:
+            del self.model
+            del self.processor
+            self.model = None
+            self.processor = None
+            gc.collect()
+
+            # Clear CUDA cache if available
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+            logger.info(
+                f"{self._name_or_path} was removed from memory after {self.timeout_seconds} seconds of inactivity"
+            )
+
+    def is_deleted(self):
+        """Check if the instances have been deleted."""
+        return not hasattr(self, "model") or self.model is None
+
+
 @dataclass
 class ServeArguments:
     r"""
@@ -289,6 +358,10 @@ class ServeArguments:
     # Serving settings
     host: str = field(default="localhost", metadata={"help": "Interface the server will listen to."})
     port: int = field(default=8000, metadata={"help": "Port the server will listen to."})
+    model_timeout: int = field(
+        default=300,
+        metadata={"help": "Time in seconds after which a model will be removed from memory."},
+    )
 
     # Other settings
     log_level: str = field(
@@ -357,16 +430,15 @@ def __init__(self, args: ServeArguments):
         cb_logger.setLevel(logging.log_levels[self.args.log_level.lower()])
 
         # Internal state:
-        # 1. Tracks the most recently used model, to prevent reloading the model unnecessarily
-        self.loaded_model: Optional[str] = None
+        # 1. Tracks models in memory, to prevent reloading the model unnecessarily
+        self.loaded_models: dict[str, TimedModel] = {}
         self.running_continuous_batching_manager: Optional[ContinuousBatchingManager] = None
-        self.model: PreTrainedModel
-        self.tokenizer: PreTrainedTokenizerFast
 
         # 2. preserves information about the last call and last KV cache, to determine whether we can reuse the KV
         # cache and avoid re-running prefil
         self.last_messages = None
         self.last_kv_cache = None
+        self.last_text_model = None
 
     def _validate_request(
         self,
@@ -433,10 +505,19 @@ def validate_chat_completion_request(self, request: dict):
             unused_fields=UNUSED_CHAT_COMPLETION_FIELDS,
         )
 
+    def validate_transcription_request(self, request: dict):
+        self._validate_request(
+            request=request,
+            schema=TransformersTranscriptionCreateParams,
+            validator=transcription_validator,
+            unused_fields=UNUSED_TRANSCRIPTION_FIELDS,
+        )
+
     def build_chat_completion_chunk(
         self,
         request_id: Optional[str] = "",
         content: Optional[str] = None,
+        model: Optional[str] = None,
         role: Optional[str] = None,
         finish_reason: Optional[str] = None,
         tool_calls: Optional[list["ChoiceDeltaToolCall"]] = None,
@@ -452,6 +533,8 @@ def build_chat_completion_chunk(
                 The request ID.
             content (`str`, *optional*):
                 Content of the response from the model.
+            model (`str`, *optional*):
+                The model that generated the content.
             role (`str`, *optional*):
                 The role of the next content, until a new role is defined.
             finish_reason (`str`, *optional*):
@@ -465,7 +548,7 @@ def build_chat_completion_chunk(
         chunk = ChatCompletionChunk(
             id=request_id,
             created=int(time.time()),
-            model=self.loaded_model,
+            model=model,
             choices=[
                 Choice(
                     delta=ChoiceDelta(
@@ -529,6 +612,26 @@ def responses(request: dict):
             output = self.generate_response(request)
             return StreamingResponse(output, media_type="text/event-stream")
 
+        from fastapi import Request
+
+        @app.post("/v1/audio/transcriptions")
+        async def audio_transcriptions(request: Request):
+            # Parses the multipart/form-data request into the request format used by other endpoints
+            async with request.form() as form:
+                parsed_request = TransformersTranscriptionCreateParams(
+                    file=await form["file"].read(),
+                    model=form["model"],
+                    # TODO: add other fields
+                )
+                logger.debug(
+                    f"Received file: {form['file'].filename}; MIME type: {form['file'].content_type}; "
+                    f"size: {form['file'].size / 1024:.2f} KiB"
+                )
+            self.validate_transcription_request(request=parsed_request)
+
+            output = self.generate_transcription(parsed_request)
+            return StreamingResponse(output, media_type="text/event-stream")
+
         @app.get("/v1/models")
         def get_all_models():
             return JSONResponse(
@@ -579,22 +682,22 @@ def continuous_batching_chat_completion(self, req: dict) -> Generator[str, None,
         Returns:
             `Generator[str, None, None]`: A generator that yields the OpenAI Chat Completion chunks.
         """
-        if self.args.force_model is not None:
-            req["model"] = self.args.force_model
 
-        update_model = self.canonicalized_model_name(req["model"]) != self.loaded_model
-        if update_model:
+        model_id_and_revision = self.process_model_name(req["model"])
+        must_discard_cache = model_id_and_revision != self.last_text_model
+        self.last_text_model = model_id_and_revision
+        if must_discard_cache:
             # When switching models, terminate a continuous batching manager if it is running.
             if self.running_continuous_batching_manager is not None:
                 self.running_continuous_batching_manager.stop(block=True, timeout=2)
                 self.running_continuous_batching_manager = None
-            self.load_model_and_tokenizer(req["model"], self.args)
+        model, tokenizer = self.load_text_model_and_tokenizer(model_id_and_revision)
 
         generation_config = create_generation_config_from_req(
             req,
-            model_generation_config=self.model.generation_config,
-            eos_token_id=self.tokenizer.eos_token_id,
-            pad_token_id=self.tokenizer.pad_token_id,
+            model_generation_config=model.generation_config,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.pad_token_id,
             use_cache=False,
             num_blocks=1,
             block_size=1024,
@@ -604,7 +707,7 @@ def continuous_batching_chat_completion(self, req: dict) -> Generator[str, None,
         )
 
         if self.running_continuous_batching_manager is None:
-            self.running_continuous_batching_manager = self.model.init_continuous_batching(
+            self.running_continuous_batching_manager = model.init_continuous_batching(
                 generation_config=generation_config, streaming=True
             )
 
@@ -614,9 +717,9 @@ def continuous_batching_chat_completion(self, req: dict) -> Generator[str, None,
             self.running_continuous_batching_manager.start()
 
         # TODO (Joao, Lysandre): this should also work with tool support
-        inputs = self.tokenizer.apply_chat_template(
-            req["messages"], return_tensors="pt", add_generation_prompt=True
-        ).to(self.model.device)
+        inputs = tokenizer.apply_chat_template(req["messages"], return_tensors="pt", add_generation_prompt=True).to(
+            model.device
+        )
 
         def stream_chat_completion(_inputs):
             try:
@@ -628,7 +731,7 @@ def stream_chat_completion(_inputs):
 
                 # Emit the assistant role to start the stream. Other chunks won't have a role, as it is implicit
                 # they come from the assistant.
-                yield self.build_chat_completion_chunk(request_id, role="assistant")
+                yield self.build_chat_completion_chunk(request_id, role="assistant", model=model_id_and_revision)
 
                 for result in self.running_continuous_batching_manager:
                     if result.request_id != request_id:
@@ -641,10 +744,14 @@ def stream_chat_completion(_inputs):
 
                     finish_reason = "stop" if result.status == RequestStatus.FINISHED else None
                     if result.status == RequestStatus.FINISHED:
-                        yield self.build_chat_completion_chunk(request_id, finish_reason=finish_reason)
+                        yield self.build_chat_completion_chunk(
+                            request_id, finish_reason=finish_reason, model=model_id_and_revision
+                        )
                         break
                     else:
-                        yield self.build_chat_completion_chunk(request_id=request_id, content=result.next_token)
+                        yield self.build_chat_completion_chunk(
+                            request_id=request_id, content=result.next_token, model=model_id_and_revision
+                        )
 
             except Exception as e:
                 logger.error(str(e))
@@ -662,22 +769,20 @@ def generate_chat_completion(self, req: dict) -> Generator[str, None, None]:
         Returns:
             `Generator[str, None, None]`: A generator that yields the OpenAI Chat Completion chunks.
         """
-        if self.args.force_model is not None:
-            req["model"] = self.args.force_model
-
-        update_model = self.canonicalized_model_name(req["model"]) != self.loaded_model
-        if update_model:
-            self.load_model_and_tokenizer(req["model"], self.args)
-
         # HACK for tiny-agents: it sends a request after the assistant message (???). Let's assume we can't have a
         # request whose last message is from the assistant.
         if req["messages"][-1]["role"] == "assistant":
             return
 
+        model_id_and_revision = self.process_model_name(req["model"])
+        must_discard_cache = model_id_and_revision != self.last_text_model
+        self.last_text_model = model_id_and_revision
+        model, tokenizer = self.load_text_model_and_tokenizer(model_id_and_revision)
+
         # ====== TOOL PREPROCESSING LOGIC ======
         tool_model_family = None
         for supported_model_families in _MODELS_WITH_TOOL_SUPPORT:
-            if supported_model_families in self.model.config.architectures[0].lower():
+            if supported_model_families in model.config.architectures[0].lower():
                 tool_model_family = supported_model_families
                 break
         # TODO: trigger 2 constrained generations after the tool call start token is emitted:
@@ -686,21 +791,19 @@ def generate_chat_completion(self, req: dict) -> Generator[str, None, None]:
         # ====== END OF TOOL PREPROCESSING LOGIC ======
 
         if tool_model_family is not None:
-            text = self.tokenizer.apply_chat_template(
+            text = tokenizer.apply_chat_template(
                 req["messages"], add_generation_prompt=True, tokenize=False, tools=req.get("tools")
             )
         else:
-            text = self.tokenizer.apply_chat_template(req["messages"], add_generation_prompt=True, tokenize=False)
-        inputs = self.tokenizer(text, return_tensors="pt").to(self.model.device)["input_ids"]
+            text = tokenizer.apply_chat_template(req["messages"], add_generation_prompt=True, tokenize=False)
+        inputs = tokenizer(text, return_tensors="pt").to(model.device)["input_ids"]
         request_id = req.get("request_id", "req_0")
 
-        generation_streamer = TextIteratorStreamer(self.tokenizer, skip_special_tokens=True, skip_prompt=True)
-        generation_config = create_generation_config_from_req(
-            req, model_generation_config=self.model.generation_config
-        )
+        generation_streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True, skip_prompt=True)
+        generation_config = create_generation_config_from_req(req, model_generation_config=model.generation_config)
 
         last_kv_cache = None
-        if self.is_continuation(req) and not update_model:
+        if self.is_continuation(req) and not must_discard_cache:
             last_kv_cache = self.last_kv_cache
 
         generation_kwargs = {
@@ -715,7 +818,7 @@ def generate_chat_completion(self, req: dict) -> Generator[str, None, None]:
         def stream_chat_completion(streamer, _request_id):
             # Thin wrapper to save the KV cache after generation
             def generate_with_cache(**kwargs):
-                generate_output = self.model.generate(**kwargs)
+                generate_output = model.generate(**kwargs)
                 self.last_kv_cache = generate_output.past_key_values
 
             thread = Thread(target=generate_with_cache, kwargs=generation_kwargs)
@@ -726,7 +829,7 @@ def generate_with_cache(**kwargs):
 
                 # Emit the assistant role to start the stream. Other chunks won't have a role, as it is implicit
                 # they come from the assistant.
-                yield self.build_chat_completion_chunk(request_id, role="assistant")
+                yield self.build_chat_completion_chunk(request_id, role="assistant", model=model_id_and_revision)
 
                 for result in streamer:
                     # ====== TOOL CALL LOGIC ======
@@ -740,10 +843,13 @@ def generate_with_cache(**kwargs):
                         if result.strip() == _TOOL_CALL_TOKENS[tool_model_family]["end"]:
                             tool_state.reset()
                             yield self.build_chat_completion_chunk(
-                                request_id=_request_id, role=None, finish_reason="tool_calls"
+                                request_id=_request_id,
+                                role=None,
+                                finish_reason="tool_calls",
+                                model=model_id_and_revision,
                             )
-                            continue
 
+                            continue
                         # Inside a tool call
                         if tool_state.inside_tool_call:
                             tool_state.buffer += result
@@ -789,15 +895,17 @@ def generate_with_cache(**kwargs):
                                 )
 
                             yield self.build_chat_completion_chunk(
-                                request_id=_request_id, role=None, tool_calls=[tool]
+                                request_id=_request_id, role=None, tool_calls=[tool], model=model_id_and_revision
                             )
                             continue
                     # ====== END OF TOOL CALL LOGIC ======
 
                     # All non-tool related tokens are emitted as assistant messages. Empty text is skipped.
                     if result != "":
-                        yield self.build_chat_completion_chunk(_request_id, content=result)
-                yield self.build_chat_completion_chunk(_request_id, finish_reason="stop")
+                        yield self.build_chat_completion_chunk(
+                            _request_id, content=result, model=model_id_and_revision
+                        )
+                yield self.build_chat_completion_chunk(_request_id, finish_reason="stop", model=model_id_and_revision)
 
                 thread.join()
             except Exception as e:
@@ -820,24 +928,20 @@ def generate_response(self, req: dict) -> Generator[str, None, None]:
             `Generator[str, None, None]`: A generator that yields the OpenAI Response events.
         """
         # TODO -- Implement non-streaming mode
-        if self.args.force_model is not None:
-            req["model"] = self.args.force_model
-
-        update_model = self.canonicalized_model_name(req["model"]) != self.loaded_model
-        if update_model:
-            self.load_model_and_tokenizer(req["model"], self.args)
+        model_id_and_revision = self.process_model_name(req["model"])
+        must_discard_cache = model_id_and_revision != self.last_text_model
+        self.last_text_model = model_id_and_revision
+        model, tokenizer = self.load_text_model_and_tokenizer(model_id_and_revision)
 
-        text = self.tokenizer.apply_chat_template(req["input"], add_generation_prompt=True, tokenize=False)
-        inputs = self.tokenizer(text, return_tensors="pt").to(self.model.device)["input_ids"]
+        text = tokenizer.apply_chat_template(req["input"], add_generation_prompt=True, tokenize=False)
+        inputs = tokenizer(text, return_tensors="pt").to(model.device)["input_ids"]
         request_id = req.get("previous_response_id", "req_0")
 
-        generation_streamer = TextIteratorStreamer(self.tokenizer, skip_special_tokens=True, skip_prompt=True)
-        generation_config = create_generation_config_from_req(
-            req, model_generation_config=self.model.generation_config
-        )
+        generation_streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True, skip_prompt=True)
+        generation_config = create_generation_config_from_req(req, model_generation_config=model.generation_config)
 
         last_kv_cache = None
-        if self.is_continuation(req) and not update_model:
+        if self.is_continuation(req) and not must_discard_cache:
             last_kv_cache = self.last_kv_cache
 
         generation_kwargs = {
@@ -850,7 +954,12 @@ def generate_response(self, req: dict) -> Generator[str, None, None]:
         }
 
         def stream_response(streamer, _request_id):
-            thread = Thread(target=self.model.generate, kwargs=generation_kwargs)
+            # Thin wrapper to save the KV cache after generation
+            def generate_with_cache(**kwargs):
+                generate_output = model.generate(**kwargs)
+                self.last_kv_cache = generate_output.past_key_values
+
+            thread = Thread(target=generate_with_cache, kwargs=generation_kwargs)
             sequence_number = 0
             output_index = 0
             content_index = 0
@@ -868,7 +977,7 @@ def stream_response(streamer, _request_id):
                         id=f"resp_{request_id}",
                         created_at=created_at,
                         status="queued",
-                        model=self.loaded_model,
+                        model=model_id_and_revision,
                         instructions=req.get("instructions"),
                         text={"format": {"type": "text"}},
                         object="response",
@@ -889,7 +998,7 @@ def stream_response(streamer, _request_id):
                         id=f"resp_{request_id}",
                         created_at=created_at,
                         status="in_progress",
-                        model=self.loaded_model,
+                        model=model_id_and_revision,
                         instructions=req.get("instructions"),
                         text={"format": {"type": "text"}},
                         object="response",
@@ -994,7 +1103,7 @@ def stream_response(streamer, _request_id):
                         id=f"resp_{request_id}",
                         created_at=created_at,
                         status="completed",
-                        model=self.loaded_model,
+                        model=model_id_and_revision,
                         instructions=req.get("instructions"),
                         text={"format": {"type": "text"}},
                         output=[response_output_item_done.item],
@@ -1026,7 +1135,7 @@ def stream_response(streamer, _request_id):
                         id=f"resp_{request_id}",
                         created_at=created_at,
                         status="failed",
-                        model=self.loaded_model,
+                        model=model_id_and_revision,
                         instructions=req.get("instructions"),
                         text={"format": {"type": "text"}},
                         output=[],
@@ -1049,6 +1158,54 @@ def stream_response(streamer, _request_id):
 
         return stream_response(generation_streamer, request_id)
 
+    def generate_transcription(self, req: dict) -> Generator[str, None, None]:
+        """
+        Generates an OpenAI Transcription using the audio file.
+
+        Args:
+            req (`dict`): The request containing the audio file and model information.
+
+        Returns:
+            `Generator[str, None, None]`: A generator that yields the transcription result.
+        """
+        # TODO: implement streaming transcription (currently, it's not streaming)
+        if not is_librosa_available():
+            raise ImportError(
+                "Missing librosa dependency for audio transcription. Please install with `pip install librosa`"
+            )
+        model_id_and_revision = self.process_model_name(req["model"])
+        audio_model, audio_processor = self.load_audio_model_and_processor(model_id_and_revision)
+
+        generation_streamer = TextIteratorStreamer(
+            audio_processor.tokenizer, skip_special_tokens=True, skip_prompt=True
+        )
+        generation_config = create_generation_config_from_req(
+            req, model_generation_config=audio_model.generation_config
+        )
+
+        # Read the binary audio file using librosa
+        model_sampling_rate = audio_processor.feature_extractor.sampling_rate
+        audio_bytes = io.BytesIO(req["file"])
+        audio_array, _ = librosa.load(audio_bytes, sr=model_sampling_rate, mono=True)
+        audio_inputs = audio_processor(audio_array, sampling_rate=model_sampling_rate, return_tensors="pt").to(
+            audio_model.device
+        )
+        audio_inputs["input_features"] = audio_inputs["input_features"].to(audio_model.dtype)
+
+        generation_kwargs = {
+            "streamer": generation_streamer,
+            "generation_config": generation_config,
+            "return_dict_in_generate": True,
+        }
+
+        def _generate_transcription():
+            generated_ids = audio_model.generate(**audio_inputs, **generation_kwargs)
+            transcription_text = audio_processor.batch_decode(generated_ids.sequences, skip_special_tokens=True)[0]
+            transcription = Transcription(text=transcription_text)
+            yield f"{transcription.model_dump_json(exclude_none=True)}"
+
+        return _generate_transcription()
+
     def is_continuation(self, req: dict) -> bool:
         """
         Determines whether the current request is a continuation of the last request. In other words, if it is the
@@ -1108,39 +1265,49 @@ def get_quantization_config(args: ServeArguments) -> Optional["BitsAndBytesConfi
 
         return quantization_config
 
-    def canonicalized_model_name(self, model_id: str) -> str:
+    def process_model_name(self, model_id: str) -> str:
         """
-        Canonicalizes the model name to the format "model_id@revision". If the model_id DOESN'T contain an @, it
-        defaults to "model_id@main".
+        Applies the `force_model` CLI argument and canonicalizes the model name to the format "model_id@revision".
+        If the model_id DOESN'T contain an @, it defaults to "model_id@main".
 
         Args:
             model_id (`str`): The model ID.
 
         Returns:
-            `str`: The canonicalized model name.
+            `str`: The canonicalized model name to be used
         """
+        if self.args.force_model is not None:
+            model_id = self.args.force_model
         if "@" in model_id:
             return model_id
         return f"{model_id}@main"
 
-    def load_model_and_tokenizer(self, model_id_and_revision: str, args: ServeArguments):
+    def _load_model_and_data_processor(
+        self, model_id_and_revision: str, model_cls: type[PreTrainedModel]
+    ) -> tuple[PreTrainedModel, Union[ProcessorMixin, PreTrainedTokenizerFast]]:
         """
-        Loads the model and tokenizer from the given model ID and revision into the ServeCommand instance.
+        Generic method to load a model and a data processor from a model ID and revision, making use of the serve CLI
+        arguments.
 
         Args:
             model_id_and_revision (`str`):
                 The model ID and revision to load.
-            args (`ServeArguments`):
-                The serve arguments. May contain quantization settings, device, etc.
+            model_cls (`type[PreTrainedModel]`):
+                The model class to load.
+
+        Returns:
+            `tuple[PreTrainedModel, Union[ProcessorMixin, PreTrainedTokenizerFast]]`: The loaded model and
+            data processor (tokenizer, audio processor, etc.).
         """
-        logger.warning(f"Loading {model_id_and_revision}")
+        args = self.args
+        logger.info(f"Loading {model_id_and_revision}")
 
         if "@" in model_id_and_revision:
             model_id, revision = model_id_and_revision.split("@", 1)
         else:
             model_id, revision = model_id_and_revision, "main"
 
-        tokenizer = AutoTokenizer.from_pretrained(
+        data_processor = AutoProcessor.from_pretrained(
             model_id,
             revision=revision,
             trust_remote_code=args.trust_remote_code,
@@ -1158,19 +1325,76 @@ def load_model_and_tokenizer(self, model_id_and_revision: str, args: ServeArgume
             "trust_remote_code": args.trust_remote_code,
         }
 
-        model = AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs)
-
-        if model.generation_config.max_new_tokens is not None and model.generation_config.max_new_tokens < 1024:
-            model.generation_config.max_new_tokens = 1024
+        model = model_cls.from_pretrained(model_id, **model_kwargs)
 
         if getattr(model, "hf_device_map", None) is None:
             model = model.to(args.device)
 
-        self.loaded_model = f"{model_id}@{revision}"
+        has_default_max_length = (
+            model.generation_config.max_new_tokens is None and model.generation_config.max_length == 20
+        )
+        has_short_max_new_tokens = (
+            model.generation_config.max_new_tokens is not None and model.generation_config.max_new_tokens < 1024
+        )
+        if has_default_max_length or has_short_max_new_tokens:
+            model.generation_config.max_new_tokens = 1024
 
-        logger.warning(f"Loaded model {self.loaded_model}")
-        self.model = model
-        self.tokenizer = tokenizer
+        logger.info(f"Loaded model {model_id_and_revision}")
+        return model, data_processor
+
+    def load_text_model_and_tokenizer(
+        self, model_id_and_revision: str
+    ) -> tuple[PreTrainedModel, PreTrainedTokenizerFast]:
+        """
+        Loads the text model and tokenizer from the given model ID and revision into the ServeCommand instance.
+
+        Args:
+            model_id_and_revision (`str`):
+                The model ID and revision to load.
+
+        Returns:
+            `tuple[PreTrainedModel, PreTrainedTokenizerFast]`: The loaded text model and tokenizer.
+        """
+        if model_id_and_revision not in self.loaded_models or self.loaded_models[model_id_and_revision].is_deleted():
+            model, tokenizer = self._load_model_and_data_processor(model_id_and_revision, AutoModelForCausalLM)
+            self.loaded_models[model_id_and_revision] = TimedModel(
+                model,
+                timeout_seconds=self.args.model_timeout,
+                processor=tokenizer,
+            )
+        else:
+            self.loaded_models[model_id_and_revision].reset_timer()
+            model = self.loaded_models[model_id_and_revision].model
+            tokenizer = self.loaded_models[model_id_and_revision].processor
+
+        return model, tokenizer
+
+    def load_audio_model_and_processor(self, model_id_and_revision: str) -> tuple[PreTrainedModel, ProcessorMixin]:
+        """
+        Loads the audio model and processor from the given model ID and revision into the ServeCommand instance.
+
+        Args:
+            model_id_and_revision (`str`):
+                The model ID and revision to load.
+
+        Returns:
+            `tuple[PreTrainedModel, ProcessorMixin]`: The loaded audio model and processor.
+        """
+        if model_id_and_revision not in self.loaded_models or self.loaded_models[model_id_and_revision].is_deleted():
+            audio_model, audio_processor = self._load_model_and_data_processor(
+                model_id_and_revision, AutoModelForSpeechSeq2Seq
+            )
+            self.loaded_models[model_id_and_revision] = TimedModel(
+                audio_model,
+                timeout_seconds=self.args.model_timeout,
+                processor=audio_processor,
+            )
+        else:
+            self.loaded_models[model_id_and_revision].reset_timer()
+            audio_model = self.loaded_models[model_id_and_revision].model
+            audio_processor = self.loaded_models[model_id_and_revision].processor
+
+        return audio_model, audio_processor
 
 
 if __name__ == "__main__":

filename: tests/commands/test_serving.py
patch: @@ -63,14 +63,13 @@ def test_build_chat_completion_chunk(self):
         """
         dummy = ServeCommand.__new__(ServeCommand)
         dummy.args = type("Args", (), {})()
-        dummy.loaded_model = "dummy_model@main"
 
         # The keys for these fields must be present in every chunk
         MANDATORY_FIELDS = ["data", "id", "choices", "created", "model", "object", "system_fingerprint"]
 
         # Case 1: most fields are provided
         chunk = ServeCommand.build_chat_completion_chunk(
-            dummy, request_id="req0", content="hello", finish_reason="stop", role="user"
+            dummy, request_id="req0", content="hello", finish_reason="stop", role="user", model="dummy_model@main"
         )
         for field in MANDATORY_FIELDS:
             self.assertIn(field, chunk)
@@ -79,13 +78,13 @@ def test_build_chat_completion_chunk(self):
         )
 
         # Case 2: only the role is provided -- other fields in 'choices' are omitted
-        chunk = dummy.build_chat_completion_chunk(request_id="req0", role="user")
+        chunk = dummy.build_chat_completion_chunk(request_id="req0", role="user", model="dummy_model@main")
         for field in MANDATORY_FIELDS:
             self.assertIn(field, chunk)
         self.assertIn('"choices":[{"delta":{"role":"user"},"index":0}]', chunk)
 
         # Case 3: only the content is provided -- other fields in 'choices' are omitted
-        chunk = dummy.build_chat_completion_chunk(request_id="req0", content="hello")
+        chunk = dummy.build_chat_completion_chunk(request_id="req0", content="hello", model="dummy_model@main")
         for field in MANDATORY_FIELDS:
             self.assertIn(field, chunk)
         self.assertIn('"choices":[{"delta":{"content":"hello"},"index":0}]', chunk)
@@ -96,7 +95,7 @@ def test_build_chat_completion_chunk(self):
             function=ChoiceDeltaToolCallFunction(name="foo_bar", arguments='{"foo1": "bar1", "foo2": "bar2"}'),
             type="function",
         )
-        chunk = dummy.build_chat_completion_chunk(request_id="req0", tool_calls=[tool_call])
+        chunk = dummy.build_chat_completion_chunk(request_id="req0", tool_calls=[tool_call], model="dummy_model@main")
         for field in MANDATORY_FIELDS:
             self.assertIn(field, chunk)
         expected_choices_content = (