[Admin maintenance] Support new ZeroGPU hardware

app.py

@@ -6,6 +6,7 @@ This source code is licensed under the license found in the
 LICENSE file in the root directory of this source tree.
 """
 
+import spaces  # noqa: F401  (must be imported before torch / any CUDA-touching module)
 from tempfile import NamedTemporaryFile
 import argparse
 import torch
@@ -250,6 +251,16 @@ def load_melody_filepath(melody_filepath, title, assigned_model, topp, temperatu
     
     return  gr.update(value=melody_name), gr.update(maximum=MAX_PROMPT_INDEX, value=-1), gr.update(value=assigned_model, interactive=True), gr.update(value=topp), gr.update(value=temperature), gr.update(value=cfg_coef), gr.update(maximum=MAX_OVERLAP)
 
+def _gpu_duration(model=None, text=None, melody_filepath=None, duration=10, *args, **kwargs):
+    """Rough per-call GPU time estimate for the ZeroGPU scheduler, scaled with requested audio duration."""
+    try:
+        d = int(duration)
+    except (TypeError, ValueError):
+        d = 10
+    return max(30, min(d * 3 + 30, 300))
+
+
+@spaces.GPU(duration=_gpu_duration)
 def predict(model, text, melody_filepath = None, duration=10, dimension=2, topk=200, topp=0, temperature=1.0, cfg_coef=4.0, background = None, title="UnlimitedMusicGen", settings_font="./assets/arial.ttf", settings_font_color = "#c87f05", seed=-1, overlap=1, prompt_index = 0, include_title = True, include_settings = True, harmony_only = False, profile = gr.OAuthProfile, segment_length = 30, settings_font_size=28, settings_animate_waveform=False, video_orientation="Landscape", excerpt_duration=3.5, return_history_json=False, progress=gr.Progress(track_tqdm=True)):
     """
     Generate music and video based on the provided parameters and model.
@@ -621,6 +632,7 @@ def fix_path(path: str) -> str:
         return "./" + path[index:].replace("\\", "/")
     return path
 # Add this wrapper function above the gr.api definitions
+@spaces.GPU(duration=_gpu_duration)
 def predict_simple(model: str, text: str,  melody_filepath: str = None, duration: int = 10, dimension: int = 2, topk: int = 200, topp: float = 0.01, temperature: float = 1.0, cfg_coef: float = 4.0, background: str = "./assets/background.png", title: str = "UnlimitedMusicGen", settings_font: str = "./assets/arial.ttf", settings_font_color: str = "#c87f05", seed: int = -1, overlap: int = 1, prompt_index: int = -1, include_title: bool = True, include_settings: bool = True, harmony_only: bool = False, profile: str = "Satoshi Nakamoto", segment_length: int = 30, settings_font_size: int = 28, settings_animate_waveform: bool = False, video_orientation: str = "Landscape",  return_history_json: bool = False) -> tp.List[tp.Tuple[str, str, str]]:
     """
         Generate music and video based on the provided parameters and model.

requirements.txt

@@ -4,14 +4,13 @@ einops
 flashy>=0.0.1
 hydra-core>=1.1
 hydra_colorlog
-torch==2.6.0 --extra-index-url https://download.pytorch.org/whl/cu124
-torchaudio>=2.0.0,<2.6.2 --extra-index-url https://download.pytorch.org/whl/cu124
+torchaudio
 soundfile
 huggingface_hub
 hf_xet
 tqdm
 transformers==4.43.4  # need Encodec there.
-xformers>=0.0.23 --index-url https://download.pytorch.org/whl/cu124
+xformers
 demucs
 librosa==0.11.0
 soundfile
@@ -29,7 +28,7 @@ mutagen
 fastapi>=0.88.0
 pydantic
 typer
-torchvision>=0.21.0 --extra-index-url https://download.pytorch.org/whl/cu124
+torchvision
 #torchtext
 pesq
 pystoi