app.py

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()