Qwen2.5-Omni 7B

"""Qwen2.5-Omni 7B, generated from a declarative template."""

from __future__ import annotations

import base64
import io
import tempfile

from pydantic import BaseModel, ConfigDict, Field

from runware_serverless import serve

MODEL = "Qwen/Qwen2.5-Omni-7B"
REVISION = "ae9e1690543ffd5c0221dc27f79834d0294cba00"
CACHE = "/weights"
SPEAKER = "Chelsie"
SAMPLE_RATE = 16000


class Request(BaseModel):
    # An input this model does not read is a line the caller believed was doing
    # something, so it is refused rather than dropped.
    model_config = ConfigDict(extra="forbid")

    prompt: str = Field(description="What to ask or say.")
    audio: str | None = Field(
        default=None,
        description="A recording to listen to before answering, as a base64 audio file. Ask about it in the prompt.",
    )
    speak: bool = Field(
        default=False,
        description="Whether to read the answer aloud as well as write it. Off by default, because the voice head is the slow half and most answers do not need it.",
    )
    max_tokens: int = Field(default=512, ge=1, le=2048, description="Upper bound on the written answer.")


class Answered(BaseModel):
    text: str = Field(description="The written answer.")
    audio: str | None = Field(description="The answer read aloud as a base64 WAV, or null where speak was off.")
    sample_rate: int | None = Field(description="The rate the speech was written at, or null where there is none.")


@serve
class OmniAnswer:
    def load(self) -> None:
        if hasattr(self, "model"):
            return

        # Imported here rather than at the top of the file. The build imports
        # this file with the platform's own dependencies, so nothing above
        # exists until a worker runs.
        import torch
        from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor

        self.torch = torch
        self.processor = Qwen2_5OmniProcessor.from_pretrained(
            MODEL, revision=REVISION, cache_dir=CACHE
        )
        # Straight onto the card. A worker's host RAM is half its card's VRAM,
        # so a model assembled on the host and moved has a ceiling below the
        # card's own.
        self.model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
            MODEL,
            revision=REVISION,
            cache_dir=CACHE,
            dtype=torch.bfloat16,
            device_map="cuda",
        )

    def predict(self, request: Request) -> Answered:
        import librosa
        import soundfile

        content: list[dict] = []
        heard = None

        if request.audio is not None:
            # Decoded and resampled here. soundfile reports the file's own rate
            # and the feature extractor is built for one rate only, so skipping
            # this listens at the wrong speed rather than failing.
            raw = base64.b64decode(request.audio)
            heard, rate = soundfile.read(io.BytesIO(raw), dtype="float32", always_2d=True)
            heard = heard.mean(axis=1)
            if rate != SAMPLE_RATE:
                heard = librosa.resample(heard, orig_sr=rate, target_sr=SAMPLE_RATE)
            content.append({"type": "audio", "audio": heard})

        content.append({"type": "text", "text": request.prompt})

        text = self.processor.apply_chat_template(
            [{"role": "user", "content": content}],
            add_generation_prompt=True,
            tokenize=False,
        )
        inputs = self.processor(
            text=text,
            audio=[heard] if heard is not None else None,
            sampling_rate=SAMPLE_RATE,
            return_tensors="pt",
            padding=True,
        ).to("cuda")

        # generation_mode, not return_audio: the older name is deprecated in
        # this release and renaming it is what keeps the talker from running on
        # every request.
        spoken = self.model.generate(
            **inputs,
            speaker=SPEAKER,
            generation_mode="audio" if request.speak else "text",
            thinker_max_new_tokens=request.max_tokens,
        )

        # Two values only when the talker ran. In text mode it returns the
        # sequences alone, so unpacking a pair raises "not enough values to
        # unpack" on every request that did not ask to be read aloud, which is
        # the default one.
        sequences, wavs = spoken if isinstance(spoken, tuple) else (spoken, None)

        # Only the tokens the model added. Decoding the whole sequence returns
        # the chat template with the answer buried at the end of it.
        new = sequences[:, inputs.input_ids.shape[1]:]
        answer = self.processor.batch_decode(new, skip_special_tokens=True)[0]

        spoken = None
        rate_out = None
        if request.speak and wavs is not None and len(wavs):
            track = wavs[0]
            if hasattr(track, "detach"):
                track = track.detach().to("cpu").float().numpy()
            rate_out = SAMPLE_RATE
            with tempfile.NamedTemporaryFile(suffix=".wav") as handle:
                soundfile.write(handle.name, track, rate_out, format="WAV")
                handle.seek(0)
                spoken = base64.b64encode(handle.read()).decode("ascii")

        return Answered(text=answer.strip(), audio=spoken, sample_rate=rate_out)

configVersion: 1
id: qwen-2-5-omni-7b
name: Qwen2.5-Omni 7B
source:
  type: code
  entry: model.py
compute:
  gpuType: rtx-pro-6000
build:
  requirements:
    - torch==2.14.1
    - torchvision==0.29.1
    - transformers==5.21.0
    - accelerate==1.15.0
    - soundfile==0.14.0
    - librosa==0.11.0
volumes:
  - /weights