Update app.py

app.py

@@ -8,7 +8,8 @@ import os
 def img2text(img_path):
     captioner = pipeline(
         "image-to-text", 
-        model="Salesforce/blip-image-captioning-base"
+        model="nlpconnect/vit-gpt2-image-captioning",
+        device_map="auto"
     )
     result = captioner(img_path)
     return result[0]["generated_text"]
@@ -16,11 +17,11 @@ def img2text(img_path):
 def text2story(scenario):
     generator = pipeline(
         "text-generation",
-        model="gpt2-medium",
+        model="gpt2",
         trust_remote_code=True,
         device_map="auto",    # 必须添加
         torch_dtype=torch.float16, 
-        max_length=150,
+        max_length=100,
         num_return_sequences=1
     )
     prompt = f"Create a children's story based on: {scenario}"