Update app.py

app.py

@@ -85,7 +85,7 @@ def write_story_fn(story_topic, main_role, scene,
     # story_data, story_accordion, story_content
     return pages, gr.update(visible=True), pages[current_page], gr.update()
 
-@spaces.GPU(duration=600)
+@spaces.GPU()
 def modality_assets_generation_fn(
         height, width, image_seed, sound_guidance_scale, sound_seed,
         n_candidate_per_text, music_duration,
@@ -119,6 +119,57 @@ def modality_assets_generation_fn(
     # image gallery
     return gr.update(visible=True, value=images, columns=[len(images)], rows=[1], height="auto")
 
+def speech_generation_fn(story_data):
+    story_gen_agent = MMStoryAgent()
+    story_gen_agent.generate_speech(config, story_data)
+
+@spaces.GPU(duration=120)
+def sound_generation_fn(sound_guidance_scale, sound_seed, n_candidate_per_text,
+                        story_data, progress=gr.Progress(track_tqdm=True)):
+    deep_update(config, {
+        "sound_generation": {
+            "call_cfg": {
+                "guidance_scale": sound_guidance_scale,
+                "seed": sound_seed,
+                "n_candidate_per_text": n_candidate_per_text
+            }
+        }
+    })
+    story_gen_agent = MMStoryAgent()
+    story_gen_agent.generate_sound(config, story_data)
+
+@spaces.GPU(duration=120)
+def music_generation_fn(music_duration,
+                        story_data, progress=gr.Progress(track_tqdm=True)):
+    deep_update(config, {
+        "music_generation": {
+            "call_cfg": {
+                "duration": music_duration
+            }
+        }
+    })
+    story_gen_agent = MMStoryAgent()
+    story_gen_agent.generate_music(config, story_data)
+
+@spaces.GPU(duration=120)
+def image_generation_fn(height, width, image_seed,
+                        story_data, progress=gr.Progress(track_tqdm=True)):
+    deep_update(config, {
+        "image_generation": {
+            "obj_cfg": {
+                "height": height,
+                "width": width,
+            },
+            "call_cfg": {
+                "seed": image_seed
+            }
+        },
+    })
+    story_gen_agent = MMStoryAgent()
+    result = story_gen_agent.generate_image(config, story_data)
+    images = result["images"]
+    return gr.update(visible=True, value=images, columns=[len(images)], rows=[1], height="auto")
+
 def compose_storytelling_video_fn(
         fade_duration, slide_duration, zoom_speed, move_ratio,
         sound_volume, music_volume, bg_speech_ratio, fps, 
@@ -228,10 +279,16 @@ if __name__ == "__main__":
             inputs=[gr.State("Generating Modality Assets")],
             outputs=video_generation_information
         ).then(
-            fn=modality_assets_generation_fn,
-            inputs=[height, width, image_seed, sound_guidance_scale, sound_seed,
-                    n_candidate_per_text, music_duration,
-                    story_data],
+            fn=speech_generation_fn,
+        ).then(
+            fn=sound_generation_fn,
+            inputs=[sound_guidance_scale, sound_seed, n_candidate_per_text, story_data]
+        ).then(
+            fn=music_generation_fn,
+            inputs=[music_duration, story_data]
+        ).then(
+            fn=image_generation_fn,
+            inputs=[height, width, image_seed, story_data],
             outputs=[image_gallery]
         ).then(
             fn=set_generating_progress_text,

mm_story_agent/__init__.py

@@ -34,6 +34,22 @@ class MMStoryAgent:
         pages = story_writer.call(config["story_setting"])
         return pages
     
+    def generate_speech(self, config, pages):
+        speech_agent = CosyVoiceAgent(config["speech_generation"])
+        speech_agent.call(pages, config["story_dir"] / "speech")
+    
+    def generate_sound(self, config, pages):
+        sound_agent = AudioLDM2Agent(config["sound_generation"])
+        sound_agent.call(pages, config["story_dir"] / "sound")
+
+    def generate_music(self, config, pages):
+        music_agent = MusicGenAgent(config["music_generation"])
+        music_agent.call(pages, config["story_dir"] / "music")
+    
+    def generate_image(self, config, pages):
+        image_agent = StoryDiffusionAgent(config["image_generation"])
+        image_agent.call(pages, config["story_dir"] / "image")
+    
     def generate_modality_assets(self, config, pages):
         script_data = {"pages": [{"story": page} for page in pages]}
         story_dir = Path(config["story_dir"])
@@ -45,22 +61,16 @@ class MMStoryAgent:
         for modality in self.modalities:
             agents[modality] = self.modality_agent_class[modality](config[modality + "_generation"])
 
-        # processes = []
-        # return_dict = mp.Manager().dict()
+        processes = []
+        return_dict = mp.Manager().dict()
 
-        # for modality in self.modalities:
-        #     p = mp.Process(target=self.call_modality_agent, args=(agents[modality], pages, story_dir / modality, return_dict), daemon=False)
-        #     processes.append(p)
-        #     p.start()
-        
-        # for p in processes:
-        #     p.join()
-
-
-        return_dict = {}
-        
         for modality in self.modalities:
-            self.call_modality_agent(agents[modality], pages, story_dir / modality, return_dict)
+            p = mp.Process(target=self.call_modality_agent, args=(agents[modality], pages, story_dir / modality, return_dict), daemon=False)
+            processes.append(p)
+            p.start()
+        
+        for p in processes:
+            p.join()
 
         for modality, result in return_dict.items():
             try: