Update app.py

app.py

@@ -4,28 +4,30 @@ from transformers import pipeline
 from PIL import Image
 
 # function part
-# img2text - Using a lighter model
+# img2text - Using the original model
 def img2text(image):
-    # Use a smaller, faster image captioning model
-    image_to_text = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
-    text = image_to_text(image, max_new_tokens=20)[0]["generated_text"]
+    # Use the specified model but with optimized parameters
+    image_to_text = pipeline("image-to-text", model="sooh-j/blip-image-captioning-base")
+    # Limiting the output length for speed
+    text = image_to_text(image, max_new_tokens=30)[0]["generated_text"]
     return text
 
-# text2story - Using a much faster model with constraints
+# text2story - Using the original model but with optimized parameters
 def text2story(text):
-    # Use a tiny model that's much faster
-    generator = pipeline("text-generation", model="distilgpt2")
+    # Using the specified TinyLlama model
+    generator = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
     
-    # Create a more constrained prompt for faster generation
-    prompt = f"A short children's story about {text}: Once upon a time, "
+    # Create a prompt for the story generation
+    prompt = f"Write a brief children's story based on this: {text}. Once upon a time, "
     
-    # Generate with strict constraints for speed
+    # Generate with more constrained parameters for speed
     story_result = generator(
         prompt,
-        max_new_tokens=100,  # Limit token generation
+        max_new_tokens=150,  # Use max_new_tokens instead of max_length for efficiency
         num_return_sequences=1,
         temperature=0.7,
         top_k=50,
+        top_p=0.95,
         do_sample=True
     )
    
@@ -33,17 +35,26 @@ def text2story(text):
     story_text = story_result[0]['generated_text']
     story_text = story_text.replace(prompt, "Once upon a time, ")
     
-    # Find a natural ending point (end of sentence)
-    last_period = story_text.rfind('.')
-    last_question = story_text.rfind('?')
-    last_exclamation = story_text.rfind('!')
-    
-    # Find the last sentence ending punctuation
-    last_end = max(last_period, last_question, last_exclamation)
-    
-    if last_end > 0:
-        # Truncate at the end of the last complete sentence
-        story_text = story_text[:last_end + 1]
+    # Find a natural ending point (end of sentence) before 100 words
+    words = story_text.split()
+    if len(words) > 100:
+        # Join the first 100 words
+        shortened_text = " ".join(words[:100])
+        
+        # Find the last complete sentence
+        last_period = shortened_text.rfind('.')
+        last_question = shortened_text.rfind('?')
+        last_exclamation = shortened_text.rfind('!')
+        
+        # Find the last sentence ending punctuation
+        last_end = max(last_period, last_question, last_exclamation)
+        
+        if last_end > 0:
+            # Truncate at the end of the last complete sentence
+            story_text = shortened_text[:last_end + 1]
+        else:
+            # If no sentence ending found, just use the shortened text
+            story_text = shortened_text
     
     return story_text