Update app.py

app.py

@@ -1,47 +1,339 @@
+# import streamlit as st
+# import re
+# from transformers import pipeline
+# from diffusers import StableDiffusionPipeline
+# from pydub import AudioSegment
+# import zipfile
+# from io import BytesIO
+# import json
+# import torch
+
+# # 初始化模型 (全部使用原始pipeline)
+# @st.cache_resource
+# def load_pipelines():
+#     try:
+#         # 1. Script generation
+#         script_pipe = pipeline(
+#             "text2text-generation",
+#             model="RUCAIBox/mvp-story",
+#             device=0 if torch.cuda.is_available() else -1
+#         )
+
+#         # 2. Storyboard generation (BART model)
+#         storyboard_pipe = pipeline(
+#             "text2text-generation",
+#             model="Jessiesj/script_gpt2-story",
+#             device=0 if torch.cuda.is_available() else -1
+#         )
+
+#         # 3. Soundtrack generation (MusicGen)
+#         music_pipe = pipeline(
+#             "text-to-audio", 
+#             model="facebook/musicgen-small",
+#             device_map="auto"
+#         )
+
+#         # 4. Storyboard image generation (Stable Diffusion)
+#         image_pipe = StableDiffusionPipeline.from_pretrained(
+#             "prompthero/openjourney-v4",
+#             safety_checker=None,
+#             torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+#         )
+
+#         return script_pipe, storyboard_pipe, image_pipe
+#         # return script_pipe, storyboard_pipe, music_pipe, image_pipe
+
+#     except Exception as e:
+#         st.error(f"Failed to load models: {str(e)}")
+#         raise
+
+# # Prompt template
+# def build_script_prompt(theme, style_or_opening):
+#     return f"""
+# {style_or_opening}
+# Generate a micro - movie script with two scenes, with the theme: {theme}.
+# ### Format requirements (must be strictly followed):
+# - Use Markdown syntax
+# - Each scene starts with ### Scene X: [Location] (INT/EXT. Time)
+# - Include action descriptions (wrapped in square brackets []) and character dialogues
+
+# ### Example:
+# ### Scene 1: Spaceship cockpit (INT. Emergency)
+# [Alarm lights are flashing wildly, sparks are coming out of the console]
+# Captain (shouting): Start the backup engine!
+# Co - pilot (typing on the keyboard): The power system has failed!
+
+# ### Scene 2: Alien jungle (EXT. Dusk)
+# [The expedition team pushes aside the dense purple bushes]
+# Team member A (pointing into the distance): Look! Is that the entrance to the ruins?
+# """
+
+
+# # Content cleaning
+# def format_script(text):
+#     # Remove invalid symbols
+#     text = re.sub(r"[():<>]{2,}", "", text)
+#     # Standardize scene titles
+#     return re.sub(r"(?<!### )Scene\d+", r"### Scene\g<0>", text)
+
+# # Load pipelines
+# script_pipe, storyboard_pipe, image_pipe = load_pipelines()
+# # script_pipe, storyboard_pipe, music_pipe, image_pipe = load_pipelines()
+
+# # Streamlit interface
+# st.title("🎥 Micro - movie Intelligent Creation System")
+
+# # User input
+# user_input = st.text_input(
+#     "Enter the movie theme (e.g., Sci - fi space adventure)",
+#     help="Recommended format: Genre + Core conflict, example: 'Jungle exploration + Search for the lost golden city'"
+# )
+
+# style_or_opening = st.text_input(
+#     "Enter the script style or opening",
+#     help="Provide a style description or the beginning of the script"
+# )
+
+# if user_input and style_or_opening:
+
+#     # 1. Generate the script
+#     with st.status("🖋️ Generating the script...", expanded=True) as status:
+#         try:
+#             # Generate content
+#             prompt = build_script_prompt(user_input, style_or_opening)
+#             response = script_pipe(
+#                 prompt,
+#                 max_length=600,
+#                 temperature=0.8,
+#                 num_beams=4,
+#                 no_repeat_ngram_size=3
+#             )[0]["generated_text"]
+
+#             # Post - processing
+#             def format_script(text):
+#                 return text
+
+#             # Display the result
+#             st.subheader("Generated Script")
+#             cleaned = format_script(response)  # Ensure cleaned is defined here
+#             st.markdown(f"```markdown\n{format_script(response)}\n```")
+#             status.update(label="✅ Generation completed", state="complete")
+
+#         except Exception as e:
+#             status.update(label="❌ Generation failed", status="error")
+#             st.error(f"Error details: {str(e)}")
+#             st.stop()
+
+#     # 2. Generate the storyboard
+#     with st.status("🎥 Converting to storyboard script...", expanded=True) as status:
+#         try:
+#             # 示例数据模板
+#             sample_storyboard = """
+#             {
+#               "scenes": [
+#                 {
+#                   "scene_number": 1,
+#                   "location": "Abandoned Space Station (INT. Dystopian Future)",
+#                   "characters": ["Eve", "R2-D9"],
+#                   "camera_angle": "Wide shot showing decaying corridors",
+#                   "action": "Eve scans walls with her holographic projector, R2-D9 beeps nervously"
+#                 },
+#                 {
+#                   "scene_number": 2,
+#                   "location": "Alien Wasteland (EXT. Toxic Sunset)",
+#                   "characters": ["Eve", "Mutant Leader"],
+#                   "camera_angle": "Low-angle shot of Eve aiming plasma rifle",
+#                   "action": "Mutant leader growls, acid rain hisses on Eve's protective suit"
+#                 }
+#               ]
+#             }
+#             """
+
+#             # JSON模板引导
+#             json_template = """{
+#               "scenes": [
+#                 {
+#                   "scene_number": 1,
+#                   "location": "Scene location (INT/EXT)",
+#                   "characters": ["Character 1", "Character 2"],
+#                   "camera_angle": "Camera angle",
+#                   "action": "Main action description"
+#                 }
+#               ]
+#             }"""
+
+#             # 生成分镜
+#             # storyboard = storyboard_pipe(
+#             #     f"Convert to storyboard JSON: {cleaned[:800]}",
+#             #     max_length=1500,
+#             #     temperature=0.3
+#             # )[0]["generated_text"]
+
+#             # 增强的格式修复
+#             # storyboard = (
+#             #     storyboard.replace("'", '"')
+#             #     .replace("None", '""')
+#             #     .replace("True", "true")
+#             #     .replace("False", "false")
+#             # )
+#             # storyboard = re.search(r'\{.*\}', storyboard, re.DOTALL).group()
+#             # parsed = json.loads(storyboard)
+#             parsed = json.loads(sample_storyboard)
+#         except Exception as e:
+#             # st.warning(f"Using sample storyboard due to error: {str(e)}")
+#             # parsed = json.loads(sample_storyboard)  # 回退到示例数据
+#             # status.update(label="⚠️ Using sample storyboard", status="warning")
+#             pass
+
+#         finally:
+#         #     status.update(label="Storyboard ready", status="complete")
+#         #     st.subheader("Storyboard Script")
+#         #     st.json(parsed)
+#             # st.caption("*Note: If using sample data, edit the JSON before proceeding*")
+#             pass
+
+
+#     # 3. 生成配乐
+#     # with st.spinner("🎵 正在生成配乐..."):
+#     #     audio = music_pipe(
+#     #         "epic sci-fi background music", 
+#     #         max_new_tokens=256
+#     #     )
+#     #     audio_buffer = BytesIO()
+#     #     audio["audio"].export(audio_buffer, format="wav")
+    
+#     # st.subheader("🎧 背景音乐")
+#     # st.audio(audio_buffer, format="audio/wav")
+
+#     # 4. 生成分镜图（可选）
+#     # if st.toggle("🎨 生成分镜预览图（需要2-3分钟）"):
+#     #         with st.spinner("🖼️ 正在渲染电影级画面..."):
+#     #             # 从剧本提取关键描述
+#     #             scene_desc = re.search(r"场景\d+：(.*?)\n", script)
+#     #             if scene_desc:
+#     #                 scene_desc = scene_desc.group(1)
+#     #             else:
+#     #                 scene_desc = f"movie scene of {user_input}"
+                
+#     #             # 精细化图片提示词
+#     #             image_prompt = f"""电影剧照，{scene_desc}，
+#     #                 超高清8K分辨率，
+#     #                 电影级打光，浅景深效果，
+#     #                 Unreal Engine 5渲染，
+#     #                 电影胶片颗粒感"""
+                
+#     #             # 优化生成参数
+#     #             image = image_pipe(
+#     #                 image_prompt,
+#     #                 num_inference_steps=50,
+#     #                 guidance_scale=9.5,
+#     #                 width=1024,
+#     #                 height=576  # 16:9电影比例
+#     #             ).images[0]
+                
+#     #             st.subheader("🎞️ 首场景预览")
+#     #             st.image(image, caption=scene_desc)
+                
+#     # except Exception as e:
+#     #     st.error("生成过程中发生错误")
+#     #     st.code(f"错误详情：{str(e)}")
+#     #     st.button("点击查看技术详情", 
+#     #              help="将以下信息提供给技术人员：\n" + str(e))
+    
+# st.subheader("🎧 Background Music")
+# for i, scene in enumerate(parsed["scenes"]):
+#     action = scene["action"]
+#     with st.spinner(f"🎵 Generating soundtrack for scene {i + 1}..."):
+#         try:
+#             audio = music_pipe(
+#                 action,
+#                 max_new_tokens=256
+#             )
+#             audio_buffer = BytesIO()
+#             audio["audio"].export(audio_buffer, format="wav")
+#             st.audio(audio_buffer, format="audio/wav", caption=f"Soundtrack for scene {i + 1}")
+#         except Exception as e:
+#             st.error(f"Failed to generate soundtrack for scene {i + 1}: {str(e)}")
+
+# st.subheader("🎨 Storyboard Images")
+# for i, scene in enumerate(parsed["scenes"]):
+#     location = scene["location"]
+#     if st.toggle(f"Generate preview for scene {i + 1} (takes 2 - 3 minutes)"):
+#         with st.spinner(f"🖼️ Rendering image for scene {i + 1}..."):
+#             try:
+#                 # 精细化图片提示词
+#                 image_prompt = f"""Movie still, {location},
+#                     Ultra - high - definition 8K resolution,
+#                     Movie - level lighting, shallow depth - of - field effect,
+#                     Unreal Engine 5 rendering,
+#                     Movie film graininess"""
+
+#                 # 优化生成参数
+#                 image = image_pipe(
+#                     image_prompt,
+#                     num_inference_steps=50,
+#                     guidance_scale=9.5,
+#                     width=1024,
+#                     height=576  # 16:9 movie ratio
+#                 ).images[0]
+
+#                 st.image(image, caption=f"Preview for scene {i + 1}")
+#             except Exception as e:
+#                 st.error(f"Failed to generate image for scene {i + 1}: {str(e)}")
+
 import streamlit as st
 import re
-from transformers import pipeline
+from transformers import pipeline, AutoModelForTextToWaveform, AutoTokenizer
 from diffusers import StableDiffusionPipeline
 from pydub import AudioSegment
 import zipfile
 from io import BytesIO
 import json
 import torch
+import torchaudio
 
 # 初始化模型 (全部使用原始pipeline)
 @st.cache_resource
 def load_pipelines():
     try:
+        device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        torch_dtype = torch.float16 if device == "cuda:0" else torch.float32
+
         # 1. Script generation
         script_pipe = pipeline(
             "text2text-generation",
             model="RUCAIBox/mvp-story",
-            device=0 if torch.cuda.is_available() else -1
+            device=device
         )
 
         # 2. Storyboard generation (BART model)
         storyboard_pipe = pipeline(
             "text2text-generation",
             model="Jessiesj/script_gpt2-story",
-            device=0 if torch.cuda.is_available() else -1
+            device=device
         )
 
         # 3. Soundtrack generation (MusicGen)
+        music_model = AutoModelForTextToWaveform.from_pretrained(
+            "facebook/musicgen-small",
+            torch_dtype=torch_dtype
+        ).to(device)
+        music_tokenizer = AutoTokenizer.from_pretrained("facebook/musicgen-small")
         music_pipe = pipeline(
-            "text-to-audio", 
-            model="facebook/musicgen-small",
-            device_map="auto"
+            "text-to-audio",
+            model=music_model,
+            tokenizer=music_tokenizer,
+            device=device
         )
 
         # 4. Storyboard image generation (Stable Diffusion)
         image_pipe = StableDiffusionPipeline.from_pretrained(
             "prompthero/openjourney-v4",
             safety_checker=None,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-        )
+            torch_dtype=torch_dtype
+        ).to(device)
 
-        return script_pipe, storyboard_pipe, image_pipe
-        # return script_pipe, storyboard_pipe, music_pipe, image_pipe
+        return script_pipe, storyboard_pipe, music_pipe, image_pipe
 
     except Exception as e:
         st.error(f"Failed to load models: {str(e)}")
@@ -77,8 +369,7 @@ def format_script(text):
     return re.sub(r"(?<!### )Scene\d+", r"### Scene\g<0>", text)
 
 # Load pipelines
-script_pipe, storyboard_pipe, image_pipe = load_pipelines()
-# script_pipe, storyboard_pipe, music_pipe, image_pipe = load_pipelines()
+script_pipe, storyboard_pipe, music_pipe, image_pipe = load_pipelines()
 
 # Streamlit interface
 st.title("🎥 Micro - movie Intelligent Creation System")
@@ -117,10 +408,10 @@ if user_input and style_or_opening:
             st.subheader("Generated Script")
             cleaned = format_script(response)  # Ensure cleaned is defined here
             st.markdown(f"```markdown\n{format_script(response)}\n```")
-            status.update(label="✅ Generation completed", state="complete")
+            status.update(label="✅ Generation completed")
 
         except Exception as e:
-            status.update(label="❌ Generation failed", status="error")
+            status.update(label="❌ Generation failed")
             st.error(f"Error details: {str(e)}")
             st.stop()
 
@@ -163,120 +454,70 @@ if user_input and style_or_opening:
             }"""
 
             # 生成分镜
-            # storyboard = storyboard_pipe(
-            #     f"Convert to storyboard JSON: {cleaned[:800]}",
-            #     max_length=1500,
-            #     temperature=0.3
-            # )[0]["generated_text"]
+            storyboard = storyboard_pipe(
+                f"Convert to storyboard JSON: {cleaned[:800]}",
+                max_length=1500,
+                temperature=0.3
+            )[0]["generated_text"]
 
             # 增强的格式修复
-            # storyboard = (
-            #     storyboard.replace("'", '"')
-            #     .replace("None", '""')
-            #     .replace("True", "true")
-            #     .replace("False", "false")
-            # )
-            # storyboard = re.search(r'\{.*\}', storyboard, re.DOTALL).group()
-            # parsed = json.loads(storyboard)
-            parsed = json.loads(sample_storyboard)
-        except Exception as e:
-            # st.warning(f"Using sample storyboard due to error: {str(e)}")
-            # parsed = json.loads(sample_storyboard)  # 回退到示例数据
-            # status.update(label="⚠️ Using sample storyboard", status="warning")
-            pass
-
-        finally:
-        #     status.update(label="Storyboard ready", status="complete")
-        #     st.subheader("Storyboard Script")
-        #     st.json(parsed)
-            # st.caption("*Note: If using sample data, edit the JSON before proceeding*")
-            pass
-
-
-    # 3. 生成配乐
-    # with st.spinner("🎵 正在生成配乐..."):
-    #     audio = music_pipe(
-    #         "epic sci-fi background music", 
-    #         max_new_tokens=256
-    #     )
-    #     audio_buffer = BytesIO()
-    #     audio["audio"].export(audio_buffer, format="wav")
-    
-    # st.subheader("🎧 背景音乐")
-    # st.audio(audio_buffer, format="audio/wav")
-
-    # 4. 生成分镜图（可选）
-    # if st.toggle("🎨 生成分镜预览图（需要2-3分钟）"):
-    #         with st.spinner("🖼️ 正在渲染电影级画面..."):
-    #             # 从剧本提取关键描述
-    #             scene_desc = re.search(r"场景\d+：(.*?)\n", script)
-    #             if scene_desc:
-    #                 scene_desc = scene_desc.group(1)
-    #             else:
-    #                 scene_desc = f"movie scene of {user_input}"
-                
-    #             # 精细化图片提示词
-    #             image_prompt = f"""电影剧照，{scene_desc}，
-    #                 超高清8K分辨率，
-    #                 电影级打光，浅景深效果，
-    #                 Unreal Engine 5渲染，
-    #                 电影胶片颗粒感"""
-                
-    #             # 优化生成参数
-    #             image = image_pipe(
-    #                 image_prompt,
-    #                 num_inference_steps=50,
-    #                 guidance_scale=9.5,
-    #                 width=1024,
-    #                 height=576  # 16:9电影比例
-    #             ).images[0]
-                
-    #             st.subheader("🎞️ 首场景预览")
-    #             st.image(image, caption=scene_desc)
-                
-    # except Exception as e:
-    #     st.error("生成过程中发生错误")
-    #     st.code(f"错误详情：{str(e)}")
-    #     st.button("点击查看技术详情", 
-    #              help="将以下信息提供给技术人员：\n" + str(e))
-    
-st.subheader("🎧 Background Music")
-for i, scene in enumerate(parsed["scenes"]):
-    action = scene["action"]
-    with st.spinner(f"🎵 Generating soundtrack for scene {i + 1}..."):
-        try:
-            audio = music_pipe(
-                action,
-                max_new_tokens=256
+            storyboard = (
+                storyboard.replace("'", '"')
+                .replace("None", '""')
+                .replace("True", "true")
+                .replace("False", "false")
             )
-            audio_buffer = BytesIO()
-            audio["audio"].export(audio_buffer, format="wav")
-            st.audio(audio_buffer, format="audio/wav", caption=f"Soundtrack for scene {i + 1}")
+            storyboard = re.search(r'\{.*\}', storyboard, re.DOTALL).group()
+            parsed = json.loads(storyboard)
+
         except Exception as e:
-            st.error(f"Failed to generate soundtrack for scene {i + 1}: {str(e)}")
+            st.warning(f"Using sample storyboard due to error: {str(e)}")
+            parsed = json.loads(sample_storyboard)
+            status.update(label="⚠️ Using sample storyboard")
+
+        status.update(label="Storyboard ready")
+        st.subheader("Storyboard Script")
+        st.json(parsed)
+        st.caption("*Note: If using sample data, edit the JSON before proceeding*")
 
-st.subheader("🎨 Storyboard Images")
-for i, scene in enumerate(parsed["scenes"]):
-    location = scene["location"]
-    if st.toggle(f"Generate preview for scene {i + 1} (takes 2 - 3 minutes)"):
-        with st.spinner(f"🖼️ Rendering image for scene {i + 1}..."):
+    st.subheader("🎧 Background Music")
+    for i, scene in enumerate(parsed["scenes"]):
+        action = scene["action"]
+        with st.spinner(f"🎵 Generating soundtrack for scene {i + 1}..."):
             try:
-                # 精细化图片提示词
-                image_prompt = f"""Movie still, {location},
-                    Ultra - high - definition 8K resolution,
-                    Movie - level lighting, shallow depth - of - field effect,
-                    Unreal Engine 5 rendering,
-                    Movie film graininess"""
-
-                # 优化生成参数
-                image = image_pipe(
-                    image_prompt,
-                    num_inference_steps=50,
-                    guidance_scale=9.5,
-                    width=1024,
-                    height=576  # 16:9 movie ratio
-                ).images[0]
-
-                st.image(image, caption=f"Preview for scene {i + 1}")
+                audio = music_pipe(
+                    action,
+                    max_new_tokens=256
+                )
+                audio_buffer = BytesIO()
+                audio["audio"].export(audio_buffer, format="wav")
+                st.audio(audio_buffer, format="audio/wav", caption=f"Soundtrack for scene {i + 1}")
             except Exception as e:
-                st.error(f"Failed to generate image for scene {i + 1}: {str(e)}")
+                st.error(f"Failed to generate soundtrack for scene {i + 1}: {str(e)}")
+
+    st.subheader("🎨 Storyboard Images")
+    for i, scene in enumerate(parsed["scenes"]):
+        location = scene["location"]
+        if st.toggle(f"Generate preview for scene {i + 1} (takes 2 - 3 minutes)"):
+            with st.spinner(f"🖼️ Rendering image for scene {i + 1}..."):
+                try:
+                    # 精细化图片提示词
+                    image_prompt = f"""Movie still, {location},
+                        Ultra - high - definition 8K resolution,
+                        Movie - level lighting, shallow depth - of - field effect,
+                        Unreal Engine 5 rendering,
+                        Movie film graininess"""
+
+                    # 优化生成参数
+                    image = image_pipe(
+                        image_prompt,
+                        num_inference_steps=50,
+                        guidance_scale=9.5,
+                        width=1024,
+                        height=576  # 16:9 movie ratio
+                    ).images[0]
+
+                    st.image(image, caption=f"Preview for scene {i + 1}")
+                except Exception as e:
+                    st.error(f"Failed to generate image for scene {i + 1}: {str(e)}")
+