Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 from transformers import pipeline
+import torch
 from gtts import gTTS
 import tempfile
 import os
@@ -17,7 +18,9 @@ def text2story(scenario):
         "text-generation",
         model="gpt2-medium",
         trust_remote_code=True,
-        max_length=300,
+        device_map="auto",    # 必须添加
+        torch_dtype=torch.float16, 
+        max_length=150,
         num_return_sequences=1
     )
     prompt = f"Create a children's story based on: {scenario}"