Update app.py

app.py

@@ -13,11 +13,17 @@ def img2text(url):
         url (str): Path to the image file.
     
     Returns:
-        str: Generated text caption from the image.
+        str: Generated text caption from the image, without words like "illustration".
     """
     image_to_text_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
     text = image_to_text_model(url)[0]["generated_text"]
-    return text
+    
+    # Remove unwanted words like "illustration"
+    unwanted_words = ["illustration", "drawing", "sketch", "picture", "dream", "imagination"]
+    for word in unwanted_words:
+        text = text.replace(word, "")
+    
+    return text.strip()
 
 # Function to generate a kid-friendly story from the text caption
 def text2story(text):
@@ -28,24 +34,24 @@ def text2story(text):
         text (str): Text caption generated from the image.
     
     Returns:
-        str: Generated story suitable for kids aged 3-10, within 95 words.
+        str: Generated story suitable for kids aged 3-10, within 100 words.
     """
-    # Add a prompt to ensure the story is kid-friendly and happy
-    prompt = f"Write a happy, fun and complete story for kids aged 3-10 based on the following scenario: {text}. Keep the story under 120 words."
+    # Add a prompt to ensure the story is happy, fun, and appropriate for kids
+    prompt = f"Write a happy and fun story for kids aged 3-10 based on the following scenario: {text}. The story should be cheerful, with no sad, violent, or scary elements. Avoid any mention of guns, weapons, or conflict. The story must have a clear beginning, middle, and end. Keep the story under 100 words."
     
     # Load the text generation model
     story_generator = pipeline("text-generation", model="aspis/gpt2-genre-story-generation")
     
     # Generate the story
-    story = story_generator(prompt, max_length=150, num_return_sequences=1)[0]["generated_text"]
+    story = story_generator(prompt, max_length=150, num_return_sequences=1, temperature=0.7, top_k=50, top_p=0.9)[0]["generated_text"]
     
     # Remove the prompt from the generated story
     story = story.replace(prompt, "").strip()
     
     # Ensure the story is within 95 words
     story_words = story.split()
-    if len(story_words) > 98:
-        story = " ".join(story_words[:98])
+    if len(story_words) > 100:
+        story = " ".join(story_words[:100])
     
     return story