Spaces:
Runtime error
Runtime error
| import streamlit as st | |
| import torch | |
| from transformers import pipeline, set_seed | |
| from transformers import AutoTokenizer | |
| from PIL import ( | |
| Image, | |
| ImageFont, | |
| ImageDraw | |
| ) | |
| import numpy as np | |
| from datetime import datetime | |
| import glob | |
| import re | |
| import random | |
| import textwrap | |
| from examples import EXAMPLES | |
| import meta | |
| from utils import ( | |
| remote_css, | |
| local_css, | |
| load_image, | |
| pure_comma_separation | |
| ) | |
| # np.random.seed(42) | |
| class TextGeneration: | |
| def __init__(self): | |
| self.debug = False | |
| self.dummy_output = { | |
| 'directions': [ | |
| "for the dough", | |
| "in a small bowl, combine the warm water and yeast.", | |
| "let it sit for 5 minutes.", | |
| "add the flour, salt, and olive oil.", | |
| "mix well and knead until the dough is smooth and elastic.", | |
| "cover the dough with a damp towel and let it rise for about 1 hour.", | |
| "for the filling", | |
| "heat a large skillet over medium high heat.", | |
| "cook the beef and onion until the beef is browned and the onion is translucent. browned and the onion is translucent.", | |
| "drain off any excess grease.", | |
| "stir in the pepper and salt and black pepper to taste.", | |
| "remove from the heat and set aside.", | |
| "preheat the oven to 425 degrees f.", | |
| "roll out the dough on a lightly floured surface into a 12 inch circle.", | |
| "spread the beef mixture over the dough, leaving a 1 inch border.", | |
| "top with the feta, parsley, and lemon juice.", | |
| "bake for 20 minutes or until the crust is golden brown.", | |
| "cut into wedges and serve.", | |
| ], | |
| 'ingredients': [ | |
| '1 potato', | |
| '1 slice cheese 1 slice cheese', | |
| '1 potato 1 slice cheese 1 slice cheese', | |
| '1 slice cheese' | |
| '1 potato', | |
| '1 slice cheese', | |
| '1 slice cheese', | |
| '1 potato', | |
| '1 slice cheese' | |
| '1 potato', | |
| '1 slice cheese', | |
| ], | |
| 'title': 'Cheese Potatoes with Some other items' | |
| } | |
| self.tokenizer = None | |
| self.generator = None | |
| self.task = "text2text-generation" | |
| self.model_name_or_path = "flax-community/t5-recipe-generation" | |
| self.frames = list(glob.glob("asset/images/frames/*.jpg")) | |
| self.fonts = { | |
| "title": ImageFont.truetype("asset/fonts/Poppins-Bold.ttf", 90), | |
| "sub_title": ImageFont.truetype("asset/fonts/Poppins-Medium.ttf", 30), | |
| "body_bold": ImageFont.truetype("asset/fonts/Montserrat-Bold.ttf", 50), | |
| "body": ImageFont.truetype("asset/fonts/Montserrat-Regular.ttf", 30), | |
| } | |
| self.list_division = 5 | |
| self.point = "•" | |
| set_seed(42) | |
| def _skip_special_tokens_and_prettify(self, text): | |
| recipe_maps = {"<sep>": "--", "<section>": "\n"} | |
| recipe_map_pattern = "|".join(map(re.escape, recipe_maps.keys())) | |
| text = re.sub( | |
| recipe_map_pattern, | |
| lambda m: recipe_maps[m.group()], | |
| re.sub("|".join(self.tokenizer.all_special_tokens), "", text) | |
| ) | |
| data = {"title": "", "ingredients": [], "directions": []} | |
| for section in text.split("\n"): | |
| section = section.strip() | |
| if section.startswith("title:"): | |
| data["title"] = " ".join( | |
| [w.strip().capitalize() for w in section.replace("title:", "").strip().split() if w.strip()] | |
| ) | |
| elif section.startswith("ingredients:"): | |
| data["ingredients"] = [s.strip() for s in section.replace("ingredients:", "").split('--')] | |
| elif section.startswith("directions:"): | |
| data["directions"] = [s.strip() for s in section.replace("directions:", "").split('--')] | |
| else: | |
| pass | |
| return data | |
| def load(self): | |
| if not self.debug: | |
| self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path) | |
| self.generator = pipeline(self.task, model=self.model_name_or_path, tokenizer=self.model_name_or_path) | |
| def prepare_frame(self, recipe, frame): | |
| im_editable = ImageDraw.Draw(frame) | |
| # Info | |
| im_editable.text( | |
| (1570, 180), | |
| "Created " + datetime.now().strftime("%Y-%m-%d %H:%M"), | |
| (61, 61, 70), | |
| font=self.fonts["sub_title"], | |
| ) | |
| im_editable.text( | |
| (1570, 170 + 70), | |
| "By " + recipe["by"], | |
| (61, 61, 70), | |
| font=self.fonts["sub_title"], | |
| ) | |
| # Title | |
| im_editable.text( | |
| (400, 650), | |
| textwrap.fill(recipe["title"], 30), | |
| (61, 61, 70), | |
| font=self.fonts["title"], | |
| ) | |
| # Ingredients | |
| im_editable.text( | |
| (100, 1000), | |
| "Ingredients", | |
| (61, 61, 70), | |
| font=self.fonts["body_bold"], | |
| ) | |
| ingredients = recipe["ingredients"] | |
| ingredients = [textwrap.fill(item, 30).replace("\n", "\n ") for item in ingredients] | |
| im_editable.text( | |
| (100, 1080), | |
| "\n".join([f"{self.point} {item}" for item in ingredients]), | |
| (61, 61, 70), | |
| font=self.fonts["body"], | |
| ) | |
| # Directions | |
| im_editable.text( | |
| (700, 1000), | |
| "Directions", | |
| (61, 61, 70), | |
| font=self.fonts["body_bold"], | |
| ) | |
| directions = recipe["directions"] | |
| directions = [textwrap.fill(item, 80).replace("\n", "\n ") for item in directions] | |
| im_editable.text( | |
| (700, 1080), | |
| "\n".join([f"{i + 1}. {item}" for i, item in enumerate(directions)]).strip(), | |
| (61, 61, 70), | |
| font=self.fonts["body"], | |
| ) | |
| # directions_col1 = [textwrap.fill(item, 30).replace("\n", "\n ") for item in directions[:self.list_division]] | |
| # directions_col2 = [textwrap.fill(item, 30).replace("\n", "\n ") for item in directions[self.list_division:]] | |
| # im_editable.text( | |
| # (700, 1130), | |
| # "\n".join([f"{i + 1}. {item}" for i, item in enumerate(directions_col1)]).strip(), | |
| # (61, 61, 70), | |
| # font=self.fonts["body"], | |
| # ) | |
| # im_editable.text( | |
| # (1300, 1130), | |
| # "\n".join([f"{i + 1 + self.list_division}. {item}" for i, item in enumerate(directions_col2)]).strip(), | |
| # (61, 61, 70), | |
| # font=self.fonts["body"], | |
| # ) | |
| return frame | |
| def generate(self, items, generation_kwargs): | |
| if not self.debug: | |
| generation_kwargs["num_return_sequences"] = 1 | |
| # generation_kwargs["return_full_text"] = False | |
| generation_kwargs["return_tensors"] = True | |
| generation_kwargs["return_text"] = False | |
| generated_ids = self.generator( | |
| items, | |
| **generation_kwargs, | |
| )[0]["generated_token_ids"] | |
| recipe = self.tokenizer.decode(generated_ids, skip_special_tokens=False) | |
| recipe = self._skip_special_tokens_and_prettify(recipe) | |
| return recipe | |
| return self.dummy_output | |
| def generate_frame(self, recipe, frame_path): | |
| frame = load_image(frame_path) | |
| return self.prepare_frame(recipe, frame) | |
| def load_text_generator(): | |
| generator = TextGeneration() | |
| generator.load() | |
| return generator | |
| chef_top = { | |
| "max_length": 512, | |
| "min_length": 64, | |
| "no_repeat_ngram_size": 3, | |
| "do_sample": True, | |
| "top_k": 60, | |
| "top_p": 0.95, | |
| "num_return_sequences": 1 | |
| } | |
| chef_beam = { | |
| "max_length": 512, | |
| "min_length": 64, | |
| "no_repeat_ngram_size": 3, | |
| "early_stopping": True, | |
| "num_beams": 5, | |
| "length_penalty": 1.5, | |
| "num_return_sequences": 1 | |
| } | |
| def main(): | |
| st.set_page_config( | |
| page_title="Chef Transformer", | |
| page_icon="🍲", | |
| layout="wide", | |
| initial_sidebar_state="expanded" | |
| ) | |
| generator = load_text_generator() | |
| if hasattr(st, "session_state"): | |
| if 'get_random_frame' not in st.session_state: | |
| st.session_state.get_random_frame = generator.frames[0] | |
| else: | |
| get_random_frame = generator.frames[0] | |
| local_css("asset/css/style.css") | |
| col1, col2 = st.beta_columns([5, 3]) | |
| with col2: | |
| st.image(load_image("asset/images/chef-transformer-transparent.png"), width=300) | |
| st.markdown(meta.SIDEBAR_INFO, unsafe_allow_html=True) | |
| with st.beta_expander("Where did this story start?"): | |
| st.markdown(meta.STORY, unsafe_allow_html=True) | |
| with col1: | |
| st.markdown(meta.HEADER_INFO, unsafe_allow_html=True) | |
| st.markdown(meta.CHEF_INFO, unsafe_allow_html=True) | |
| chef = st.selectbox("Choose your chef", index=0, options=["Chef Scheherazade", "Chef Giovanni"]) | |
| prompts = list(EXAMPLES.keys()) + ["Custom"] | |
| prompt = st.selectbox( | |
| 'Examples (select from this list)', | |
| prompts, | |
| # index=len(prompts) - 1, | |
| index=0 | |
| ) | |
| if prompt == "Custom": | |
| prompt_box = "" | |
| else: | |
| prompt_box = EXAMPLES[prompt] | |
| items = st.text_area( | |
| 'Insert your ingredients here (separated by `,`): ', | |
| pure_comma_separation(prompt_box, return_list=False), | |
| ) | |
| items = pure_comma_separation(items, return_list=False) | |
| entered_items = st.empty() | |
| recipe_button = st.button('Get Recipe!') | |
| st.markdown( | |
| "<hr />", | |
| unsafe_allow_html=True | |
| ) | |
| if recipe_button: | |
| if hasattr(st, "session_state"): | |
| st.session_state.get_random_frame = generator.frames[random.randint(0, len(generator.frames)) - 1] | |
| else: | |
| get_random_frame = generator.frames[random.randint(0, len(generator.frames)) - 1] | |
| entered_items.markdown("**Generate recipe for:** " + items) | |
| with st.spinner("Generating recipe..."): | |
| if not isinstance(items, str) or not len(items) > 1: | |
| entered_items.markdown( | |
| f"**{chef}** would like to know what ingredients do you like to use in " | |
| f"your food? " | |
| ) | |
| else: | |
| gen_kw = chef_top if chef == "Chef Scheherazade" else chef_beam | |
| generated_recipe = generator.generate(items, gen_kw) | |
| title = generated_recipe["title"] | |
| ingredients = generated_recipe["ingredients"] | |
| directions = [textwrap.fill(item, 70).replace("\n", "\n ") for item in | |
| generated_recipe["directions"]] | |
| generated_recipe["by"] = chef | |
| r1, r2 = st.beta_columns([3, 5]) | |
| with r1: | |
| # st.write(st.session_state.get_random_frame) | |
| if hasattr(st, "session_state"): | |
| recipe_post = generator.generate_frame(generated_recipe, st.session_state.get_random_frame) | |
| else: | |
| recipe_post = generator.generate_frame(generated_recipe, get_random_frame) | |
| st.image( | |
| recipe_post, | |
| # width=500, | |
| caption="Click 🔎 to enlarge", | |
| use_column_width="auto", | |
| output_format="PNG" | |
| ) | |
| with r2: | |
| st.markdown( | |
| " ".join([ | |
| "<div class='r-text-recipe'>", | |
| f"<h2>{title}</h2>", | |
| "<h3>Ingredients</h3>", | |
| "<ul class='ingredients-list'>", | |
| " ".join([f'<li>{item}</li>' for item in ingredients]), | |
| "</ul>", | |
| "<h3>Directions</h3>", | |
| "<ol class='ingredients-list'>", | |
| " ".join([f'<li>{item}</li>' for item in directions]), | |
| "</ol>", | |
| "</div>" | |
| ]), | |
| unsafe_allow_html=True | |
| ) | |
| if __name__ == '__main__': | |
| main() | |