Transformers MusicGen

app.py

@@ -97,7 +97,7 @@ def write_story_fn(story_topic, main_role, scene,
 
 def modality_assets_generation_fn(
         height, width, image_seed, sound_guidance_scale, sound_seed,
-        n_candidate_per_text, music_duration,
+        n_candidate_per_text,
         config,
         story_data):
     deep_update(config, {
@@ -117,11 +117,6 @@ def modality_assets_generation_fn(
                 "n_candidate_per_text": n_candidate_per_text
             }
         },
-        "music_generation": {
-            "call_cfg": {
-                "duration": music_duration
-            }
-        }
     })
     story_gen_agent = MMStoryAgent()
     images = story_gen_agent.generate_modality_assets(config, story_data)
@@ -180,9 +175,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 sound_seed = gr.Number(label="Sound Seed", value=default_sound_config["call_cfg"]['seed'])
                 n_candidate_per_text = gr.Slider(label="Number of Candidates per Text", minimum=0, maximum=5, step=1, value=default_sound_config["call_cfg"]['n_candidate_per_text'])
 
-            with gr.Accordion("Detailed Music Configuration (Optional)", open=False):
-                music_duration = gr.Number(label="Music Duration", min_width=30.0, maximum=120.0, value=default_music_config["call_cfg"]["duration"])
-
             with gr.Accordion("Detailed Slideshow Effect (Optional)", open=False):
                 fade_duration = gr.Slider(label="Fade Duration", minimum=0.1, maximum=1.5, step=0.1, value=default_slideshow_effect['fade_duration'])
                 slide_duration = gr.Slider(label="Slide Duration", minimum=0.1, maximum=1.0, step=0.1, value=default_slideshow_effect['slide_duration'])
@@ -244,7 +236,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     ).then(
         fn=modality_assets_generation_fn,
         inputs=[height, width, image_seed, sound_guidance_scale, sound_seed,
-                n_candidate_per_text, music_duration,
+                n_candidate_per_text,
                 config,
                 story_data],
         outputs=[image_gallery]

configs/mm_story_agent.yaml

@@ -56,8 +56,7 @@ music_generation:
     revise_cfg:
         num_turns: 3
     obj_cfg: {}
-    call_cfg:
-        duration: 60.0
+    call_cfg: {}
 
 slideshow_effect:
     fade_duration: 0.8

mm_story_agent/modality_agents/music_agent.py

@@ -2,9 +2,9 @@ from pathlib import Path
 import json
 from typing import List, Union
 
+import soundfile as sf
 import torchaudio
-from audiocraft.models import MusicGen
-from audiocraft.data.audio import audio_write
+from transformers import AutoProcessor, MusicgenForConditionalGeneration
 
 from mm_story_agent.modality_agents.llm import QwenAgent
 from mm_story_agent.prompts_en import story_to_music_reviser_system, story_to_music_reviewer_system
@@ -17,19 +17,23 @@ class MusicGenSynthesizer:
                  device: str = 'cuda',
                  sample_rate: int = 16000,
                  ) -> None:
-        self.model = MusicGen.get_pretrained(model_name, device=device).to(device)
+        self.device = device
+        self.processor = AutoProcessor.from_pretrained(model_name)
+        self.model = MusicgenForConditionalGeneration.from_pretrained(model_name).to(device)
         self.sample_rate = sample_rate
     
     def call(self,
              prompt: Union[str, List[str]],
              save_path: Union[str, Path],
-             duration: float = 60.0,
              ):
-        self.model.set_generation_params(duration=duration)
-        wav = self.model.generate([prompt], progress=True)[0].cpu()
-        wav = torchaudio.functional.resample(wav, self.model.sample_rate, self.sample_rate)
-        save_path = Path(save_path).parent / Path(save_path).stem
-        audio_write(save_path, wav, self.sample_rate)
+        inputs = self.processor(
+            text=[prompt],
+            padding=True,
+            return_tensors="pt",
+        ).to(self.device)
+        wav = self.model.generate(**inputs, max_new_tokens=1536)[0, 0].cpu()
+        wav = torchaudio.functional.resample(wav, self.model.config.audio_encoder.sampling_rate, self.sample_rate)
+        sf.write(save_path, wav.numpy(), self.sample_rate)
 
 
 class MusicGenAgent: