import sys import os # Monkey-patch for huggingface_hub breaking change removing 'HfFolder' try: import huggingface_hub if not hasattr(huggingface_hub, "HfFolder"): class HfFolder: @staticmethod def get_token(): return None @staticmethod def save_token(token): pass huggingface_hub.HfFolder = HfFolder sys.modules["huggingface_hub"].HfFolder = HfFolder except Exception: pass # Monkey-patch for gradio_client schema parsing bug with Pydantic 2.10+ (TypeError: argument of type 'bool' is not iterable) try: import gradio_client.utils orig_get_type = gradio_client.utils.get_type def safe_get_type(schema): if isinstance(schema, bool): return "Any" return orig_get_type(schema) gradio_client.utils.get_type = safe_get_type if hasattr(gradio_client.utils, "_json_schema_to_python_type"): orig_json_schema_to_python_type = gradio_client.utils._json_schema_to_python_type def safe_json_schema_to_python_type(schema, defs=None): if isinstance(schema, bool): return "Any" return orig_json_schema_to_python_type(schema, defs) gradio_client.utils._json_schema_to_python_type = safe_json_schema_to_python_type except Exception as e: print(f"[MLOps Warning] Failed to apply gradio_client monkey-patch: {e}") # Try importing Hugging Face Spaces for ZeroGPU compatibility try: import spaces has_spaces = True except ImportError: has_spaces = False # Mock spaces decorator for local execution class spaces: @staticmethod def GPU(func=None, **kwargs): if func is not None: return func return lambda f: f import time import tempfile import torch import torchaudio import numpy as np import gradio as gr from utils.config import load_config from models.generator import Generator # 1. Device Setup (Always start globally on CPU to prevent Hugging Face ZeroGPU startup violations) device = torch.device("cpu") print("[Hugging Face Space] Initializing model on CPU...") # 2. Load Configuration and Model CONFIG_PATH = "config.yaml" CHECKPOINT_PATH = "best_model.pth" config = load_config(CONFIG_PATH) target_sr = config.audio.target_sr # 16000 Hz input_sr = config.audio.input_sr # 8000 Hz # Initialize Generator generator = Generator(config) if os.path.exists(CHECKPOINT_PATH): print(f"[Hugging Face Space] Loading lightweight checkpoint: {CHECKPOINT_PATH}") state_dict = torch.load(CHECKPOINT_PATH, map_location=device, weights_only=False) if "generator_state" in state_dict: generator.load_state_dict(state_dict["generator_state"]) else: generator.load_state_dict(state_dict) print("[Hugging Face Space] Checkpoint loaded successfully!") else: print(f"[Hugging Face Space WARNING] Checkpoint not found at '{CHECKPOINT_PATH}'. Using uninitialized weights.") generator = generator.to(device) generator.eval() @spaces.GPU def process_audio(audio_input): """Enhance narrowband audio input (8 kHz) into wideband speech (16 kHz). Args: audio_input: Tuple of (sample_rate, numpy_ndarray) or filepath string from Gradio. Returns: Tuple of (output_audio_tuple_or_path, status_markdown) """ if audio_input is None: return None, "⚠️ **Error:** Please upload an audio file or record speech first." try: t0 = time.time() # Dynamic device selection inside the decorated function run_device = torch.device("cuda" if torch.cuda.is_available() else "cpu") generator.to(run_device) # Handle Gradio Audio Input (tuple or filepath) if isinstance(audio_input, tuple): sr, waveform_np = audio_input if waveform_np.dtype == np.int16: waveform_np = waveform_np.astype(np.float32) / 32768.0 elif waveform_np.dtype == np.int32: waveform_np = waveform_np.astype(np.float32) / 2147483648.0 # Ensure float32 waveform_np = waveform_np.astype(np.float32) # Check channels (channels last -> channels first) if waveform_np.ndim == 2: if waveform_np.shape[1] in [1, 2]: waveform_np = waveform_np.T waveform_tensor = torch.from_numpy(waveform_np) elif waveform_np.ndim == 1: waveform_tensor = torch.from_numpy(waveform_np).unsqueeze(0) else: generator.to(torch.device("cpu")) return None, "⚠️ **Error:** Unsupported audio array dimension." elif isinstance(audio_input, str): waveform_tensor, sr = torchaudio.load(audio_input) else: generator.to(torch.device("cpu")) return None, "⚠️ **Error:** Invalid audio format." if waveform_tensor.numel() == 0: generator.to(torch.device("cpu")) return None, "⚠️ **Error:** Input audio file is empty." # Convert stereo to mono if waveform_tensor.shape[0] > 1: waveform_tensor = torch.mean(waveform_tensor, dim=0, keepdim=True) # Resample input audio to target sample rate (16 kHz) if sr != target_sr: resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=target_sr) waveform_16k = resampler(waveform_tensor) else: waveform_16k = waveform_tensor orig_len_16k = waveform_16k.shape[-1] # Pad to multiple of 256 to fit neural network conv steps from utils.audio import pad_to_multiple waveform_16k = pad_to_multiple(waveform_16k, 256) # Add batch dimension [1, 1, T] x_nb = waveform_16k.unsqueeze(0).to(run_device) # Run Generator Inference with torch.no_grad(): y_fake = generator(x_nb) # Remove batch dimension -> [1, T_out] and slice back to original 16 kHz length y_fake = y_fake.squeeze(0).cpu()[:, :orig_len_16k] # Move model back to CPU to release GPU resource on ZeroGPU generator.to(torch.device("cpu")) # Normalize audio peak to prevent clipping max_val = torch.max(torch.abs(y_fake)) if max_val > 1.0: y_fake = y_fake / max_val # Save output to temporary WAV file temp_dir = tempfile.gettempdir() output_path = os.path.join(temp_dir, f"enhanced_speech_{int(time.time()*1000)}.wav") torchaudio.save(output_path, y_fake, target_sr) processing_time = (time.time() - t0) * 1000.0 audio_duration = y_fake.shape[-1] / target_sr rtf = (processing_time / 1000.0) / max(audio_duration, 1e-5) status_report = f"""### ✅ Speech Restoration Complete! * **Input Sampling Rate:** `{sr} Hz` (Resampled to `{input_sr} Hz` Narrowband) * **Output Sampling Rate:** `{target_sr} Hz` (High-Fidelity Wideband) * **Audio Duration:** `{audio_duration:.2f} seconds` * **Inference Latency:** `{processing_time:.2f} ms` * **Real-Time Factor (RTF):** `{rtf:.4f}` ({f"{1.0/max(rtf, 1e-4):.1f}x Faster than Real-Time" if rtf < 1.0 else "Processing Complete"}) * **Hardware Device:** `{run_device.type.upper()}` """ return output_path, status_report except Exception as e: try: generator.to(torch.device("cpu")) except Exception: pass return None, f"⚠️ **Error during processing:** `{str(e)}`" # 3. Gradio Interface Construction description = """ ### 🎙️ Real-Time Speech Bandwidth Extension & Telephony Restoration Upload a narrowband telephone recording, mobile call, or low-quality microphone audio to extend its bandwidth to **16 kHz Wideband Speech**. **Model Architecture:** Hybrid Time-Frequency GAN (Dual-Path Spectral Encoder + Multi-Head Cross-Attention + 67% Parameter-Squeezed Discriminator Group). """ css = """ footer {visibility: hidden} .output-audio {margin-top: 15px;} """ with gr.Blocks(title="HybridGAN-BWE: Speech Bandwidth Extension") as demo: gr.Markdown("# 🎧 HybridGAN-BWE: Speech Bandwidth Extension") gr.Markdown(description) with gr.Row(): with gr.Column(): audio_in = gr.Audio( label="Input Audio (Narrowband / Telephone / Microphone)", type="filepath", sources=["upload", "microphone"] ) btn_enhance = gr.Button("✨ Restore & Extend Bandwidth", variant="primary", size="lg") with gr.Column(): audio_out = gr.Audio( label="Restored Wideband Output Audio (16 kHz)", type="filepath" ) status_out = gr.Markdown(value="*Upload audio and click 'Restore & Extend Bandwidth' to begin.*") btn_enhance.click( fn=process_audio, inputs=[audio_in], outputs=[audio_out, status_out] ) if __name__ == "__main__": demo.queue().launch()