|
import gradio as gr |
|
import time |
|
import re |
|
import os |
|
from ebooklib import epub |
|
import random |
|
import base64 |
|
import requests |
|
import json |
|
from huggingface_hub import InferenceClient |
|
client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1") |
|
|
|
|
|
|
|
|
|
ANTHROPIC_API_KEY = "YOUR KEY HERE" |
|
stability_api_key = "YOUR KEY HERE" |
|
|
|
def remove_first_line(test_string): |
|
print("removing first line") |
|
if test_string.startswith("Here") and test_string.split("\n")[0].strip().endswith(":"): |
|
return re.sub(r'^.*\n', '', test_string, count=1) |
|
return test_string |
|
def format_prompt(message, history): |
|
prompt = "<s>" |
|
for user_prompt, bot_response in history: |
|
prompt += f"[INST] {user_prompt} [/INST]" |
|
prompt += f" {bot_response}</s> " |
|
prompt += f"[INST] {message} [/INST]" |
|
return prompt |
|
|
|
def generate_text(prompt,max_tokens=2000,history=""): |
|
generate_kwargs = dict( |
|
temperature=0.9, |
|
max_new_tokens=max_tokens, |
|
top_p=0.95, |
|
repetition_penalty=1.0, |
|
do_sample=True, |
|
seed=random.randint(1,1000000000) |
|
, |
|
) |
|
|
|
|
|
|
|
|
|
system_prompt="You are a world-class author. Write the requested content with great skill and attention to detail." |
|
|
|
formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history) |
|
|
|
|
|
stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False) |
|
resp = "" |
|
for response in stream: |
|
resp += response.token.text |
|
|
|
return resp |
|
|
|
def generate_text_og(prompt, model="claude-3-haiku-20240307", max_tokens=2000, temperature=0.7): |
|
headers = { |
|
"x-api-key": ANTHROPIC_API_KEY, |
|
"anthropic-version": "2023-06-01", |
|
"content-type": "application/json" |
|
} |
|
data = { |
|
"model": model, |
|
"max_tokens": max_tokens, |
|
"temperature": temperature, |
|
"system": "You are a world-class author. Write the requested content with great skill and attention to detail.", |
|
"messages": [{"role": "user", "content": prompt}], |
|
} |
|
response = requests.post("https://api.anthropic.com/v1/messages", headers=headers, json=data) |
|
response_text = response.json()['content'][0]['text'] |
|
return response_text.strip() |
|
|
|
def generate_cover_prompt(plot): |
|
response = generate_text(f"Plot: {plot}\n\n--\n\nDescribe the cover we should create, based on the plot. This should be two sentences long, maximum.") |
|
return response |
|
|
|
def generate_title(plot): |
|
response = generate_text(f"Here is the plot for the book: {plot}\n\n--\n\nRespond with a great title for this book. Only respond with the title, nothing else is allowed.") |
|
|
|
return response |
|
|
|
|
|
def create_cover_image(plot): |
|
|
|
plot = str(generate_cover_prompt(plot)) |
|
|
|
with open(f"cover.png", "wb") as f: |
|
f.write(base64.b64decode(image["base64"])) |
|
|
|
|
|
def create_cover_image_OG(plot): |
|
|
|
plot = str(generate_cover_prompt(plot)) |
|
|
|
engine_id = "stable-diffusion-xl-beta-v2-2-2" |
|
api_host = os.getenv('API_HOST', 'https://api.stability.ai') |
|
api_key = stability_api_key |
|
|
|
if api_key is None: |
|
raise Exception("Missing Stability API key.") |
|
|
|
response = requests.post( |
|
f"{api_host}/v1/generation/{engine_id}/text-to-image", |
|
headers={ |
|
"Content-Type": "application/json", |
|
"Accept": "application/json", |
|
"Authorization": f"Bearer {api_key}" |
|
}, |
|
json={ |
|
"text_prompts": [ |
|
{ |
|
"text": plot |
|
} |
|
], |
|
"cfg_scale": 7, |
|
"clip_guidance_preset": "FAST_BLUE", |
|
"height": 768, |
|
"width": 512, |
|
"samples": 1, |
|
"steps": 30, |
|
}, |
|
) |
|
|
|
if response.status_code != 200: |
|
raise Exception("Non-200 response: " + str(response.text)) |
|
|
|
data = response.json() |
|
|
|
for i, image in enumerate(data["artifacts"]): |
|
with open(f"/content/cover.png", "wb") as f: |
|
f.write(base64.b64decode(image["base64"])) |
|
|
|
def generate_chapter_title(chapter_content): |
|
response = generate_text(f"Chapter Content:\n\n{chapter_content}\n\n--\n\nGenerate a concise and engaging title for this chapter based on its content. Respond with the title only, nothing else.") |
|
return remove_first_line(response) |
|
|
|
def create_epub(title, author, chapters, cover_image_path='cover.png'): |
|
book = epub.EpubBook() |
|
|
|
book.set_identifier('id123456') |
|
book.set_title(title) |
|
book.set_language('en') |
|
book.add_author(author) |
|
|
|
with open(cover_image_path, 'rb') as cover_file: |
|
cover_image = cover_file.read() |
|
book.set_cover('cover.png', cover_image) |
|
|
|
epub_chapters = [] |
|
for i, chapter_content in enumerate(chapters): |
|
chapter_title = generate_chapter_title(chapter_content) |
|
chapter_file_name = f'chapter_{i+1}.xhtml' |
|
epub_chapter = epub.EpubHtml(title=chapter_title, file_name=chapter_file_name, lang='en') |
|
|
|
formatted_content = ''.join(f'{paragraph.strip()}' for paragraph in chapter_content.split('\n') if paragraph.strip()) |
|
epub_chapter.content = f'{chapter_title}{formatted_content}' |
|
book.add_item(epub_chapter) |
|
epub_chapters.append(epub_chapter) |
|
|
|
|
|
|
|
book.toc = (epub_chapters) |
|
|
|
|
|
book.add_item(epub.EpubNcx()) |
|
book.add_item(epub.EpubNav()) |
|
|
|
|
|
style = ''' |
|
@namespace epub "http://www.idpf.org/2007/ops"; |
|
body { |
|
font-family: Cambria, Liberation Serif, serif; |
|
} |
|
h1 { |
|
text-align: left; |
|
text-transform: uppercase; |
|
font-weight: 200; |
|
} |
|
''' |
|
|
|
|
|
nav_css = epub.EpubItem(uid="style_nav", file_name="style/nav.css", media_type="text/css", content=style) |
|
book.add_item(nav_css) |
|
|
|
|
|
book.spine = ['nav'] + epub_chapters |
|
|
|
|
|
epub.write_epub(f'{title}.epub', book) |
|
|
|
|
|
def generate_book(writing_style, book_description, num_chapters): |
|
print("Generating plot outline...") |
|
plot_prompt = f"Create a detailed plot outline for a {num_chapters}-chapter book in the {writing_style} style, based on the following description:\n\n{book_description}\n\nEach chapter should be at least 10 pages long." |
|
plot_outline = generate_text(plot_prompt) |
|
print("Plot outline generated.") |
|
|
|
chapters = [] |
|
for i in range(num_chapters): |
|
print(f"Generating chapter {i+1}...") |
|
chapter_prompt = f"Previous Chapters:\n\n{' '.join(chapters)}\n\nWriting style: `{writing_style}`\n\nPlot Outline:\n\n{plot_outline}\n\nWrite chapter {i+1} of the book, ensuring it follows the plot outline and builds upon the previous chapters. The chapter should be at least 256 paragraphs long... we're going for lengthy yet exciting chapters here." |
|
chapter = generate_text(chapter_prompt, max_tokens=4000) |
|
chapters.append(remove_first_line(chapter)) |
|
print(f"Chapter {i+1} generated.") |
|
time.sleep(1) |
|
|
|
print("Compiling the book...") |
|
book = "\n\n".join(chapters) |
|
print("Book generated!") |
|
|
|
return plot_outline, book, chapters |
|
|
|
|
|
def main(writing_style, book_description, num_chapters): |
|
try: |
|
|
|
plot_outline, book, chapters = generate_book(writing_style, book_description, num_chapters) |
|
|
|
title = generate_title(plot_outline) |
|
|
|
|
|
with open(f"{title}.txt", "w") as file: |
|
file.write(book) |
|
|
|
|
|
|
|
|
|
|
|
|
|
print(f"Book saved as '{title}.txt'.") |
|
return f'{title}.txt',[f'{title}.txt'] |
|
except Exception as e: |
|
return e,None |
|
|
|
with gr.Blocks() as app: |
|
with gr.Row(): |
|
with gr.Column(): |
|
writing_style=gr.Textbox(label="Enter the desired writing style: ") |
|
book_description = gr.Textbox(label="Enter a high-level description of the book: ") |
|
num_chapters = gr.Number(label="Enter the number of chapters: ",precision=0) |
|
with gr.Column(): |
|
gen_btn=gr.Button("Generate Book") |
|
|
|
outp=gr.HTML() |
|
outf=gr.Files() |
|
gen_btn.click(main,[writing_style,book_description,num_chapters],[outp,outf]) |
|
app.queue(default_concurrency_limit=10).launch() |