""" Ohm Audio Studio ====================== A professional interface for Qwen2-Audio ASR and TTS models. This application uses Daggr and Gradio to provide a seamless user experience for Voice Design, Voice Cloning, Custom Voice Synthesis, and Automatic Speech Recognition. Author: Ohm Date: 2026 """ import os import gc import base64 import io import logging import numpy as np import torch import torchaudio import soundfile as sf import gradio as gr from typing import Any, Dict, List, Optional, Tuple, Union from dataclasses import dataclass from huggingface_hub import snapshot_download, login from daggr import FnNode, Graph # Configure Logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # --- Configuration --- @dataclass class AppConfig: HF_TOKEN: Optional[str] = os.environ.get('HF_TOKEN') OMP_NUM_THREADS: str = "1" MODEL_SIZES = ["0.6B", "1.7B"] SPEAKERS = [ "Aiden", "Dylan", "Eric", "Ono_anna", "Ryan", "Serena", "Sohee", "Uncle_fu", "Vivian" ] TTS_LANGUAGES = [ "Auto", "English", "Japanese", "Korean", "French", "German", "Spanish", "Portuguese", "Russian" ] ASR_SUPPORTED_LANGUAGES = [ "English", "Arabic", "German", "French", "Spanish", "Portuguese", "Indonesian", "Italian", "Korean", "Russian", "Thai", "Vietnamese", "Japanese", "Turkish", "Hindi", "Malay", "Dutch", "Swedish", "Danish", "Finnish", "Polish", "Czech", "Filipino", "Persian", "Greek", "Romanian", "Hungarian", "Macedonian" ] # Apply Environment Variables os.environ["OMP_NUM_THREADS"] = AppConfig.OMP_NUM_THREADS if AppConfig.HF_TOKEN: login(token=AppConfig.HF_TOKEN) # --- Utilities --- class AudioUtils: """Utilities for audio processing and normalization.""" @staticmethod def title_case_display(s: str) -> str: s = (s or "").strip() s = s.replace("_", " ") return " ".join([w[:1].upper() + w[1:] if w else "" for w in s.split()]) @staticmethod def build_choices_and_map(items: Optional[List[str]]) -> Tuple[List[str], Dict[str, str]]: if not items: return [], {} display = [AudioUtils.title_case_display(x) for x in items] mapping = {d: r for d, r in zip(display, items)} return display, mapping @staticmethod def normalize_audio(wav: np.ndarray, eps: float = 1e-12, clip: bool = True) -> np.float32: """Normalize audio to float32 in [-1, 1] range.""" x = np.asarray(wav) if np.issubdtype(x.dtype, np.integer): info = np.iinfo(x.dtype) if info.min < 0: y = x.astype(np.float32) / max(abs(info.min), info.max) else: mid = (info.max + 1) / 2.0 y = (x.astype(np.float32) - mid) / mid elif np.issubdtype(x.dtype, np.floating): y = x.astype(np.float32) m = np.max(np.abs(y)) if y.size else 0.0 if m > 1.0 + 1e-6: y = y / (m + eps) else: y = x.astype(np.float32) if clip: y = np.clip(y, -1.0, 1.0) if y.ndim > 1: y = np.mean(y, axis=-1).astype(np.float32) return y @staticmethod def process_input(audio_input: Any) -> Optional[Tuple[np.float32, int]]: """ Handles Filepaths, Data URIs (base64), and Numpy arrays. Returns (numpy_float32, sample_rate_int) """ if audio_input is None: return None try: # Handle Path or Base64 if isinstance(audio_input, str): if audio_input.startswith("data:"): try: header, encoded = audio_input.split(",", 1) data = base64.b64decode(encoded) wav, sr = sf.read(io.BytesIO(data)) return AudioUtils.normalize_audio(wav), int(sr) except Exception as e: logger.error(f"Failed to decode base64 audio: {e}") return None if os.path.exists(audio_input): wav_tensor, sr = torchaudio.load(audio_input) wav = wav_tensor.mean(dim=0).numpy() return AudioUtils.normalize_audio(wav), int(sr) else: logger.error(f"Input string is not a file or valid data URI: {audio_input[:50]}...") return None # Handle Tuple (sample_rate, data) or (data, sample_rate) if isinstance(audio_input, tuple) and len(audio_input) == 2: a0, a1 = audio_input if isinstance(a0, int): return AudioUtils.normalize_audio(a1), int(a0) else: return AudioUtils.normalize_audio(a0), int(a1) # Handle Dictionary if isinstance(audio_input, dict): if "name" in audio_input: return AudioUtils.process_input(audio_input["name"]) if "path" in audio_input: return AudioUtils.process_input(audio_input["path"]) if "sampling_rate" in audio_input and "data" in audio_input: return AudioUtils.normalize_audio(audio_input["data"]), int(audio_input["sampling_rate"]) return None except Exception as e: logger.error(f"Audio Processing Error: {e}") return None # --- Model Management --- class ModelManager: """Manages loading and unloading of AI models.""" def __init__(self): self._loaded_models = {} def _get_model_path(self, model_type: str, model_size: str) -> str: """Download/Get model path based on type and size.""" if model_type == "ASR": return "Qwen/Qwen3-ASR-1.7B" return snapshot_download(f"Qwen/Qwen3-TTS-12Hz-{model_size}-{model_type}") def get_model(self, model_type: str, model_size: str): """ Lazy load models. Unloads previous models if VRAM is tight. """ key = (model_type, model_size) if key not in self._loaded_models: logger.info(f"Clearing Cache before loading {model_type}...") self._loaded_models.clear() gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() logger.info(f"Loading Model: {model_type} {model_size}...") device = "cuda" if torch.cuda.is_available() else "cpu" dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32 if model_type == "ASR": from qwen_asr import Qwen3ASRModel self._loaded_models[key] = Qwen3ASRModel.from_pretrained( "Qwen/Qwen3-ASR-1.7B", dtype=dtype, device_map=device, forced_aligner="Qwen/Qwen3-ForcedAligner-0.6B", forced_aligner_kwargs=dict(dtype=dtype, device_map=device), max_inference_batch_size=4, attn_implementation="sdpa", ) else: from qwen_tts import Qwen3TTSModel model_path = self._get_model_path(model_type, model_size) self._loaded_models[key] = Qwen3TTSModel.from_pretrained( model_path, device_map=device, dtype=dtype, token=AppConfig.HF_TOKEN, ) return self._loaded_models[key] # --- Core Service --- class QwenService: """Core service logic connecting the ModelManager and AudioUtils.""" def __init__(self): self.models = ModelManager() def _cleanup_resources(self): gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() def voice_design(self, text, language, voice_description): """Voice Design (Prompt-to-Speech)""" self._cleanup_resources() if not text: return None, "Text required" if not voice_description: return None, "Description required" try: tts = self.models.get_model("VoiceDesign", "1.7B") wavs, sr = tts.generate_voice_design( text=text.strip(), language=language, instruct=voice_description.strip(), non_streaming_mode=True, max_new_tokens=2048, ) return (sr, wavs[0]), "Success" except Exception as e: logger.exception("Voice Design Error") return None, f"Error: {str(e)}" def voice_clone(self, ref_audio, ref_text, target_text, language, use_xvector_only, model_size): """Voice Cloning (Zero-Shot)""" self._cleanup_resources() if not target_text: return None, "Target text required" audio_tuple = AudioUtils.process_input(ref_audio) if audio_tuple is None: return None, "Error: Could not process reference audio. Please upload a valid WAV/MP3." if not use_xvector_only and not ref_text: return None, "Error: Reference text required (or check 'Use x-vector only')" try: tts = self.models.get_model("Base", model_size) wavs, sr = tts.generate_voice_clone( text=target_text.strip(), language=language, ref_audio=audio_tuple, ref_text=ref_text.strip() if ref_text else None, x_vector_only_mode=use_xvector_only, max_new_tokens=2048, ) return (sr, wavs[0]), "Success" except Exception as e: logger.exception("Voice Clone Error") return None, f"Error: {str(e)}" def custom_voice(self, text, language, speaker, instruct, model_size): """Standard TTS""" self._cleanup_resources() if not text: return None, "Text required" try: tts = self.models.get_model("CustomVoice", model_size) wavs, sr = tts.generate_custom_voice( text=text.strip(), language=language, speaker=speaker.lower().replace(" ", "_"), instruct=instruct.strip() if instruct else None, non_streaming_mode=True, max_new_tokens=2048, ) return (sr, wavs[0]), "Success" except Exception as e: logger.exception("Custom Voice Error") return None, f"Error: {str(e)}" def asr(self, audio_upload, lang_disp): """Automatic Speech Recognition""" self._cleanup_resources() if audio_upload is None: return "", "", "No Audio" processed_audio = AudioUtils.process_input(audio_upload) if processed_audio is None: return "", "", "Error processing audio" language = None if lang_disp and lang_disp != "Auto": # Assuming ASR_LANG_MAP is globally available or we rebuild it # For efficiency let's reuse if possible, or rebuild locally _, mapping = AudioUtils.build_choices_and_map(AppConfig.ASR_SUPPORTED_LANGUAGES) language = mapping.get(lang_disp, lang_disp) try: asr_model = self.models.get_model("ASR", "1.7B") results = asr_model.transcribe( audio=processed_audio, language=language, return_time_stamps=False, ) if not isinstance(results, list) or len(results) != 1: return "", "", "Unexpected result format" r = results[0] detected_lang = getattr(r, "language", "") or "" transcribed_text = getattr(r, "text", "") or "" return detected_lang, transcribed_text, "Success" except Exception as e: logger.exception("ASR Error") return "", "", f"Error: {str(e)}" # --- Graph Construction --- # Initialize Service service = QwenService() ASR_LANG_DISPLAY, _ = AudioUtils.build_choices_and_map(AppConfig.ASR_SUPPORTED_LANGUAGES) ASR_LANG_CHOICES = ["Auto"] + ASR_LANG_DISPLAY # Define Nodes voice_design_node = FnNode( fn=service.voice_design, inputs={ "text": gr.Textbox( label="Text to Synthesize (Voice Design)", lines=4, value="Welcome to Ohm Audio Studio. Experience the future of voice design." ), "language": gr.Dropdown( label="Language (Voice Design)", choices=AppConfig.TTS_LANGUAGES, value="Auto" ), "voice_description": gr.Textbox( label="Voice Description (Voice Design)", lines=3, value="A professional, warm and inviting voice with a clear, confident tone." ), }, outputs={ "generated_audio": gr.Audio(label="Generated Audio", type="numpy"), "status": gr.Textbox(label="Status", interactive=False), }, name="Voice Design" ) custom_voice_node = FnNode( fn=service.custom_voice, inputs={ "text": gr.Textbox( label="Text to Synthesize (Custom Voice)", lines=4, value="Welcome to Ohm Audio Studio coverage of the latest in AI audio technology." ), "language": gr.Dropdown( label="Language (Custom Voice)", choices=AppConfig.TTS_LANGUAGES, value="English" ), "speaker": gr.Dropdown( label="Speaker (Custom Voice)", choices=AppConfig.SPEAKERS, value="Ryan" ), "instruct": gr.Textbox( label="Style Instruction (Custom Voice)", lines=2, placeholder="e.g. Happy, Sad", value="Neutral" ), "model_size": gr.Dropdown( label="Model Size (Custom Voice)", choices=AppConfig.MODEL_SIZES, value="1.7B" ), }, outputs={ "tts_audio": gr.Audio(label="Generated Audio", type="numpy"), "status": gr.Textbox(label="Status", interactive=False), }, name="Custom Voice" ) voice_clone_node = FnNode( fn=service.voice_clone, inputs={ "ref_audio": gr.Audio(label="Reference Audio (Voice Clone)", type="filepath"), "ref_text": gr.Textbox(label="Reference Transcript (Voice Clone)", lines=2), "target_text": gr.Textbox(label="Target Text (Voice Clone)", lines=4), "language": gr.Dropdown( label="Language (Voice Clone)", choices=AppConfig.TTS_LANGUAGES, value="Auto" ), "use_xvector_only": gr.Checkbox(label="Use x-vector only (Voice Clone)", value=False), "model_size": gr.Dropdown( label="Model Size (Voice Clone)", choices=AppConfig.MODEL_SIZES, value="1.7B" ), }, outputs={ "cloned_audio": gr.Audio(label="Cloned Audio", type="numpy"), "status": gr.Textbox(label="Status", interactive=False), }, name="Voice Clone" ) asr_node = FnNode( fn=service.asr, inputs={ "audio_upload": gr.Audio( label="Upload Audio (Qwen3 ASR)", type="numpy", sources=["upload", "microphone"] ), "lang_disp": gr.Dropdown( label="Language (Qwen3 ASR)", choices=ASR_LANG_CHOICES, value="Auto" ), }, outputs={ "detected_lang": gr.Textbox(label="Detected Language", interactive=False), "transcription": gr.Textbox(label="Transcription Result", lines=6, interactive=True), "status": gr.Textbox(label="Status", interactive=False), }, name="Qwen3 ASR" ) # Create and Launch Graph graph = Graph( name="Ohm-Audio-Studio", nodes=[voice_design_node, custom_voice_node, voice_clone_node, asr_node] ) if __name__ == "__main__": port = int(os.environ.get("PORT", 7860)) graph.launch(host="0.0.0.0", port=port)