storybook-generator / generate_plot.py
ethanrom's picture
Upload 3 files
9ffa7ce
import os
import requests
import io
from PIL import Image
from langchain import PromptTemplate, LLMChain
from PIL import Image, ImageDraw, ImageFont, ImageFilter
from langchain.llms import OpenAI
import openai
from g4f import Provider, Model
from langchain_g4f import G4FLLM
def set_openai_api_key(api_key):
openai.api_key = api_key
os.environ["OPENAI_API_KEY"] = openai.api_key
template = """Write a very short and simple children storybook suitable for children of {age} years and below.
Your answer should be structred like this with <Text> and <Image> tags.
<Text> first page text here </Text>
<Image> descibe image here without including names so that prompt can be used to generate an image using a ML model.</Image>
NUMBER OF PAGES THE BOOK MUST HAVE:
{number_of_pages}
=============
CHARACTERS BOOK MUST HAVE:
{character_names}
Answer:"""
prompt = PromptTemplate(template=template, input_variables=["number_of_pages", "character_names", "age"])
def query(payload):
API_URL = "https://api-inference.huggingface.co/models/prompthero/openjourney"
headers = {"Authorization": "Bearer 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"}
response = requests.post(API_URL, headers=headers, json=payload)
return response.content
def query_alt(payload):
API_URL = "https://api-inference.huggingface.co/models/stablediffusionapi/anything-v5"
headers = {"Authorization": "Bearer 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"}
response = requests.post(API_URL, headers=headers, json=payload)
return response.content
def generate_plot(number_of_pages, character_names, age, selected_style, provider, selected_provider=None):
if provider == "OpenAI":
llm = OpenAI(temperature=0)
elif provider == "G4F":
if selected_provider == "Ails":
llm = G4FLLM(
model=Model.gpt_35_turbo,
provider=Provider.Ails,
)
elif selected_provider == "You":
llm = G4FLLM(
model=Model.gpt_35_turbo,
provider=Provider.You,
)
elif selected_provider == "GetGpt":
llm = G4FLLM(
model=Model.gpt_35_turbo,
provider=Provider.GetGpt,
)
elif selected_provider == "DeepAi":
llm = G4FLLM(
model=Model.gpt_35_turbo,
provider=Provider.DeepAi,
)
elif selected_provider == "Forefront":
llm = G4FLLM(
model=Model.gpt_35_turbo,
provider=Provider.Forefront,
)
elif selected_provider == "Aichat":
llm = G4FLLM(
model=Model.gpt_35_turbo,
provider=Provider.Aichat,
)
elif selected_provider == "Bard":
llm = G4FLLM(
model=Model.gpt_35_turbo,
provider=Provider.Bard,
)
# Add other providers here
else:
raise ValueError("Invalid G4F provider selected.")
else:
raise ValueError("Invalid provider selected.")
llm_chain = LLMChain(prompt=prompt, llm=llm)
response = llm_chain.run(number_of_pages=number_of_pages, character_names=character_names, age=age)
pages = response.split("<Text>")
plot_result = []
additional_texts = {
"Style 1" : " clean cel shaded vector art. shutterstock. behance hd by lois van baarle, artgerm, helen huang, by makoto shinkai and ilya kuvshinov, rossdraws, illustration, art by ilya kuvshinov and gustav klimt 4.",
"Style 2" : " storybook illustration, monochromatic, white background, pinterest, trending on artstation behance, digital illustration, vector illustration, sharp focus, pastel palette, graphic novel, 8k",
"Style 3" : " storybook illustration, monochromatic, white background, vector art illustration, award winning, stunning, trending on artstation, detailed, high quality resolution, cartoon character design, adobe illustrator."
}
for i, page in enumerate(pages[1:]):
text, img_text = page.split("<Image>", 1)
selected_additional_text = additional_texts.get(selected_style, "")
new_img_text = img_text.strip() + " " + selected_additional_text
text = text.replace("</Text>", "")
new_img_text = new_img_text.replace("</Image>", "")
plot_result.append((i + 1, "<Text>" + text.strip() + "</Text>", "<Image>" + new_img_text + "</Image>"))
return plot_result
def generate_storybook(plot_result):
storybook = []
for page_number, text, image_text in plot_result:
image_bytes = query({"inputs": image_text})
image = Image.open(io.BytesIO(image_bytes))
blurred_image = image.filter(ImageFilter.GaussianBlur(8))
draw = ImageDraw.Draw(blurred_image)
font_size = 30
font = ImageFont.truetype("DeliusSwashCaps-Regular.ttf", font_size)
first_letter_font_size = 60
first_letter_font = ImageFont.truetype("DeliusSwashCaps-Regular.ttf", first_letter_font_size)
text_x, text_y = 50, 50
max_width = blurred_image.width - text_x * 2
text = text.replace("<Text>", "").replace("</Text>", "")
wrapped_text = ""
words = text.split()
for i, word in enumerate(words):
if i == 0:
first_letter_width = draw.textsize(word[0], font=first_letter_font)[0]
draw.text((text_x, text_y), word[0], fill="white", font=first_letter_font, stroke_width=2, stroke_fill="black")
text_x += first_letter_width + 10
wrapped_text += word[1:] + " "
elif draw.textsize(wrapped_text + word, font=font)[0] < max_width:
wrapped_text += word + " "
else:
draw.text((text_x, text_y), wrapped_text.strip(), fill="white", font=font, stroke_width=2, stroke_fill="black")
text_y += font.getsize(wrapped_text)[1] + 10
wrapped_text = word + " "
draw.text((text_x, text_y), wrapped_text.strip(), fill="white", font=font, stroke_width=2, stroke_fill="black")
combined_image = Image.new('RGB', (image.width * 2, image.height))
combined_image.paste(blurred_image, (0, 0))
combined_image.paste(image, (image.width, 0))
storybook.append((page_number, combined_image))
return storybook
def generate_book_cover(title, author, image_text):
book_cover = []
image_bytes = query({"inputs": image_text})
image = Image.open(io.BytesIO(image_bytes))
cover_image = Image.new('RGB', (image.width * 2, image.height), color='white')
cover_image.paste(image, (0, 0))
draw = ImageDraw.Draw(cover_image)
font_size = 50
title_font = ImageFont.truetype("DeliusSwashCaps-Regular.ttf", font_size)
author_font = ImageFont.truetype("DeliusSwashCaps-Regular.ttf", 30)
title_x, title_y = image.width + 50, 50
author_x, author_y = image.width + 50, title_y + 100
draw.text((title_x, title_y), title, fill="black", font=title_font)
draw.text((author_x, author_y), "By " + author, fill="black", font=author_font)
book_cover.append(cover_image)
return book_cover