miltonc commited on
Commit
599621d
·
verified ·
1 Parent(s): e07695d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +9 -2
app.py CHANGED
@@ -7,7 +7,14 @@ from PIL import Image
7
  # Load models
8
  def load_models():
9
  image_to_text = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
10
- storyteller = pipeline("text2text-generation", model="google/flan-t5-small", max_length=100) # Max tokens set to 100
 
 
 
 
 
 
 
11
  return image_to_text, storyteller
12
 
13
  # Process image to text
@@ -17,7 +24,7 @@ def generate_caption(image, image_to_text):
17
 
18
  # Generate a narrative story using an optimized Flan-T5 prompt
19
  def generate_story(text, storyteller):
20
- prompt = f"Write a 75-word creative and engaging short story based on the following description: {text}."
21
  story = storyteller(prompt)
22
  return story[0]["generated_text"] if story else "No story generated."
23
 
 
7
  # Load models
8
  def load_models():
9
  image_to_text = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
10
+ storyteller = pipeline(
11
+ "text2text-generation",
12
+ model="google/flan-t5-xl",
13
+ max_length=100,
14
+ temperature=1,
15
+ no_repeat_ngram_size=2, # Prevents repeating 2-grams
16
+ penalty_alpha=0.5 # Penalty for repeated n-grams
17
+ )
18
  return image_to_text, storyteller
19
 
20
  # Process image to text
 
24
 
25
  # Generate a narrative story using an optimized Flan-T5 prompt
26
  def generate_story(text, storyteller):
27
+ prompt = f'''Write a 50-word story about "{text}".'''
28
  story = storyteller(prompt)
29
  return story[0]["generated_text"] if story else "No story generated."
30