import gradio as gr
import torch
import torchaudio
from torchaudio.transforms import Resample
import moviepy.editor as mp
import numpy as np
from denoiser.pretrained import master64  # Import Facebook denoiser pre-trained model
from denoiser.denoiser import Denoiser  # Denoising wrapper

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print("Device - ", DEVICE)

# Load Facebook denoiser model
model = master64()
den = Denoiser(model).to(DEVICE)
den.eval()

def identity(video_path):
    print(video_path)
    
    # Extract audio from the video
    video = mp.VideoFileClip(video_path)
    audio = video.audio
    wav_file = "tmp.wav"
    audio.write_audiofile(wav_file)
    print("Wav stored.")

    # Load audio
    waveform, sr = torchaudio.load(wav_file)
    waveform = waveform.to(DEVICE)

    # Resample if necessary
    target_sr = 48000
    if sr != target_sr:
        resampler = Resample(orig_freq=sr, new_freq=target_sr).to(DEVICE)
        waveform = resampler(waveform)
        sr = target_sr

    # Process audio in chunks to avoid memory issues
    chunk_duration = 10  # seconds
    chunk_size = sr * chunk_duration
    num_chunks = int(np.ceil(waveform.shape[1] / chunk_size))

    enhanced_chunks = []
    for i in range(num_chunks):
        start = i * chunk_size
        end = min((i + 1) * chunk_size, waveform.shape[1])
        chunk = waveform[:, start:end]
        enhanced_chunk = den(chunk.unsqueeze(0))
        enhanced_chunks.append(enhanced_chunk.squeeze(0))

    # Combine enhanced audio
    enhanced_audio = torch.cat(enhanced_chunks, dim=1).cpu()

    # Save enhanced audio
    output_audio_path = "enhanced_aud.wav"
    torchaudio.save(output_audio_path, enhanced_audio, sr)

    # Replace audio in video
    enhanced_audio_clip = mp.AudioFileClip(output_audio_path)
    final_video = video.set_audio(enhanced_audio_clip)
    output_video_path = "output_video.mp4"
    final_video.write_videofile(output_video_path, 
                                codec='libx264', 
                                audio_codec='aac', 
                                temp_audiofile='temp-audio.m4a', 
                                remove_temp=True)
    return output_video_path

demo = gr.Interface(
    fn=identity,
    title="NoNoise - THE BEST AUDIO DENOISER",
    description="NoNoise is the only platform you need for removing all kinds of background noise from your videos!!",
    examples=[['audiopure_og.mov'], ['example.mp4']],
    cache_examples=True,
    inputs=gr.Video(label="Input Video", source="upload"),
    outputs=gr.Video(label="Output Video"),
)
demo.launch()