Update app.py

app.py

@@ -12,74 +12,95 @@ import torch
 @st.cache_resource
 def load_pipelines():
     try:
-        # 1. 剧本生成
+        # 1. Script generation
         script_pipe = pipeline(
             "text2text-generation",
             model="RUCAIBox/mvp-story",
             device=0 if torch.cuda.is_available() else -1
         )
-        
-        # 2. 分镜生成（BART模型）
+
+        # 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. 配乐生成（MusicGen）
+
+        # # 3. Soundtrack generation (MusicGen)
         # # music_pipe = pipeline(
         # #     "text-to-audio", 
         # #     model="facebook/musicgen-small",
         # #     device_map="auto"
         # # )
-        
-        # 4. 分镜图片生成（Stable Diffusion）
+
+        # 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"模型加载失败: {str(e)}")
+        st.error(f"Failed to load models: {str(e)}")
         raise
 
-#提示词模板
-def build_script_prompt(theme):
-    return f"""{theme}"""
+# 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)
-    # 规范场景标题
-    return re.sub(r"(?<!### )场景\d+", r"### 场景\g<0>", text)
+    # Standardize scene titles
+    return re.sub(r"(?<!### )Scene\d+", r"### Scene\g<0>", text)
 
-# 加载pipeline
+# Load pipelines
 script_pipe, storyboard_pipe, image_pipe = load_pipelines()
 # script_pipe, storyboard_pipe, music_pipe, image_pipe = load_pipelines()
 
-# Streamlit界面
+# Streamlit interface
 st.title("🎥 Micro - movie Intelligent Creation System")
-    
-#用户输入
+
+# User input
 user_input = st.text_input(
-    "Enter the movie theme or the start of the movie (e.g., Sci - fi space adventure or A girl is missing)",
+    "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'"
 )
 
-if user_input:
-       
+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. 生成剧本
+    # 1. Generate the script
     with st.status("🖋️ Generating the script...", expanded=True) as status:
         try:
-            # 生成内容
-            prompt = build_script_prompt(user_input)
+            # Generate content
+            prompt = build_script_prompt(user_input, style_or_opening)
             response = script_pipe(
                 prompt,
                 max_length=600,
@@ -87,34 +108,33 @@ if user_input:
                 num_beams=4,
                 no_repeat_ngram_size=3
             )[0]["generated_text"]
-                    
-            # 后处理
+
+            # Post - processing
             def format_script(text):
-                # # 强制修正场景标题格式（如 "场景1" -> "### 场景1"）
-                # text = re.sub(r"场景(\d+)", r"### 场景\1", text)
-                # # 删除多余符号
+                # # Force - correct the scene title format (e.g., "Scene 1" -> "### Scene 1")
+                # text = re.sub(r"Scene(\d+)", r"### Scene\1", text)
+                # # Remove extra symbols
                 # text = re.sub(r"[<>()]{2,}|\[?\]?", "", text)
-                # # 确保动作描述包裹在方括号内
+                # # Ensure action descriptions are wrapped in square brackets
                 # return re.sub(r"(?<=\n)\s*([^\[\n]+?)\s*(?=\n)", r"[\1]", text)
-                return
-                    
-            # 显示结果
+                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", state="error")
             st.error(f"Error details: {str(e)}")
             st.stop()
 
-    
-    #2. 生成分镜
+    # 2. Generate the storyboard
     with st.status("🎥 Converting to a storyboard script...", expanded=True) as status:
         try:
-            # JSON模板引导
-             json_template = """{
+            # JSON template guidance
+            json_template = """{
               "scenes": [
                 {
                   "scene_number": 1,
@@ -125,7 +145,7 @@ if user_input:
                 }
               ]
             }"""
-            
+
             # Generate the storyboard
             storyboard = storyboard_pipe(
                 f"Convert the following script into a storyboard JSON, strictly following the template:\n{json_template}\nScript content: {cleaned[:800]}",  # Key modification point