Update app.py

app.py

@@ -1,78 +1,111 @@
 import streamlit as st
 from transformers import pipeline
 from diffusers import StableDiffusionPipeline
-import torch
 from pydub import AudioSegment
-import base64
-import json
 import zipfile
 from io import BytesIO
+import json
 
-# 初始化模型Pipeline（使用修正后的模型名称）
+# 初始化模型 (全部使用原始pipeline)
 @st.cache_resource
-def load_models():
-    # 剧本生成（T5模型）
-    script_pipe = pipeline("text2text-generation", model="mrm8488/t5-base-finetuned-common_gen")
-    
-    # 分镜生成（BART模型）
-    storyboard_pipe = pipeline("text-generation", model="philschmid/bart-large-cnn-samsum")
-    
-    # 配乐生成（MusicGen）
-    music_pipe = pipeline("text-to-audio", model="facebook/musicgen-small")
-    
-    # 分镜图片生成（Stable Diffusion）
-    image_pipe = StableDiffusionPipeline.from_pretrained("prompthero/openjourney-v4")
+def load_pipelines():
+    try:
+        # 1. 剧本生成（T5模型）
+        script_pipe = pipeline(
+            "text2text-generation", 
+            model="mrm8488/t5-base-finetuned-common_gen",
+            tokenizer="t5-base"  # 显式指定tokenizer
+        )
+        
+        # 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"
+        )
+        
+        # 4. 分镜图片生成（Stable Diffusion）
+        image_pipe = StableDiffusionPipeline.from_pretrained(
+            "prompthero/openjourney-v4",
+            safety_checker=None  # 禁用安全检查以加速
+        )
+        
+        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"模型加载失败: {str(e)}")
+        raise
 
-# 加载模型
-script_pipe, storyboard_pipe, music_pipe, image_pipe = load_models()
+# 加载pipeline
+script_pipe, storyboard_pipe, music_pipe, image_pipe = load_pipelines()
 
 # Streamlit界面
-st.title("🎬 微电影创作助手（基础版）")
-user_input = st.text_input("输入电影主题或关键词（例如：浪漫的咖啡馆相遇）")
+st.title("🎥 微电影智能创作系统")
+user_input = st.text_input("输入电影主题（如：科幻太空冒险）")
 
 if user_input:
-    # 生成剧本
-    with st.spinner("正在生成剧本..."):
+    # 1. 生成剧本
+    with st.spinner("🤖 正在生成剧本..."):
         script = script_pipe(
-            f"Generate a movie script about: {user_input}",
-            max_length=300
+            f"generate movie script: {user_input}",
+            max_length=500,
+            num_return_sequences=1
         )[0]["generated_text"]
     
-    st.subheader("📜 生成的剧本")
-    st.markdown(f"```\n{script}\n```")
-    
-    # 生成分镜
-    with st.spinner("正在生成分镜脚本..."):
+    st.subheader("📜 电影剧本")
+    st.code(script, language="markdown")
+
+    # 2. 生成分镜
+    with st.spinner("🎬 正在生成分镜..."):
         storyboard = storyboard_pipe(
-            f"Convert this script to storyboard JSON: {script}",
-            max_length=500
+            f"convert to storyboard JSON: {script}",
+            max_length=1000
         )[0]["generated_text"]
     
-    st.subheader("🎥 分镜脚本")
+    st.subheader("📋 分镜脚本")
     try:
-        st.json(json.loads(storyboard))
+        st.json(json.loads(storyboard))  # 验证JSON格式
     except:
-        st.error("分镜生成格式错误，请重试！")
-    
-    # 生成配乐
-    with st.spinner("正在生成背景音乐..."):
-        audio = music_pipe("Calm background music for a romantic scene", max_new_tokens=200)
+        st.warning("分镜格式异常，原始输出：")
+        st.text(storyboard)
+
+    # 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.audio(audio_buffer, format="audio/wav")
     
-    # 生成分镜图片（可选）
-    if st.checkbox("生成分镜预览图"):
-        with st.spinner("正在生成图片..."):
-            image = image_pipe("cinematic scene, 4k film still, " + user_input).images[0]
+    st.subheader("🎧 背景音乐")
+    st.audio(audio_buffer, format="audio/wav")
+
+    # 4. 生成分镜图（可选）
+    if st.toggle("生成分镜预览图"):
+        with st.spinner("🖼️ 正在生成图片..."):
+            image = image_pipe(
+                f"cinematic still from a {user_input} movie, 4k",
+                num_inference_steps=25  # 减少步数以加速
+            ).images[0]
             st.image(image)
-    
-    # 打包下载功能
+
+    # 打包下载
     zip_buffer = BytesIO()
     with zipfile.ZipFile(zip_buffer, "w") as zf:
         zf.writestr("script.txt", script)
         zf.writestr("storyboard.json", storyboard)
         zf.writestr("bgm.wav", audio_buffer.getvalue())
-    st.download_button("📥 下载完整包", data=zip_buffer.getvalue(), file_name="film_package.zip")
+    
+    st.download_button(
+        "📦 下载创作包",
+        data=zip_buffer.getvalue(),
+        file_name="movie_kit.zip",
+        mime="application/zip"
+    )