Update app.py

app.py

@@ -14,7 +14,7 @@ import numpy as np
 from PIL import Image
 
 # 初始化模型（CPU优化版）
-@st.cache_resource(show_spinner="🔮 Loading magic models...")
+@st.cache_resource
 def load_models():
     """加载所有需要的AI模型"""
     try:
@@ -26,8 +26,7 @@ def load_models():
         story_generator = pipeline(
             "text-generation",
             model="openai-community/gpt2",
-            device_map="auto",
-            torch_dtype=torch.float32
+            device_map="auto"
         )
         
         # 语音合成模型
@@ -70,8 +69,8 @@ Story:"""
     # 使用GPT-2生成故事
     generated = story_generator(
         prompt,
-        max_length=300,
-        min_length=150,
+        max_length=100,
+        min_length=50,
         num_return_sequences=1,
         temperature=0.85,
         repetition_penalty=2.0
@@ -80,25 +79,20 @@ Story:"""
     # 提取生成文本并清理
     full_text = generated[0]['generated_text']
     story = full_text.split("Story:")[-1].strip()
-    return story[:580].replace(caption, "").strip()
+    return story[:600].replace(caption, "").strip()
 
 def text_to_speech(text, processor, model, vocoder, embeddings_dataset):
-    """文本转语音（修复版）"""
+    """文本转语音"""
     try:
-        # 输入预处理
         inputs = processor(
-            text=text,
+            text=text, 
             return_tensors="pt",
-            padding="max_length",
-            max_length=600,
-            truncation=True,
             voice_preset=None
         )
         input_ids = inputs["input_ids"].to(torch.int64)
         
-        # 动态调整说话者嵌入维度
-        speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"])
-        speaker_embeddings = speaker_embeddings.unsqueeze(0).repeat(1, input_ids.shape[1], 1)
+        # 随机选择一个说话者嵌入
+        speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)
         
         with torch.no_grad():
             speech = model.generate_speech(
@@ -107,11 +101,8 @@ def text_to_speech(text, processor, model, vocoder, embeddings_dataset):
                 vocoder=vocoder
             )
         
-        # 音频处理
-        audio_array = speech.numpy().astype(np.float32)
-        max_val = np.max(np.abs(audio_array)) + 1e-8
-        audio_array = 0.9 * audio_array / max_val
-        
+        audio_array = speech.numpy()
+        audio_array = audio_array / np.max(np.abs(audio_array))
         return audio_array, 16000
     except Exception as e:
         st.error(f"语音生成失败: {str(e)}")
@@ -133,11 +124,46 @@ def main():
     if 'generated' not in st.session_state:
         st.session_state.generated = False
     
-    # 加载模型（保持不变...）
+    # 加载模型
+    try:
+        (blip_proc, blip_model, story_gen, 
+         tts_proc, tts_model, vocoder, embeddings) = load_models()
+    except:
+        return
     
-    # 文件上传组件（保持不变...）
+    # 文件上传组件
+    uploaded_file = st.file_uploader(
+        "Choose your magic image",
+        type=["jpg", "png", "jpeg"],
+        help="Upload photos of pets, toys or adventures!",
+        key="uploader"
+    )
     
-    # 处理上传文件（保持不变...）
+    # 处理上传文件
+    if uploaded_file and not st.session_state.generated:
+        try:
+            image = Image.open(uploaded_file).convert("RGB")
+            st.image(image, caption="Your Magic Picture ✨", use_container_width=True)
+            
+            with st.status("Creating Magic...", expanded=True) as status:
+                # 生成故事
+                st.write("🔍 Reading the image...")
+                story = generate_story(image, blip_proc, blip_model, story_gen)
+                
+                # 生成语音
+                st.write("🔊 Adding sounds...")
+                audio_array, sr = text_to_speech(story, tts_proc, tts_model, vocoder, embeddings)
+                
+                # 保存结果
+                st.session_state.story = story
+                st.session_state.audio = (audio_array, sr)
+                status.update(label="Ready!", state="complete", expanded=False)
+            
+            st.session_state.generated = True
+            st.rerun()
+            
+        except Exception as e:
+            st.error(f"Magic failed: {str(e)}")
     
     # 显示结果
     if st.session_state.generated: