Update audiocraft/models/musicgen.py

audiocraft/models/musicgen.py

@@ -78,11 +78,14 @@ class MusicGen:
         self.generation_params: dict = {}
         self.set_generation_params(duration=15)  # 15 seconds by default
         self._progress_callback: tp.Optional[tp.Callable[[int, int], None]] = None
+        '''
         if self.device.type == 'cpu':
             self.autocast = TorchAutocast(enabled=False)
         else:
             self.autocast = TorchAutocast(
                 enabled=True, device_type=self.device.type, dtype=torch.bfloat16)
+        '''
+        self.autocast = TorchAutocast(enabled=False)
 
     @property
     def frame_rate(self) -> float:
@@ -173,7 +176,7 @@ class MusicGen:
         """Override the default progress callback."""
         self._progress_callback = progress_callback
 
-    def generate_unconditional(self, num_samples: int, progress: bool = False,
+    def generate_unconditional(self, num_samples: int, progress: bool = True,
                                return_tokens: bool = False) -> tp.Union[torch.Tensor,
                                                                         tp.Tuple[torch.Tensor, torch.Tensor]]:
         """Generate samples in an unconditional manner.
@@ -189,7 +192,7 @@ class MusicGen:
             return self.generate_audio(tokens), tokens
         return self.generate_audio(tokens)
 
-    def generate(self, descriptions: tp.List[str], progress: bool = False, return_tokens: bool = False) \
+    def generate(self, descriptions: tp.List[str], progress: bool = True, return_tokens: bool = False) \
             -> tp.Union[torch.Tensor, tp.Tuple[torch.Tensor, torch.Tensor]]:
         """Generate samples conditioned on text.
 
@@ -205,7 +208,7 @@ class MusicGen:
         return self.generate_audio(tokens)
 
     def generate_with_chroma(self, descriptions: tp.List[str], melody_wavs: MelodyType,
-                             melody_sample_rate: int, progress: bool = False,
+                             melody_sample_rate: int, progress: bool = True,
                              return_tokens: bool = False) -> tp.Union[torch.Tensor,
                                                                       tp.Tuple[torch.Tensor, torch.Tensor]]:
         """Generate samples conditioned on text and melody.
@@ -244,7 +247,7 @@ class MusicGen:
 
     def generate_continuation(self, prompt: torch.Tensor, prompt_sample_rate: int,
                               descriptions: tp.Optional[tp.List[tp.Optional[str]]] = None,
-                              progress: bool = False, return_tokens: bool = False) \
+                              progress: bool = True, return_tokens: bool = False) \
             -> tp.Union[torch.Tensor, tp.Tuple[torch.Tensor, torch.Tensor]]:
         """Generate samples conditioned on audio prompts.
 
@@ -328,7 +331,7 @@ class MusicGen:
         return attributes, prompt_tokens
 
     def _generate_tokens(self, attributes: tp.List[ConditioningAttributes],
-                         prompt_tokens: tp.Optional[torch.Tensor], progress: bool = False) -> torch.Tensor:
+                         prompt_tokens: tp.Optional[torch.Tensor], progress: bool = True) -> torch.Tensor:
         """Generate discrete audio tokens given audio prompt and/or conditions.
 
         Args: