Update app.py

app.py

@@ -15,7 +15,7 @@ except ImportError:
     import torch
 
 # Load the image captioning model
-caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
+caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-large")
 
 # Load the text generation model
 text_generation_model = AutoModelForCausalLM.from_pretrained("gpt2")
@@ -28,8 +28,9 @@ def generate_caption(image):
 
 def generate_story(caption):
     # Generate the story based on the caption
-    input_ids = tokenizer.encode(caption, return_tensors="pt")
-    output = text_generation_model.generate(input_ids, max_length=100, num_return_sequences=1)
+    prompt = f"Imagine you are a storyteller for young children. Based on the image described as '{caption}', create a short and interesting story for children aged 3-10. Keep it positive and happy in tone."
+    input_ids = tokenizer.encode(prompt, return_tensors="pt")
+    output = text_generation_model.generate(input_ids, max_length=200, num_return_sequences=1)
     story = tokenizer.decode(output[0], skip_special_tokens=True)
     return story