Update app.py

app.py

@@ -1,111 +1,78 @@
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
-from PIL import Image
-import json
-import io
-from zipfile import ZipFile
-from musicgen.model import MusicGen
-
-# 模型加载（使用pipeline）
-script_generator = pipeline(
-    "text2text-generation",
-    model="mrm8488/t5-base-finetuned-common-gen",
-    tokenizer="t5-base"
-)
-
-shot_generator = pipeline(
-    "text2text-generation",
-    model="philschmid/bart-large-cnn-samsum",
-    tokenizer="facebook/bart-large"
-)
-
-# 音乐生成器（需要单独处理）
-musicgen = MusicGen.get_pretrained("facebook/musicgen-small")
-
-# 图片生成器（需要diffusers）
+from transformers import pipeline
 from diffusers import StableDiffusionPipeline
-image_generator = StableDiffusionPipeline.from_pretrained(
-    "prompthero/openjourney-v4",
-    torch_dtype=torch.float16
-).to("cuda")  # 或"cpu"
-
-
-def generate_script(topic):
-    """生成剧本"""
-    input_text = f"generate script: {topic}"
-    output = script_generator(input_text, max_length=200)[0]['generated_text']
-    return output
+import torch
+from pydub import AudioSegment
+import base64
+import json
+import zipfile
+from io import BytesIO
+
+# 初始化模型Pipeline（全部使用Hugging Face预训练模型）
+@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")
+    
+    return script_pipe, storyboard_pipe, music_pipe, image_pipe
 
+# 加载模型
+script_pipe, storyboard_pipe, music_pipe, image_pipe = load_models()
 
-def generate_shots(script):
-    """生成分镜"""
-    input_text = f"summarize shot: {script}"
-    output = shot_generator(input_text, max_length=200)[0]['generated_text']
+# Streamlit界面
+st.title("🎬 微电影创作助手（基础版）")
+user_input = st.text_input("输入电影主题或关键词（例如：浪漫的咖啡馆相遇）")
+
+if user_input:
+    # 生成剧本
+    with st.spinner("正在生成剧本..."):
+        script = script_pipe(
+            f"Generate a movie script about: {user_input}",
+            max_length=300
+        )[0]["generated_text"]
+    
+    st.subheader("📜 生成的剧本")
+    st.markdown(f"```\n{script}\n```")
     
-    # 简单解析JSON（实际需更健壮的解析）
+    # 生成分镜
+    with st.spinner("正在生成分镜脚本..."):
+        storyboard = storyboard_pipe(
+            f"Convert this script to storyboard JSON: {script}",
+            max_length=500
+        )[0]["generated_text"]
+    
+    st.subheader("🎥 分镜脚本")
     try:
-        return json.loads(output.replace("```json", "").replace("```", ""))
+        st.json(json.loads(storyboard))
     except:
-        return {"镜头1": {"类型": "中景", "角度": "水平视角", "时长": "5秒"}}
-
-
-def generate_music(shot_description):
-    """生成音乐（简化版）"""
-    descriptions = [f"background music for {shot_description} scene"]
-    wav = musicgen.generate(descriptions, duration=10)  # 10秒示例
-    return wav[0]
-
-
-def generate_image(prompt):
-    """生成图片"""
-    image = image_generator(prompt).images[0]
-    return image
-
-
-# Streamlit界面
-st.title("电影分镜生成器")
-
-# 输入区
-user_topic = st.text_input("请输入场景主题（例如：咖啡馆相遇场景）")
-
-if st.button("生成内容"):
-    if not user_topic:
-        st.warning("请输入场景主题！")
-    else:
-        with st.spinner("生成中..."):
-            # 并行生成（简化为顺序执行）
-            script = generate_script(user_topic)
-            shots = generate_shots(script)
-            
-            # 生成示例图片和音乐
-            sample_image = generate_image(f"电影分镜：{user_topic}")
-            sample_audio = generate_music(user_topic)
-
-        # 结果展示
-        st.header("生成剧本")
-        st.markdown(script)
-
-        st.header("分镜脚本")
-        st.json(shots)
-
-        st.header("配乐预览")
-        audio_bytes = io.BytesIO()
-        sample_audio.save(audio_bytes, format="WAV")
-        st.audio(audio_bytes)
-
-        st.header("分镜图片预览")
-        st.image(sample_image, caption="示例分镜画面")
-
-        # 下载功能
-        with ZipFile("output.zip", "w") as zipf:
-            zipf.writestr("script.md", script)
-            zipf.writestr("shots.json", json.dumps(shots))
-            sample_image.save(io.BytesIO(), format="PNG")
-        
-        with open("output.zip", "rb") as f:
-            st.download_button(
-                label="下载所有文件",
-                data=f,
-                file_name="movie_scenes.zip",
-                mime="application/zip"
-            )
+        st.error("分镜生成格式错误，请重试！")
+    
+    # 生成配乐
+    with st.spinner("正在生成背景音乐..."):
+        audio = music_pipe("Calm background music for a romantic scene", max_new_tokens=200)
+        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.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")