Update app.py

app.py

@@ -12,7 +12,7 @@ def load_models():
     """Load all models and cache them for faster execution"""
     models = {
         "image_captioner": pipeline("image-to-text", model="sooh-j/blip-image-captioning-base"),
-        "story_generator": pipeline("text-generation", model="Qwen/Qwen2-0.5B-Instruct")
+        "story_generator": pipeline("text-generation", model="trl-internal-testing/tiny-MistralForCausalLM-0.2")
     }
     return models
 
@@ -49,9 +49,9 @@ def text2story(caption, _models):
     # Updated prompt includes explicit instructions to incorporate details from the caption
     prompt = (
         "<|system|>\n"
-        "You are a creative story generating assistant that only gives me a happy story based on image caption given, no reply or quotation are needed which means that you don't need to answer but just give me the story with no quotation.\n"
+        "You are a creative story generating assistant.\n"
         "<|user|>\n"
-        f"Please generate a short story within 100 words that clearly utilizes the following image caption and expands upon it thoroughly:\n\"{caption}\"\n"
+        f"Please generate a short story within 100 words that clearly utilizes the following image caption and expands upon it with no introduction or closing remarks:\n\"{caption}\"\n"
         "<|assistant|>"
     )