Spaces:
Running on Zero
Running on Zero
| """ | |
| src/inference.py - Audio8 0.1b inference wrapper (free-tier: ZeroGPU / CPU fallback) | |
| Single model load, BF16 on CUDA, float32 on CPU. 1.7GB total. | |
| """ | |
| import os | |
| import time | |
| import tempfile | |
| from pathlib import Path | |
| from typing import Tuple, Optional | |
| import torch | |
| import soundfile as sf | |
| import librosa | |
| # ZeroGPU support for HF Spaces free tier | |
| try: | |
| import spaces | |
| HAS_SPACES = True | |
| except ImportError: | |
| HAS_SPACES = False | |
| spaces = None | |
| def _gpu(func): | |
| """Wrap with spaces.GPU if available (ZeroGPU), else no-op.""" | |
| if HAS_SPACES: | |
| return spaces.GPU(duration=90)(func) | |
| return func | |
| MODEL_ID = "Audio8/Audio8-TTS-Preview-0.1b" | |
| # Global singleton | |
| _processor = None | |
| _model = None | |
| _device = None | |
| _dtype = None | |
| _sample_rate = None | |
| _model_loaded = False | |
| _load_error: Optional[str] = None | |
| def get_device() -> str: | |
| if torch.cuda.is_available(): | |
| return "cuda" | |
| return "cpu" | |
| def get_dtype(device: str): | |
| if device == "cuda": | |
| # bfloat16 halves VRAM vs float32, fits free 16GB | |
| return torch.bfloat16 | |
| return torch.float32 | |
| def is_model_loaded() -> bool: | |
| return _model_loaded and _model is not None | |
| def get_model_info() -> dict: | |
| return { | |
| "model_id": MODEL_ID, | |
| "device": _device or get_device(), | |
| "dtype": str(_dtype) if _dtype else str(get_dtype(get_device())), | |
| "sample_rate": _sample_rate or 44100, | |
| "loaded": is_model_loaded(), | |
| "error": _load_error, | |
| "cuda_available": torch.cuda.is_available(), | |
| } | |
| def load_model(): | |
| """Load processor + model once. Thread-safe singleton.""" | |
| global _processor, _model, _device, _dtype, _sample_rate, _model_loaded, _load_error | |
| if is_model_loaded(): | |
| return _processor, _model | |
| _device = get_device() | |
| _dtype = get_dtype(_device) | |
| print(f"[inference] Loading {MODEL_ID} on {_device} ({_dtype}) ...") | |
| t0 = time.time() | |
| try: | |
| from transformers import AutoProcessor, AutoModel | |
| _processor = AutoProcessor.from_pretrained( | |
| MODEL_ID, trust_remote_code=True | |
| ) | |
| _model = AutoModel.from_pretrained( | |
| MODEL_ID, trust_remote_code=True, dtype=_dtype | |
| ).eval().to(_device) | |
| _sample_rate = getattr(_model.config, "codec_sample_rate", 44100) | |
| _model_loaded = True | |
| _load_error = None | |
| print(f"[inference] Loaded in {time.time()-t0:.1f}s | sr={_sample_rate} | device={_device}") | |
| # Warmup compile check (optional) | |
| # torch.compile not used on free CPU to avoid high RAM | |
| except Exception as e: | |
| _load_error = str(e) | |
| print(f"[inference] Load failed: {e}") | |
| raise | |
| return _processor, _model | |
| def validate_reference_audio(path: str) -> Tuple[bool, str]: | |
| """Check duration 3-20s, exists, <10MB.""" | |
| if not path or not os.path.exists(path): | |
| return False, "Reference audio file not found." | |
| try: | |
| size_mb = os.path.getsize(path) / (1024 * 1024) | |
| if size_mb > 10: | |
| return False, f"File too large ({size_mb:.1f}MB > 10MB)." | |
| dur = librosa.get_duration(path=path) | |
| if dur < 1.0: | |
| return False, f"Audio too short ({dur:.1f}s). Use 3-15s for best clone." | |
| if dur > 20: | |
| return False, f"Audio too long ({dur:.1f}s). Trim to <20s." | |
| if dur < 3: | |
| # Warning, not blocking | |
| return True, f"Warning: short audio ({dur:.1f}s) may reduce similarity. 5-15s recommended." | |
| return True, "OK" | |
| except Exception as e: | |
| return False, f"Cannot read audio: {e}" | |
| def synthesize( | |
| reference_audio_path: str, | |
| reference_text: str, | |
| text: str, | |
| temperature: float = 0.7, | |
| top_p: float = 0.9, | |
| top_k: int = 50, | |
| max_new_tokens: int = 512, | |
| ) -> Tuple[str, str]: | |
| """ | |
| Returns (output_wav_path, info_message) | |
| Raises on failure. | |
| """ | |
| if not reference_audio_path: | |
| raise ValueError("Please upload a reference audio (3-15s).") | |
| if not reference_text or not reference_text.strip(): | |
| raise ValueError("Reference transcript is required and must match the audio exactly.") | |
| if not text or not text.strip(): | |
| raise ValueError("Please enter text to synthesize.") | |
| if len(text) > 800: | |
| raise ValueError(f"Text too long ({len(text)} chars). Limit 800 characters for free tier.") | |
| ok, msg = validate_reference_audio(reference_audio_path) | |
| if not ok: | |
| raise ValueError(msg) | |
| processor, model = load_model() | |
| device = _device | |
| # Processor expects lists | |
| inputs = processor( | |
| text=[text.strip()], | |
| reference_audio=[reference_audio_path], | |
| reference_text=[reference_text.strip()], | |
| return_tensors="pt", | |
| ) | |
| inputs = {k: v.to(device) for k, v in inputs.items()} | |
| t0 = time.time() | |
| output = model.generate( | |
| **inputs, | |
| max_new_tokens=max_new_tokens, | |
| temperature=temperature, | |
| top_p=top_p, | |
| top_k=top_k, | |
| do_sample=True, | |
| return_dict_in_generate=True, | |
| ) | |
| waveforms, waveform_lengths = model.decode_audio(output.codes) | |
| gen_time = time.time() - t0 | |
| # Save temp wav | |
| tmp = tempfile.mktemp(suffix=".wav") | |
| audio = waveforms[0, : int(waveform_lengths[0])].float().cpu().numpy() | |
| sf.write(tmp, audio, _sample_rate) | |
| dur = len(audio) / _sample_rate | |
| info = f"Generated {dur:.1f}s in {gen_time:.1f}s on {device} | temp={temperature} top_p={top_p} | sr={_sample_rate}" | |
| if msg.startswith("Warning"): | |
| info = msg + " | " + info | |
| return tmp, info | |
| def synthesize_no_clone( | |
| text: str, | |
| temperature: float = 0.7, | |
| top_p: float = 0.9, | |
| top_k: int = 50, | |
| max_new_tokens: int = 512, | |
| ) -> Tuple[str, str]: | |
| """TTS without voice cloning (no reference).""" | |
| if not text or not text.strip(): | |
| raise ValueError("Please enter text.") | |
| processor, model = load_model() | |
| device = _device | |
| inputs = processor(text=[text.strip()], return_tensors="pt") | |
| inputs = {k: v.to(device) for k, v in inputs.items()} | |
| with torch.inference_mode(): | |
| out = model.generate(**inputs, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k, do_sample=True, return_dict_in_generate=True) | |
| wavs, lens = model.decode_audio(out.codes) | |
| tmp = tempfile.mktemp(suffix=".wav") | |
| audio = wavs[0, : int(lens[0])].float().cpu().numpy() | |
| sf.write(tmp, audio, _sample_rate) | |
| return tmp, f"Generated without cloning on {device}" | |