""" Kimi-Audio ZeroGPU Space — audio-only. ASR, audio Q&A, and speech generation via moonshotai/Kimi-Audio-7B-Instruct. Vision (Kimi-VL) is a separate Space: hf_spaces/kimi-vl-zero-gpu-test/ ZeroGPU notes: - Load only this model in this Space (~14 GB with detokenizer). - Do NOT install flash-attn at runtime; SDPA fallback is slower but reliable. - Pin transformers<5; v5 breaks Kimi-Audio remote modeling code. """ from __future__ import annotations import os import subprocess import sys import tempfile import warnings from typing import Any, Optional, Tuple warnings.filterwarnings("ignore") # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- KIMI_AUDIO_REPO = os.getenv( "KIMI_AUDIO_REPO", "https://github.com/MoonshotAI/Kimi-Audio.git" ) KIMI_AUDIO_DIR = os.getenv("KIMI_AUDIO_DIR", "/tmp/Kimi-Audio") KIMI_AUDIO_MODEL = os.getenv( "KIMI_AUDIO_MODEL", "moonshotai/Kimi-Audio-7B-Instruct" ) LOAD_DETOKENIZER = os.getenv("KIMI_LOAD_DETOKENIZER", "1") == "1" # --------------------------------------------------------------------------- # Kimi-Audio install (once per Space container) # --------------------------------------------------------------------------- def _ensure_kimia_infer(): """Clone Moonshot repo with submodules and install editable.""" try: from kimia_infer.api.kimia import KimiAudio # type: ignore return KimiAudio except ImportError: pass if not os.path.isdir(KIMI_AUDIO_DIR): print(f"📦 Cloning Kimi-Audio → {KIMI_AUDIO_DIR}") subprocess.run( [ "git", "clone", "--recursive", "--depth", "1", KIMI_AUDIO_REPO, KIMI_AUDIO_DIR, ], check=True, ) if KIMI_AUDIO_DIR not in sys.path: sys.path.insert(0, KIMI_AUDIO_DIR) # ------------------------------------------------------------------ # Install a prebuilt flash-attn wheel that matches this Space: # torch 2.8.0+cu128 + Python 3.12 + Linux x86_64 # ------------------------------------------------------------------ print("⚡ Installing prebuilt flash-attn (no compile)...") wheel_url = ( "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/" "download/v0.4.11/flash_attn-2.8.3+cu128torch2.8-cp312-cp312-linux_x86_64.whl" ) fa_result = subprocess.run( [sys.executable, "-m", "pip", "install", "-q", wheel_url], capture_output=True, text=True, ) if fa_result.returncode != 0: print("⚠️ flash-attn wheel install failed:") print(fa_result.stderr[-1500:] if fa_result.stderr else fa_result.stdout) # We continue anyway — the later import will tell us if it is fatal else: print("✅ flash-attn wheel installed") print("🎵 Installing kimia_infer (editable, --no-deps)...") result = subprocess.run( [ sys.executable, "-m", "pip", "install", "-q", "-e", KIMI_AUDIO_DIR, "--no-deps", "--no-build-isolation", ], capture_output=True, text=True, ) if result.returncode != 0: print("pip install -e failed:") print(result.stderr[-1500:] if result.stderr else result.stdout) try: from kimia_infer.api.kimia import KimiAudio # type: ignore print("✅ kimia_infer imported successfully") return KimiAudio except ImportError as e: print(f"❌ Still cannot import kimia_infer: {e}") print(f"sys.path[0:5] = {sys.path[:5]}") subprocess.run(["ls", "-la", KIMI_AUDIO_DIR], check=False) raise def _apply_transformers_patches() -> None: """Bridge Kimi-Audio remote code with common HF Space transformers builds.""" import torch import transformers version = transformers.__version__ print(f"transformers {version}") if version.startswith("5."): raise RuntimeError( f"transformers {version} is incompatible with Kimi-Audio. " "Pin transformers==4.44.1 in requirements.txt and restart the Space." ) # EncoderDecoderCache — absent in some 4.x builds try: import transformers.cache_utils as cache_utils if not hasattr(cache_utils, "EncoderDecoderCache"): class EncoderDecoderCache(torch.nn.Module): def __init__(self, self_attention_cache, cross_attention_cache): super().__init__() self.self_attention_cache = self_attention_cache self.cross_attention_cache = cross_attention_cache def get_seq_length(self, layer_idx=0): return self.self_attention_cache.get_seq_length(layer_idx) def to(self, device): self.self_attention_cache.to(device) self.cross_attention_cache.to(device) return self cache_utils.EncoderDecoderCache = EncoderDecoderCache sys.modules["transformers.cache_utils"].EncoderDecoderCache = ( EncoderDecoderCache ) print("✅ Patched EncoderDecoderCache") except Exception as exc: print(f"⚠️ EncoderDecoderCache patch skipped: {exc}") # apply_rotary_pos_emb signature drift (qwen2 backbone) try: from transformers.models.qwen2 import modeling_qwen2 original = modeling_qwen2.apply_rotary_pos_emb def patched_apply_rotary_pos_emb(q, k, cos, sin, *args, **kwargs): return original(q, k, cos, sin, unsqueeze_dim=1) modeling_qwen2.apply_rotary_pos_emb = patched_apply_rotary_pos_emb for name, module in list(sys.modules.items()): if "modeling_moonshot_kimia" in name and hasattr( module, "apply_rotary_pos_emb" ): module.apply_rotary_pos_emb = patched_apply_rotary_pos_emb print("✅ Patched apply_rotary_pos_emb") except Exception as exc: print(f"⚠️ apply_rotary_pos_emb patch skipped: {exc}") # Deferred imports after optional install import gradio as gr import spaces import soundfile as sf import torch from huggingface_hub import snapshot_download _apply_transformers_patches() KimiAudio = _ensure_kimia_infer() # --------------------------------------------------------------------------- # Model state # --------------------------------------------------------------------------- class AudioModelHolder: def __init__(self) -> None: self.model: Any = None self.device: str = "cpu" self.loaded: bool = False def unload(self) -> None: self.model = None self.loaded = False self.device = "cpu" if torch.cuda.is_available(): torch.cuda.empty_cache() holder = AudioModelHolder() @spaces.GPU(duration=180) def load_audio_model() -> str: """Download + load Kimi-Audio on ZeroGPU.""" if holder.loaded and holder.model is not None: return f"✅ Kimi-Audio already loaded on {holder.device}" try: print(f"⬇️ Downloading {KIMI_AUDIO_MODEL}...") model_path = snapshot_download( repo_id=KIMI_AUDIO_MODEL, local_dir="./kimi-audio-model", local_dir_use_symlinks=False, resume_download=True, ) device = "cuda" if torch.cuda.is_available() else "cpu" print(f"🚀 Loading Kimi-Audio on {device} (detokenizer={LOAD_DETOKENIZER})...") model = KimiAudio( model_path=model_path, load_detokenizer=LOAD_DETOKENIZER, ) if hasattr(model, "to"): model = model.to(device) holder.model = model holder.device = device holder.loaded = True vram = "" if torch.cuda.is_available(): vram = f" ({torch.cuda.memory_allocated(0) / 1e9:.1f} GB VRAM)" return f"✅ Kimi-Audio loaded on {device}{vram}" except Exception as exc: holder.unload() return f"❌ Audio load failed: {exc}" @spaces.GPU(duration=120) def generate_audio_response( audio_path: str, prompt: str, mode: str, want_speech: bool, ) -> Tuple[str, Optional[str]]: if not holder.loaded or holder.model is None: return "⚠️ Model not loaded. Click ‘Load Kimi-Audio’ first.", None if not audio_path: return "Please upload or record audio.", None # Default instructions per mode if the user left the box empty defaults = { "Transcribe": "Transcribe this audio accurately.", "Describe / Caption": "Describe this audio in detail. Include any speech, music, environmental sounds, mood, and setting.", "Answer questions about the audio": "Answer the user's question about this audio.", "Chat / Reply with voice": "Respond naturally to the audio as a helpful assistant.", } instruction = (prompt or "").strip() or defaults.get(mode, "Respond to the audio.") # Force speech generation only when the user asked for it or chose the chat mode output_type = "both" if (want_speech or mode == "Chat / Reply with voice") else "text" try: messages = [ { "role": "user", "message_type": "text", "content": instruction, }, { "role": "user", "message_type": "audio", "content": audio_path, }, ] sampling_params = { "audio_temperature": 0.8, "audio_top_k": 10, "text_temperature": 0.7, "text_top_k": 5, "audio_repetition_penalty": 1.0, "audio_repetition_window_size": 64, "text_repetition_penalty": 1.0, "text_repetition_window_size": 16, } result = holder.model.generate( messages, **sampling_params, output_type=output_type ) if output_type == "both": wav_output, text_output = result else: wav_output, text_output = None, result audio_file: Optional[str] = None if wav_output is not None and output_type == "both": with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as handle: audio_file = handle.name if isinstance(wav_output, torch.Tensor): wav_output = wav_output.detach().cpu().view(-1).numpy() sf.write(audio_file, wav_output, 24000) return str(text_output), audio_file except Exception as exc: return f"❌ Inference error: {exc}", None def unload_model() -> str: holder.unload() return "Unloaded" # --------------------------------------------------------------------------- # UI # --------------------------------------------------------------------------- with gr.Blocks(title="Kimi-Audio • ZeroGPU", theme=gr.themes.Soft()) as demo: gr.Markdown( """ # 🎵 Kimi-Audio on ZeroGPU **Audio understanding + speech conversation** > Best at: transcription, describing audio, answering questions about audio, short speech replies. > Music generation is possible but secondary — results vary. """ ) with gr.Row(): load_btn = gr.Button("🚀 Load Kimi-Audio", variant="primary", size="lg") unload_btn = gr.Button("🗑️ Unload", variant="secondary") status = gr.Textbox(label="Status", value="Not loaded", interactive=False) with gr.Row(): with gr.Column(scale=1): audio_in = gr.Audio( label="Upload / Record", sources=["microphone", "upload"], type="filepath", ) mode = gr.Dropdown( label="Mode", choices=[ "Transcribe", "Describe / Caption", "Answer questions about the audio", "Chat / Reply with voice", ], value="Transcribe", ) # Quick presets that fill the instruction box with gr.Row(): preset_transcribe = gr.Button("Transcribe", size="sm") preset_describe = gr.Button("Describe", size="sm") preset_music = gr.Button("Music analysis", size="sm") preset_emotion = gr.Button("Emotion / Mood", size="sm") prompt = gr.Textbox( label="Instruction (optional – leave blank to use mode default)", placeholder="e.g. What instruments are playing? What is the mood?", lines=2, ) want_speech = gr.Checkbox( label="Generate voice reply (needs detokenizer, more VRAM)", value=False, ) run_btn = gr.Button("▶️ Run", variant="primary") with gr.Column(scale=1): text_out = gr.Textbox(label="Text output", lines=8) audio_out = gr.Audio(label="Voice output", type="filepath") # --- Event bindings --- load_btn.click(load_audio_model, outputs=status) unload_btn.click(unload_model, outputs=status) # Preset buttons just fill the prompt box preset_transcribe.click( lambda: "Transcribe this audio accurately.", outputs=prompt, ) preset_describe.click( lambda: "Describe this audio in detail. Include any speech, music, sounds, mood, and setting.", outputs=prompt, ) preset_music.click( lambda: "Analyze the music: genre, instruments, tempo, mood, and structure.", outputs=prompt, ) preset_emotion.click( lambda: "What emotion or mood does this audio convey? Explain briefly.", outputs=prompt, ) run_btn.click( generate_audio_response, inputs=[audio_in, prompt, mode, want_speech], outputs=[text_out, audio_out], ) gr.Markdown( """ --- **Tips** - First load downloads ~7 GB. Allow 2–3 minutes. - For pure text answers turn off “Generate voice reply”. - Set Space secret `KIMI_LOAD_DETOKENIZER=0` if you only need text and keep hitting OOM. """ ) demo.queue(max_size=1) demo.launch(ssr_mode=False, show_error=True)