Update app.py

app.py

@@ -8,12 +8,10 @@ from PIL import Image
 def load_models():
     image_to_text = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
     storyteller = pipeline(
-        "text2text-generation",
-        model="google/flan-t5-xl",
-        max_length=100,
-        temperature=1,
-        no_repeat_ngram_size=2,   # Prevents repeating 2-grams
-        penalty_alpha=0.5         # Penalty for repeated n-grams
+        "text-generation",
+        model="pranavpsv/gpt2-genre-story-generator",
+        temperature=0.75,
+        max_length=100
     )
     return image_to_text, storyteller
 
@@ -22,11 +20,13 @@ def generate_caption(image, image_to_text):
     result = image_to_text(image)
     return result[0]["generated_text"] if result else "No caption generated."
 
-# Generate a narrative story using an optimized Flan-T5 prompt
+# Generate a narrative story using the GPT-2 genre-based story generator
 def generate_story(text, storyteller):
-    prompt = f'''Write a 50-word story about "{text}".'''
-    story = storyteller(prompt)
-    return story[0]["generated_text"] if story else "No story generated."
+    prompt = f"<BOS> <superhero> {text}"
+    story = storyteller(prompt, max_length=100, num_return_sequences=1)
+    # Clean the generated text by removing the <BOS> <superhero> prefix
+    generated_story = story[0]["generated_text"].replace("<BOS> <superhero>", "").strip()
+    return generated_story if generated_story else "No story generated."
 
 # Convert text to speech
 def text_to_speech(text, filename="output.mp3"):