Update app.py

app.py

@@ -9,36 +9,36 @@ import json
 
 # 初始化模型 (全部使用原始pipeline)
 @st.cache_resource
-def load_pipelines():
+def load_model():
     try:
         # 1. 剧本生成（T5模型）
         script_pipe = pipeline(
-            "text2text-generation", 
+            "text2text-generation",
             model="mrm8488/t5-base-finetuned-common_gen",
-            tokenizer="t5-base"
+            device=0 if torch.cuda.is_available() else -1
         )
         
-        # 2. 分镜生成（BART模型）
-        storyboard_pipe = pipeline(
-            "text2text-generation",  # 使用text2text而非summarization
-            model="philschmid/bart-large-cnn-samsum"
-        )
-        
-        # 3. 配乐生成（MusicGen）
-        # music_pipe = pipeline(
-        #     "text-to-audio", 
-        #     model="facebook/musicgen-small",
-        #     device_map="auto"
+        # # 2. 分镜生成（BART模型）
+        # storyboard_pipe = pipeline(
+        #     "text2text-generation",  # 使用text2text而非summarization
+        #     model="philschmid/bart-large-cnn-samsum"
         # )
         
-        # 4. 分镜图片生成（Stable Diffusion）
-        image_pipe = StableDiffusionPipeline.from_pretrained(
-            "prompthero/openjourney-v4",
-            safety_checker=None  # 禁用安全检查以加速
-        )
+        # # 3. 配乐生成（MusicGen）
+        # # music_pipe = pipeline(
+        # #     "text-to-audio", 
+        # #     model="facebook/musicgen-small",
+        # #     device_map="auto"
+        # # )
         
-        return script_pipe, storyboard_pipe, image_pipe
-        # return script_pipe, storyboard_pipe, music_pipe, image_pipe
+        # # 4. 分镜图片生成（Stable Diffusion）
+        # image_pipe = StableDiffusionPipeline.from_pretrained(
+        #     "prompthero/openjourney-v4",
+        #     safety_checker=None  # 禁用安全检查以加速
+        # )
+        
+        # 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)}")
@@ -59,41 +59,37 @@ if user_input:
     # 1. 生成剧本
     with st.status("🖋️ 剧本生成中...", expanded=True) as status:
         try:
-            # 生成提示词
-            prompt = f"""生成电影剧本（主题：{user_input}）
-            ### 要求：
-            - Markdown格式
-            - 3个场景（INT/EXT）
-            - 包含动作描述和对话
-            
-            示例：
-            ### 场景1：实验室（INT.夜）
-            [警报灯闪烁]  
-            博士（擦汗）：实验体要失控了！"""
-
-            # 生成并处理
-            script = script_pipe(
-                prompt,
-                max_length=400,  # 缩短生成长度
-                num_beams=1,     # 禁用束搜索
-                do_sample=True,  # 启用随机采样
-                temperature=0.7  # 平衡创造性与稳定性
-            )[0]["generated_text"]
-            
-            script = re.sub(r"场景\d+", r"### \g<0>", script)  # 统一标题格式
-            
-            # 更新状态
-            status.update(label="✅ 生成完成", state="complete")
-            st.subheader("生成剧本")
-            st.markdown(f"```markdown\n{script}\n```")
+            # 加载模型
+                pipe = load_model()
                 
-        except torch.cuda.OutOfMemoryError:
-            status.update(label="❌ 显存不足", state="error")
-            st.error("请尝试更简短的输入")
-        except Exception as e:
-            status.update(label="❌ 生成错误", state="error")
-            st.error(str(e))
-
+                # 生成内容
+                prompt = build_prompt(theme)
+                response = pipe(
+                    prompt,
+                    max_length=600,
+                    temperature=0.6,
+                    num_beams=2,
+                    no_repeat_ngram_size=2
+                )[0]["generated_text"]
+                
+                # 清洗与验证
+                cleaned = format_script(response)
+                if "###" not in cleaned:
+                    raise ValueError("生成内容不符合格式要求")
+                
+                # 显示结果
+                st.subheader("生成剧本")
+                st.markdown(f"```markdown\n{cleaned}\n```")
+                status.update(label="✅ 生成完成", state="complete")
+                
+            except Exception as e:
+                status.update(label="❌ 生成失败", state="error")
+                st.error(f"错误详情：{str(e)}")
+                st.code(f"原始输出：{response[:200]}")
+# 运行程序
+if __name__ == "__main__":
+    main()
+    
     # 2. 生成分镜
      # with st.status("🎥 正在转换分镜脚本...", expanded=True) as status:
      #    try: