Spaces:
Sleeping
Sleeping
| from io import BytesIO | |
| from pathlib import Path | |
| from typing import cast | |
| import av | |
| import gradio as gr | |
| import soundfile as sf | |
| import torch | |
| from huggingface_hub import hf_hub_download | |
| from safetensors.torch import load_file | |
| from silero_vad import get_speech_timestamps | |
| from silero_vad.utils_vad import OnnxWrapper | |
| from transformers import AutoFeatureExtractor, AutoTokenizer, SequenceFeatureExtractor | |
| from transformers.models.lfm2.configuration_lfm2 import Lfm2Config | |
| from src.modeling.config import ModelConfig | |
| from src.modeling.lambda_audio import LambdaAudioModel | |
| from src.modeling.lfm2_decoder import AudioLfm2ForCausalLM | |
| from src.modeling.loading import ModelComponents | |
| from src.modeling.zipformer import ZipformerConfig, ZipformerModel | |
| from src.shared.transcription import decode_audio | |
| class SpaceConfig(ModelConfig): | |
| max_audio_seconds: float = 30.0 | |
| max_new_tokens: int = 256 | |
| def load_space_components(token: str) -> ModelComponents: | |
| # --------------------------------------------------------- | |
| # Build the architecture from pinned configs, then load ASR weights. | |
| # Keep FP32 computation to match the existing inference path. | |
| # --------------------------------------------------------- | |
| config = SpaceConfig() | |
| encoder_revision = "5d82479188439e0f7a12d18a2e0ec0716785ab78" | |
| decoder_revision = "473d1e77c805524f4bfca67613146adc084521ec" | |
| encoder_config = ZipformerConfig.from_pretrained(config.encoder_model_id, revision=encoder_revision, token=token) | |
| decoder_config = Lfm2Config.from_pretrained(config.decoder_model_id, revision=decoder_revision, token=token) | |
| decoder_config._attn_implementation = "sdpa" | |
| model = LambdaAudioModel(ZipformerModel(encoder_config), AudioLfm2ForCausalLM(decoder_config)) | |
| checkpoint = hf_hub_download( | |
| "KeisukeMiyamoto/lambda-audio-v6", | |
| "model-fp16.safetensors", | |
| revision="4746ea39b9454d86fb6f4ae5ddda826d5fb38d6a", | |
| token=token, | |
| ) | |
| state = {name.removeprefix("model."): tensor for name, tensor in load_file(checkpoint).items()} | |
| model.load_state_dict(state, strict=True) | |
| model.eval() | |
| return ModelComponents( | |
| model=model, | |
| feature_extractor=cast( | |
| SequenceFeatureExtractor, | |
| AutoFeatureExtractor.from_pretrained(config.encoder_model_id, revision=encoder_revision, token=token), | |
| ), | |
| tokenizer=AutoTokenizer.from_pretrained(config.decoder_model_id, revision=decoder_revision, token=token), | |
| ) | |
| def read_recording(audio_path: str | None, vad_model: OnnxWrapper) -> list[bytes]: | |
| # --------------------------------------------------------- | |
| # Decode and validate audio before requesting a GPU. | |
| # --------------------------------------------------------- | |
| if audio_path is None: | |
| raise gr.Error("音声をアップロードするか、録音してください。") | |
| path = Path(audio_path) | |
| if path.stat().st_size > 200 * 1024 * 1024: | |
| raise gr.Error("200 MB以下の音声ファイルを使用してください。") | |
| audio = path.read_bytes() | |
| config = SpaceConfig() | |
| try: | |
| waveform = decode_audio(audio, config.sampling_rate) | |
| except (av.FFmpegError, IndexError, ValueError) as error: | |
| raise gr.Error("再生できる音声ファイルを使用してください。") from error | |
| if waveform.size <= config.sampling_rate * config.max_audio_seconds: | |
| return [audio] | |
| # --------------------------------------------------------- | |
| # Split longer recordings at speech boundaries, at most 30 seconds. | |
| # Encode lossless clips in chronological order for batched ASR. | |
| # --------------------------------------------------------- | |
| timestamps = get_speech_timestamps( | |
| torch.from_numpy(waveform), | |
| vad_model, | |
| sampling_rate=config.sampling_rate, | |
| max_speech_duration_s=config.max_audio_seconds, | |
| min_silence_duration_ms=3000, | |
| speech_pad_ms=100, | |
| ) | |
| segments: list[bytes] = [] | |
| segment_samples = int(config.max_audio_seconds * config.sampling_rate) | |
| for timestamp in timestamps: | |
| for start in range(timestamp["start"], timestamp["end"], segment_samples): | |
| end = min(start + segment_samples, timestamp["end"]) | |
| buffer = BytesIO() | |
| sf.write(buffer, waveform[start:end], config.sampling_rate, format="WAV", subtype="FLOAT") | |
| segments.append(buffer.getvalue()) | |
| return segments | |