Update app.py
Browse files
app.py
CHANGED
@@ -1,4 +1,4 @@
|
|
1 |
-
|
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 |
-
|
23 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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
|
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 |
-
|
|
|
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 |
-
#
|
56 |
-
|
57 |
-
story = " ".join(words[:100])
|
58 |
|
59 |
st.write("**π Your funny story: π**")
|
60 |
st.write(story)
|
61 |
|
62 |
-
# Generate audio
|
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:
|