app.py

import shlex
import subprocess
import spaces
import torch
import os
import shutil
import glob
import gradio as gr

# install packages for mamba
def install_mamba():
    #subprocess.run(shlex.split("pip install torch==2.2.2 torchvision==0.17.2 torchaudio==2.2.2 --index-url https://download.pytorch.org/whl/cu118"))
    subprocess.run(shlex.split("pip install https://github.com/state-spaces/mamba/releases/download/v2.2.2/mamba_ssm-2.2.2+cu122torch2.3cxx11abiFALSE-cp310-cp310-linux_x86_64.whl"))
    #subprocess.run(shlex.split("pip install numpy==1.26.4"))

def clone_github():
    subprocess.run([
        "git", "clone",
        f"https://RoyChao19477:{os.environ['GITHUB_TOKEN']}@github.com/RoyChao19477/for_HF_AVSEMamba.git"
    ])
    # move all files except README.md
    for item in glob.glob("tmp_repo/*"):
        if os.path.basename(item) != "README.md":
            if os.path.isdir(item):
                shutil.move(item, ".")
            else:
                shutil.move(item, os.path.join(".", os.path.basename(item)))

    shutil.rmtree("tmp_repo")

install_mamba()
clone_github()

ABOUT = """
# SEMamba: Speech Enhancement
A Mamba-based model that denoises real-world audio.
Upload or record a noisy clip and click **Enhance** to hear + see its spectrogram.
"""


import torch
import ffmpeg
import torchaudio
import torchaudio.transforms as T
import yaml
import librosa
import librosa.display
import matplotlib
import numpy as np
import soundfile as sf
import matplotlib.pyplot as plt
from models.stfts    import mag_phase_stft, mag_phase_istft
from models.generator import SEMamba
from models.pcs400   import cal_pcs
from ultralytics import YOLO
import supervision as sv

import gradio as gr
import cv2
import os
import tempfile
from ultralytics import YOLO
from moviepy import ImageSequenceClip

# Load face detector
model = YOLO("yolov8n-face.pt").cuda()  # assumes CUDA available


from decord import VideoReader, cpu
from model import AVSEModule
from config import sampling_rate
import spaces

# Load model once globally
ckpt_path = "ckpts/ep215_0906.oat.ckpt"
model = AVSEModule.load_from_checkpoint(ckpt_path)
model.to("cuda")
model.eval()

@spaces.GPU
def run_avse_inference(video_path, audio_path):
    # Load audio
    noisy, _ = sf.read(audio_path, dtype='float32')  # (N, )
    noisy = torch.tensor(noisy).unsqueeze(0)  # (1, N)

    # Load grayscale video
    vr = VideoReader(video_path, ctx=cpu(0))
    frames = vr.get_batch(list(range(len(vr)))).asnumpy()
    bg_frames = np.array([cv2.cvtColor(f, cv2.COLOR_RGB2GRAY) for f in frames]).astype(np.float32) / 255.0
    bg_frames = torch.tensor(bg_frames).unsqueeze(0).unsqueeze(0)  # (1, 1, T, H, W)

    # Combine into input dict (match what model.enhance expects)
    data = {
        "noisy_audio": noisy,
        "video_frames": bg_frames
    }

    with torch.no_grad():
        estimated = model.enhance(data).reshape(-1).cpu().numpy()

    # Save result
    tmp_wav = audio_path.replace(".wav", "_enhanced.wav")
    sf.write(tmp_wav, estimated, samplerate=sampling_rate)

    return tmp_wav


def extract_resampled_audio(video_path, target_sr=16000):
    # Step 1: extract audio via torchaudio
    # (moviepy will still extract it to wav temp file)
    tmp_audio_path = tempfile.mktemp(suffix=".wav")
    subprocess.run(["ffmpeg", "-y", "-i", video_path, "-vn", "-acodec", "pcm_s16le", "-ar", "44100", tmp_audio_path])

    # Step 2: Load and resample
    waveform, sr = torchaudio.load(tmp_audio_path)
    if sr != target_sr:
        resampler = T.Resample(orig_freq=sr, new_freq=target_sr)
        waveform = resampler(waveform)

    # Step 3: Save resampled audio
    resampled_audio_path = tempfile.mktemp(suffix="_16k.wav")
    torchaudio.save(resampled_audio_path, waveform, sample_rate=target_sr)
    return resampled_audio_path

@spaces.GPU
def extract_faces(video_file):
    cap = cv2.VideoCapture(video_file)
    fps = cap.get(cv2.CAP_PROP_FPS)
    frames = []

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        # Inference
        results = model(frame, verbose=False)[0]
        for box in results.boxes:
            # version 1
            # x1, y1, x2, y2 = map(int, box.xyxy[0])
            # version 2
            h, w, _ = frame.shape
            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
            pad_ratio = 0.3  # 30% padding

            dx = (x2 - x1) * pad_ratio
            dy = (y2 - y1) * pad_ratio

            x1 = int(max(0, x1 - dx))
            y1 = int(max(0, y1 - dy))
            x2 = int(min(w, x2 + dx))
            y2 = int(min(h, y2 + dy))
            face_crop = frame[y1:y2, x1:x2]
            if face_crop.size != 0:
                resized = cv2.resize(face_crop, (224, 224))
                frames.append(resized)
            break  # only one face per frame

    cap.release()

    # Save as video
    tmpdir = tempfile.mkdtemp()
    output_path = os.path.join(tmpdir, "face_only_video.mp4")
    #clip = ImageSequenceClip([cv2.cvtColor(f, cv2.COLOR_BGR2RGB) for f in frames], fps=25)
    clip = ImageSequenceClip([cv2.cvtColor(f, cv2.COLOR_BGR2RGB) for f in frames], fps=fps)
    clip.write_videofile(output_path, codec="libx264", audio=False, fps=25)

    # Save audio from original, resampled to 16kHz
    audio_path = os.path.join(tmpdir, "audio_16k.wav")

    # Extract audio using ffmpeg-python (more robust than moviepy)
    ffmpeg.input(video_file).output(
        audio_path,
        ar=16000,  # resample to 16k
        ac=1,      # mono
        format='wav',
        vn=None    # no video
    ).run(overwrite_output=True)




    # ------------------------------- #
    # AVSE models
    noisy = self.load_wav(audio_path)

    vr = VideoReader(output_path, ctx=cpu(0))
    frames = vr.get_batch(list(range(len(vr)))).asnumpy()
    bg_frames = np.array([
        cv2.cvtColor(frames[i], cv2.COLOR_RGB2GRAY) for i in range(len(frames))
    ]).astype(np.float32)
    bg_frames /= 255.0

    enhanced_audio_path = run_avse_inference(output_path, audio_path)


    return output_path, enhanced_audio_path
    #return output_path, audio_path

iface = gr.Interface(
    fn=extract_faces,
    inputs=gr.Video(label="Upload or record your video"),
    outputs=[
        gr.Video(label="Detected Face Only Video"),
        #gr.Audio(label="Extracted Audio (16kHz)", type="filepath"),
        gr.Audio(label="Enhanced Audio", type="filepath")
    ],
    title="Face Detector",
    description="Upload or record a video. We'll crop face regions and return a face-only video and its 16kHz audio."
)

iface.launch()



ckpt = "ckpts/SEMamba_advanced.pth"
cfg_f = "recipes/SEMamba_advanced.yaml"

# load config
with open(cfg_f, 'r') as f:
    cfg = yaml.safe_load(f)


# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = "cuda"
model  = SEMamba(cfg).to(device)
#sdict  = torch.load(ckpt, map_location=device)
#model.load_state_dict(sdict["generator"])
#model.eval()

@spaces.GPU
def enhance(filepath, model_name):
    # Load model based on selection
    ckpt_path = {
        "VCTK-Demand": "ckpts/SEMamba_advanced.pth",
        "VCTK+DNS": "ckpts/vd.pth"
    }[model_name]

    print("Loading:", ckpt_path)
    model.load_state_dict(torch.load(ckpt_path, map_location=device)["generator"])
    model.eval()
    with torch.no_grad():
        # load & resample
        wav, orig_sr = librosa.load(filepath, sr=None)
        noisy_wav = wav.copy()
        if orig_sr != 16000:
            wav = librosa.resample(wav, orig_sr=orig_sr, target_sr=16000)
        x = torch.from_numpy(wav).float().to(device)
        norm = torch.sqrt(len(x)/torch.sum(x**2))
        #x = (x * norm).unsqueeze(0)
        x = (x * norm)

        # split into 4s segments (64000 samples)
        segment_len = 4 * 16000
        chunks = x.split(segment_len)
        enhanced_chunks = []

        for chunk in chunks:
            if len(chunk) < segment_len:
                #pad = torch.zeros(segment_len - len(chunk), device=chunk.device)
                pad = (torch.randn(segment_len - len(chunk), device=chunk.device) * 1e-4)
                chunk = torch.cat([chunk, pad])
            chunk = chunk.unsqueeze(0)

            amp, pha, _ = mag_phase_stft(chunk, 400, 100, 400, 0.3)
            amp2, pha2, _ = model(amp, pha)
            out = mag_phase_istft(amp2, pha2, 400, 100, 400, 0.3)
            out = (out / norm).squeeze(0)
            enhanced_chunks.append(out)

        out = torch.cat(enhanced_chunks)[:len(x)].cpu().numpy()  # trim padding

        # back to original rate
        if orig_sr != 16000:
            out = librosa.resample(out, orig_sr=16000, target_sr=orig_sr)

        # Normalize
        peak = np.max(np.abs(out))
        if peak > 0.05:
            out = out / peak * 0.85

        # write file
        sf.write("enhanced.wav", out, orig_sr)

        # spectrograms
        fig, axs = plt.subplots(1, 2, figsize=(16, 4))

        # noisy
        D_noisy = librosa.stft(noisy_wav, n_fft=512, hop_length=256)
        S_noisy = librosa.amplitude_to_db(np.abs(D_noisy), ref=np.max)
        librosa.display.specshow(S_noisy, sr=orig_sr, hop_length=256, x_axis="time", y_axis="hz", ax=axs[0], vmax=0)
        axs[0].set_title("Noisy Spectrogram")

        # enhanced
        D_clean = librosa.stft(out, n_fft=512, hop_length=256)
        S_clean = librosa.amplitude_to_db(np.abs(D_clean), ref=np.max)
        librosa.display.specshow(S_clean, sr=orig_sr, hop_length=256, x_axis="time", y_axis="hz", ax=axs[1], vmax=0)
        #librosa.display.specshow(S_clean, sr=16000, hop_length=512, x_axis="time", y_axis="hz", ax=axs[1], vmax=0)
        axs[1].set_title("Enhanced Spectrogram")

        plt.tight_layout()

    return "enhanced.wav", fig

#with gr.Blocks() as demo:
#    gr.Markdown(ABOUT)
#    input_audio = gr.Audio(label="Input Audio", type="filepath", interactive=True)
#    enhance_btn = gr.Button("Enhance")
#    output_audio = gr.Audio(label="Enhanced Audio", type="filepath")
#    plot_output = gr.Plot(label="Spectrograms")
#
#    enhance_btn.click(fn=enhance, inputs=input_audio, outputs=[output_audio, plot_output])
#
#demo.queue().launch()

with gr.Blocks() as demo:
    gr.Markdown(ABOUT)
    input_audio = gr.Audio(label="Input Audio", type="filepath", interactive=True)
    model_choice = gr.Radio(
        label="Choose Model (The use of VCTK+DNS is recommended)",
        choices=["VCTK-Demand", "VCTK+DNS"],
        value="VCTK-Demand"
    )
    enhance_btn = gr.Button("Enhance")
    output_audio = gr.Audio(label="Enhanced Audio", type="filepath")
    plot_output = gr.Plot(label="Spectrograms")

    enhance_btn.click(
        fn=enhance,
        inputs=[input_audio, model_choice],
        outputs=[output_audio, plot_output]
    )
    gr.Markdown("**Note**: The current models are trained on 16kHz audio. Therefore, any input audio not sampled at 16kHz will be automatically resampled before enhancement.")

demo.queue().launch()