Adapt Space for ZeroGPU

app.py

@@ -2,17 +2,14 @@ import argparse
 from functools import partial
 
 import gradio as gr
+import spaces
 import torch
 import torchaudio
 
 from resemble_enhance.enhancer.inference import denoise, enhance
 
-if torch.cuda.is_available():
-    device = "cuda"
-else:
-    device = "cpu"
-
 
+@spaces.GPU(duration=300)
 def _fn(path, solver, nfe, tau, denoising, unlimited):
     if path is None:
         gr.Warning("Please upload an audio file.")
@@ -30,6 +27,7 @@ def _fn(path, solver, nfe, tau, denoising, unlimited):
     dwav, sr = torchaudio.load(path)
     dwav = dwav.mean(dim=0)
 
+    device = "cuda" if torch.cuda.is_available() else "cpu"
     wav1, new_sr = denoise(dwav, sr, device)
     wav2, new_sr = enhance(dwav, sr, device, nfe=nfe, solver=solver, lambd=lambd, tau=tau)
 

requirements.txt

@@ -1,5 +1,4 @@
 celluloid==0.2.0
-deepspeed==0.12.4
 librosa==0.10.1
 matplotlib==3.8.1
 numpy==1.26.2
@@ -9,9 +8,10 @@ ptflops==0.7.1.2
 rich==13.7.0
 scipy==1.11.4
 soundfile==0.12.1
-torch==2.1.1
-torchaudio==2.1.1
-torchvision==0.16.1
+spaces==0.50.4
+torch==2.8.0
+torchaudio==2.8.0
+torchvision==0.23.0
 tqdm==4.66.1
 resampy==0.4.2
 tabulate==0.8.10

resemble_enhance/denoiser/inference.py

@@ -4,7 +4,8 @@ from functools import cache
 import torch
 
 from ..inference import inference
-from .train import Denoiser, HParams
+from .denoiser import Denoiser
+from .hparams import HParams
 
 logger = logging.getLogger(__name__)
 

resemble_enhance/enhancer/enhancer.py

@@ -1,6 +1,5 @@
 import logging
 
-import matplotlib.pyplot as plt
 import pandas as pd
 import torch
 from torch import Tensor, nn
@@ -9,8 +8,6 @@ from torch.distributions import Beta
 from ..common import Normalizer
 from ..denoiser.inference import load_denoiser
 from ..melspec import MelSpectrogram
-from ..utils.distributed import global_leader_only
-from ..utils.train_loop import TrainLoop
 from .hparams import HParams
 from .lcfm import CFM, IRMAE, LCFM
 from .univnet import UnivNet
@@ -106,24 +103,9 @@ class Enhancer(nn.Module):
             return self.mel_fn(x)[..., :-1]  # (b d t)
         return self.mel_fn(x)
 
-    @global_leader_only
     @torch.no_grad()
     def _visualize(self, original_mel, denoised_mel):
-        loop = TrainLoop.get_running_loop()
-        if loop is None or loop.global_step % 100 != 0:
-            return
-
-        plt.figure(figsize=(6, 6))
-        plt.subplot(211)
-        plt.title("Original")
-        plt.imshow(original_mel[0].cpu().numpy(), origin="lower", interpolation="none")
-        plt.subplot(212)
-        plt.title("Denoised")
-        plt.imshow(denoised_mel[0].cpu().numpy(), origin="lower", interpolation="none")
-        plt.tight_layout()
-
-        path = loop.get_running_loop_viz_path("input", ".png")
-        plt.savefig(path, dpi=300)
+        return
 
     def _may_denoise(self, x: Tensor, y: Tensor | None = None):
         if self.hp.lcfm_training_mode == "cfm":

resemble_enhance/enhancer/inference.py

@@ -5,7 +5,8 @@ import torch
 
 from ..inference import inference
 from .download import download
-from .train import Enhancer, HParams
+from .enhancer import Enhancer
+from .hparams import HParams
 
 logger = logging.getLogger(__name__)