Xiaomi-CocktailASR-1-safetensors

Xiaomi-CocktailASR-1 (Safetensors / BF16 & INT8 Ready)

This repository provides the safetensors edition of Ease3/Xiaomi-CocktailASR-1, converted to 4-shard bfloat16 safetensors format (~16.5 GB total) with out-of-the-box support for INT8 mixed-precision quantization (bitsandbytes).

Why this version?

1. Windows Compatibility & Bug Fix: The original upstream checkpoint is a single 33GB pytorch_model.bin zip archive. On Windows platforms running PyTorch, loading large .bin checkpoints causes a critical memory mapping access violation (0xC0000005) due to a PyTorch Windows mmap limitation. This safetensors version resolves the issue completely by using native pread block streaming.
2. Reduced Footprint: Weight storage size is reduced from 33.1 GB (float32 bin) to 16.5 GB (bfloat16 safetensors) without loss of precision.
3. Optimized for 16GB GPUs: Supports INT8 mixed-precision loading (audio_encoder & lm_head in BF16, Qwen3-8B backbone in INT8), fitting comfortably in 9.99 GB VRAM with RTF < 0.5 on modern GPUs (e.g., RTX 5060 Ti / RTX 4080 / RTX 3090).

---

Quickstart

1. Requirements

pip install torch torchaudio transformers accelerate bitsandbytes soundfile

2. INT8 Mixed-Precision Loading (Recommended for 16GB GPUs)

import torch
import soundfile as sf
import torchaudio
from transformers import AutoModel, BitsAndBytesConfig

# Windows-safe audio decoding
def _patched_load(audio, **kwargs):
    wav, sr = sf.read(audio, dtype="float32")
    t = torch.from_numpy(wav)
    return (t.unsqueeze(0) if t.ndim == 1 else t.t()), sr
torchaudio.load = _patched_load

# Configure INT8 quantization (preserve Audio Encoder and LM Head in BF16)
quant_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_threshold=6.0,
    llm_int8_skip_modules=[
        "audio_encoder",
        "audio_encoder.encoder",
        "audio_encoder.in_proj",
        "audio_encoder.out_proj",
        "lm_head",
    ],
)

model = AutoModel.from_pretrained(
    "STIAMO/Xiaomi-CocktailASR-1-safetensors",
    trust_remote_code=True,
    quantization_config=quant_config,
    torch_dtype=torch.bfloat16,
    device_map={"": 0},
    low_cpu_mem_usage=True,
).eval()

# Standard Transcription
text = model("target.wav", "ref_speaker.wav", cot=False)
print("Transcription:", text)

# Chain-of-Thought (CoT) Transcription
text_cot = model("target.wav", "ref_speaker.wav", cot=True)
print("CoT Output:", text_cot)

---

Benchmark Results (NVIDIA RTX 5060 Ti 16GB)

---

License & Citation

This model is distributed under the Apache License 2.0.

Original model by Xiaomi Research:

• Base Model: Ease3/Xiaomi-CocktailASR-1
• GitHub: xiaomi-research/xiaomi-cocktailasr-1