Spaces:
Sleeping
Sleeping
File size: 4,449 Bytes
4df8719 82e11f0 4df8719 82e11f0 4df8719 82e11f0 4df8719 82e11f0 4df8719 82e11f0 4df8719 82e11f0 4df8719 82e11f0 4df8719 82e11f0 4df8719 82e11f0 4df8719 70c0bb9 4df8719 70c0bb9 4df8719 70c0bb9 4df8719 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | 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
|