Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions app.py
Original file line number Diff line number Diff line change
Expand Up @@ -631,7 +631,9 @@ def _do_predictions(gen_type, texts, melodies, sample, trim_start, trim_end, dur
if UNLOAD_MODEL:
MODEL = None
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
# fix not install cuda driver
if torch.cuda.is_available():
torch.cuda.ipc_collect()
return res, res_audio, res_backup, input_length


Expand Down Expand Up @@ -1832,4 +1834,4 @@ def ui_batched(launch_kwargs):
USE_DIFFUSION = False
ui_batched(launch_kwargs)
else:
ui_full(launch_kwargs)
ui_full(launch_kwargs)
9 changes: 2 additions & 7 deletions audiocraft/models/audiogen.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,16 +71,11 @@ def audio_channels(self) -> int:
return self.compression_model.channels

@staticmethod
def get_pretrained(name: str = 'facebook/audiogen-medium', device=None):
def get_pretrained(name: str = 'facebook/audiogen-medium', device=('cuda' if torch.cuda.is_available() else 'cpu')):
"""Return pretrained model, we provide a single model for now:
- facebook/audiogen-medium (1.5B), text to sound,
# see: https://huggingface.co/facebook/audiogen-medium
"""
if device is None:
if torch.cuda.device_count():
device = 'cuda'
else:
device = 'cpu'

if name == 'debug':
# used only for unit tests
Expand Down Expand Up @@ -273,4 +268,4 @@ def _progress_callback(generated_tokens: int, tokens_to_generate: int):
def to(self, device: str):
self.compression_model.to(device)
self.lm.to(device)
return self
return self
9 changes: 2 additions & 7 deletions audiocraft/models/musicgen.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ def audio_channels(self) -> int:
return self.compression_model.channels

@staticmethod
def get_pretrained(name: str = 'GrandaddyShmax/musicgen-melody', device=None):
def get_pretrained(name: str = 'GrandaddyShmax/musicgen-melody', device=('cuda' if torch.cuda.is_available() else 'cpu')):
"""Return pretrained model, we provide four models:
- facebook/musicgen-small (300M), text to music,
# see: https://huggingface.co/facebook/musicgen-small
Expand All @@ -96,11 +96,6 @@ def get_pretrained(name: str = 'GrandaddyShmax/musicgen-melody', device=None):
- facebook/musicgen-large (3.3B), text to music,
# see: https://huggingface.co/facebook/musicgen-large
"""
if device is None:
if torch.cuda.device_count():
device = 'cuda'
else:
device = 'cpu'

if name == 'debug':
# used only for unit tests
Expand Down Expand Up @@ -406,4 +401,4 @@ def generate_audio(self, gen_tokens: torch.Tensor):
def to(self, device: str):
self.compression_model.to(device)
self.lm.to(device)
return self
return self