miltonc commited on
Commit
b352af0
·
verified ·
1 Parent(s): d8861bf

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -7,7 +7,7 @@ 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("text-generation", model="databricks/dolly-v2-3b", max_new_tokens=300)
11
  return image_to_text, storyteller
12
 
13
  # Process image to text
@@ -17,7 +17,7 @@ def generate_caption(image, image_to_text):
17
 
18
  # Generate a narrative story
19
  def generate_story(text, storyteller):
20
- prompt = f"Write a short and engaging story inspired by this image description: {text}"
21
  story = storyteller(prompt, do_sample=True, temperature=0.7, max_new_tokens=300)
22
  return story[0]["generated_text"] if story else "No story generated."
23
 
@@ -36,6 +36,7 @@ def main():
36
  uploaded_file = st.file_uploader("Upload an image...", type=["jpg", "png", "jpeg"])
37
 
38
  if uploaded_file is not None:
 
39
  image = Image.open(uploaded_file)
40
  st.image(image, caption="Uploaded Image", use_container_width=True)
41
 
 
7
  # Load models
8
  def load_models():
9
  image_to_text = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
10
+ storyteller = pipeline("text-generation", model="tiiuae/falcon-7b-instruct", max_new_tokens=300)
11
  return image_to_text, storyteller
12
 
13
  # Process image to text
 
17
 
18
  # Generate a narrative story
19
  def generate_story(text, storyteller):
20
+ prompt = f"Write a short, engaging, and imaginative story based on this image description: {text}"
21
  story = storyteller(prompt, do_sample=True, temperature=0.7, max_new_tokens=300)
22
  return story[0]["generated_text"] if story else "No story generated."
23
 
 
36
  uploaded_file = st.file_uploader("Upload an image...", type=["jpg", "png", "jpeg"])
37
 
38
  if uploaded_file is not None:
39
+ # Convert uploaded file to a PIL image
40
  image = Image.open(uploaded_file)
41
  st.image(image, caption="Uploaded Image", use_container_width=True)
42