KingNish commited on
Commit
b1860c5
·
1 Parent(s): 59e8f28

modified: app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -19,6 +19,7 @@ from torchaudio.transforms import Resample
19
  from einops import rearrange
20
  from tqdm import tqdm
21
  from omegaconf import OmegaConf
 
22
 
23
  # --- Constants and Environment Setup ---
24
  IS_SHARED_UI = "innova-ai/YuE-music-generator-demo" in os.environ.get('SPACE_ID', '')
@@ -165,6 +166,7 @@ class BlockTokenRangeProcessor(LogitsProcessor):
165
  return scores
166
 
167
  # --- Music Generation Core Function ---
 
168
  def generate_music(
169
  stage1_model,
170
  mmtokenizer,
 
19
  from einops import rearrange
20
  from tqdm import tqdm
21
  from omegaconf import OmegaConf
22
+ import spaces
23
 
24
  # --- Constants and Environment Setup ---
25
  IS_SHARED_UI = "innova-ai/YuE-music-generator-demo" in os.environ.get('SPACE_ID', '')
 
166
  return scores
167
 
168
  # --- Music Generation Core Function ---
169
+ @spaces.GPU(duration=120)
170
  def generate_music(
171
  stage1_model,
172
  mmtokenizer,