Szeyu commited on
Commit
bc7ea9a
Β·
verified Β·
1 Parent(s): be6c328

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +42 -22
app.py CHANGED
@@ -1,4 +1,4 @@
1
- # import part
2
  import streamlit as st
3
  from transformers import pipeline
4
  import textwrap
@@ -19,24 +19,52 @@ def load_pipelines():
19
 
20
  captioner, storyer, tts = load_pipelines()
21
 
22
- # Function part
23
- # Function to generate content from an image
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
24
  def generate_content(image):
25
  pil_image = Image.open(image)
26
 
27
- # Generate caption
28
  caption = captioner(pil_image)[0]["generated_text"]
29
  st.write("**🌟 What's in the picture: 🌟**")
30
  st.write(caption)
31
 
32
- # Create prompt for story
 
 
33
  prompt = (
34
- f"Write a funny, interesting children's story that is precisely centered on this scene {caption}\nStory:"
35
  f" mention the exact place, location or venue within {caption}"
36
- f" in third-person narrative, that describes this scene exactly: {caption} "
37
- f" Avoid numbers, random letter combinations, and single-letter combinations."
38
  )
39
- # Generate raw story
 
40
  raw = storyer(
41
  prompt,
42
  max_new_tokens=150,
@@ -45,38 +73,30 @@ def generate_content(image):
45
  no_repeat_ngram_size=2,
46
  return_full_text=False
47
  )[0]["generated_text"].strip()
48
-
49
- # Define allowed characters to keep (removes digits and symbols like * and ~)
50
- allowed_chars = string.ascii_letters + " .,!?\"'-"
51
-
52
- # Clean the raw story by keeping only allowed characters (this filters out any digits)
53
- clean_raw = ''.join(c for c in raw if c in allowed_chars)
54
 
55
- # Split into words and trim to 100 words
56
- words = clean_raw.split()
57
- story = " ".join(words[:100])
58
 
59
  st.write("**πŸ“– Your funny story: πŸ“–**")
60
  st.write(story)
61
 
62
- # Generate audio from cleaned story
63
  chunks = textwrap.wrap(story, width=200)
64
  audio = np.concatenate([tts(chunk)["audio"].squeeze() for chunk in chunks])
65
 
66
- # Save audio to temporary file
67
  with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_file:
68
  sf.write(temp_file.name, audio, tts.model.config.sampling_rate)
69
  temp_file_path = temp_file.name
70
 
71
  return caption, story, temp_file_path
72
 
73
- # Streamlit UI
74
  st.title("✨ Magic Story Maker ✨")
75
  st.markdown("Upload a picture to make a funny story and hear it too! πŸ“Έ")
76
 
77
  uploaded_image = st.file_uploader("Choose your picture", type=["jpg", "jpeg", "png"])
78
 
79
- # Streamlit UI (modified image display section)
80
  if uploaded_image is None:
81
  st.image("https://example.com/placeholder_image.jpg", caption="Upload your picture here! πŸ“·", use_container_width=True)
82
  else:
 
1
+ import re
2
  import streamlit as st
3
  from transformers import pipeline
4
  import textwrap
 
19
 
20
  captioner, storyer, tts = load_pipelines()
21
 
22
+ def clean_generated_story(raw_story: str) -> str:
23
+ """
24
+ Cleans the generated story by:
25
+ 1. Removing digits.
26
+ 2. Removing words that are likely random letter combinations based on having no vowels.
27
+ 3. Removing single-letter words unless they are allowed (such as 'a' or 'I').
28
+ """
29
+ # Remove all digits using regex
30
+ story_without_numbers = re.sub(r'\d+', '', raw_story)
31
+
32
+ vowels = set('aeiouAEIOU')
33
+
34
+ def is_valid_word(word: str) -> bool:
35
+ # Allow "a" and "I" for single-letter words
36
+ if len(word) == 1 and word.lower() not in ['a', 'i']:
37
+ return False
38
+ # For words longer than one letter, filter out those that do not contain any vowels
39
+ if len(word) > 1 and not any(char in vowels for char in word):
40
+ return False
41
+ return True
42
+
43
+ # Split the story into words, apply filtering, and recombine into a clean story
44
+ words = story_without_numbers.split()
45
+ filtered_words = [word for word in words if is_valid_word(word)]
46
+
47
+ # Optionally, you can trim the clean story to a certain word count
48
+ clean_story = " ".join(filtered_words[:100])
49
+ return clean_story
50
+
51
  def generate_content(image):
52
  pil_image = Image.open(image)
53
 
54
+ # Generate caption from the image
55
  caption = captioner(pil_image)[0]["generated_text"]
56
  st.write("**🌟 What's in the picture: 🌟**")
57
  st.write(caption)
58
 
59
+ # Create prompt for the story
60
+ # Notice there’s no need to include the extra cleaning instructions in this prompt,
61
+ # because our code handles them later.
62
  prompt = (
63
+ f"Write a funny, interesting story for young children precisely centered on this scene {caption}\nStory:"
64
  f" mention the exact place, location or venue within {caption}"
 
 
65
  )
66
+
67
+ # Generate raw story from the model
68
  raw = storyer(
69
  prompt,
70
  max_new_tokens=150,
 
73
  no_repeat_ngram_size=2,
74
  return_full_text=False
75
  )[0]["generated_text"].strip()
 
 
 
 
 
 
76
 
77
+ # Clean the raw story using our custom function
78
+ story = clean_generated_story(raw)
 
79
 
80
  st.write("**πŸ“– Your funny story: πŸ“–**")
81
  st.write(story)
82
 
83
+ # Generate audio for the story
84
  chunks = textwrap.wrap(story, width=200)
85
  audio = np.concatenate([tts(chunk)["audio"].squeeze() for chunk in chunks])
86
 
87
+ # Save audio to a temporary file
88
  with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_file:
89
  sf.write(temp_file.name, audio, tts.model.config.sampling_rate)
90
  temp_file_path = temp_file.name
91
 
92
  return caption, story, temp_file_path
93
 
94
+ # Streamlit UI section
95
  st.title("✨ Magic Story Maker ✨")
96
  st.markdown("Upload a picture to make a funny story and hear it too! πŸ“Έ")
97
 
98
  uploaded_image = st.file_uploader("Choose your picture", type=["jpg", "jpeg", "png"])
99
 
 
100
  if uploaded_image is None:
101
  st.image("https://example.com/placeholder_image.jpg", caption="Upload your picture here! πŸ“·", use_container_width=True)
102
  else: