quazim commited on
Commit
d153695
·
1 Parent(s): 5471a60

added elastic model

Browse files
Files changed (1) hide show
  1. app.py +2 -3
app.py CHANGED
@@ -48,7 +48,7 @@ import os
48
  # setup_flash_attention()
49
 
50
  # Load model and processor
51
- # @gr.cache()
52
  def load_model():
53
  """Load the musicgen model and processor"""
54
  processor = AutoProcessor.from_pretrained("facebook/musicgen-large")
@@ -71,8 +71,7 @@ def generate_music(text_prompt, duration=10, temperature=1.0, top_k=250, top_p=0
71
  text=[text_prompt],
72
  padding=True,
73
  return_tensors="pt",
74
- )
75
-
76
  # Generate audio
77
  with torch.no_grad():
78
  audio_values = model.generate(
 
48
  # setup_flash_attention()
49
 
50
  # Load model and processor
51
+ @gr.cache()
52
  def load_model():
53
  """Load the musicgen model and processor"""
54
  processor = AutoProcessor.from_pretrained("facebook/musicgen-large")
 
71
  text=[text_prompt],
72
  padding=True,
73
  return_tensors="pt",
74
+ ).to("cuda")
 
75
  # Generate audio
76
  with torch.no_grad():
77
  audio_values = model.generate(