Update app.py

app.py

@@ -1,7 +1,12 @@
-import os
 import json
+import random
+import string
 import requests
+import torch
+
 import gradio as gr
+from PIL import Image
+from diffusers import StableDiffusionPipeline
 
 from pingpong import PingPong
 from pingpong.pingpong import PPManager
@@ -9,6 +14,64 @@ from pingpong.pingpong import PromptFmt
 from pingpong.pingpong import UIFmt
 from pingpong.gradio import GradioChatUIFmt
 
+from fpdf import FPDF
+
+class PDF(FPDF):
+    def header(self):
+        # Arial bold 15
+        self.set_font('Arial', 'B', 15)
+        # Calculate width of title and position
+        w = self.get_string_width(self.title) + 6
+        self.set_x((210 - w) / 2)
+        # Colors of frame, background and text
+        self.set_draw_color(255, 255, 255)
+        self.set_fill_color(255, 255, 255)
+        # self.set_text_color(220, 50, 50)
+        # Thickness of frame (1 mm)
+        self.set_line_width(1)
+        # Title
+        self.cell(w, 9, self.title, 1, 1, 'C', 1)
+        # Line break
+        self.ln(10)
+
+        if self.art is not None:
+          self.image(self.art, x=self.w/2.0-25, w=50)
+          self.ln(10)
+
+    def footer(self):
+        # Position at 1.5 cm from bottom
+        self.set_y(-15)
+        # Arial italic 8
+        self.set_font('Arial', 'I', 8)
+        # Text color in gray
+        self.set_text_color(128)
+        # Page number
+        self.cell(0, 10, 'Page ' + str(self.page_no()), 0, 0, 'C')
+
+    def chapter_title(self, num, label):
+        # Arial 12
+        self.set_font('Arial', '', 12)
+        # Background color
+        self.set_fill_color(200, 220, 255)
+        # Title
+        self.cell(0, 6, 'Chapter %d : %s' % (num, label), 0, 1, 'L', 1)
+        # Line break
+        self.ln(4)
+
+    def chapter_body(self, content):
+        # Times 12
+        self.set_font('Times', '', 12)
+        # Output justified text
+        self.multi_cell(0, 5, content)
+        # Line break
+        self.ln()
+        # Mention in italics
+        self.set_font('', 'I')
+
+    def print_chapter(self, content):
+        self.add_page()
+        self.chapter_body(content)
+
 class LLaMA2ChatPromptFmt(PromptFmt):
     @classmethod
     def ctx(cls, context):
@@ -54,7 +117,15 @@ class GradioLLaMA2ChatPPManager(LLaMA2ChatPPManager):
 TOKEN = os.getenv('HF_TOKEN')
 MODEL_ID = 'meta-llama/Llama-2-70b-chat-hf'
 
+pipe = StableDiffusionPipeline.from_pretrained(
+  "nota-ai/bk-sdm-small", torch_dtype=torch.float16
+)
+
 STYLES = """
+.left-panel {
+  min-width: min(290px, 100%) !important;
+}
+
 .small-big {
   font-size: 12pt !important;
 }
@@ -110,13 +181,17 @@ STYLES = """
 .no-label-chatbot {
   border: none !important;
   box-shadow: none !important;
-  height: 540px !important;
+  height: 520px !important;
 }
 
 .no-label-chatbot > div > div:nth-child(1) {
   display: none;
 }
 
+.no-label-image > div:nth-child(2) {
+  display: none;
+}
+
 .left-margin-30 {
   padding-left: 30px !important;
 }
@@ -138,6 +213,9 @@ STYLES = """
 }
 """
 
+def id_generator(size=6, chars=string.ascii_uppercase + string.digits):
+  return ''.join(random.choice(chars) for _ in range(size))
+
 def get_new_ppm(ping):
   ppm = LLaMA2ChatPPManager()
   ppm.ctx = """\
@@ -192,6 +270,44 @@ def gen_text(prompt, hf_model='meta-llama/Llama-2-70b-chat-hf', hf_token=None, p
 
   return json.loads(r.content.decode("utf-8"))[0]['generated_text']
 
+def gen_art(editor, cover_art_image, gen_cover_art_prompt):
+  if gen_cover_art_prompt.strip() == "":
+    ppm = get_new_ppm(f"""describe the story below as a movie poster. give me the caption ONLY.
+--------------------------------
+{editor}""")
+    cover_art_prompt = gen_text(ppm.build_prompts(), hf_model=MODEL_ID, hf_token=TOKEN)
+
+    return [
+        cover_art_image,
+        cover_art_prompt
+    ]
+  else:
+    global pipe
+    pipe = pipe.to("cuda")
+    
+    return [
+        pipe(gen_cover_art_prompt).images[0],
+        gen_cover_art_prompt
+    ]
+
+def generate_pdf(title, editor, concept_art):
+  tmp_filename = id_generator()
+
+  if concept_art is not None:
+    im = Image.fromarray(concept_art)
+    im.save(f"{tmp_filename}.png")
+
+  pdf = PDF()
+  pdf.title = "Untitled" if title.strip() == "" else title
+  pdf.art = None if concept_art is None else f"{tmp_filename}.png"
+  pdf.print_chapter(editor)
+  pdf.output(f'{tmp_filename}.pdf', 'F')
+
+  return (
+      gr.update(value=f'{tmp_filename}.pdf', visible=True),
+      "   "
+  )
+
 def select(editor, evt: gr.SelectData):
   return [
       evt.value,
@@ -199,9 +315,9 @@ def select(editor, evt: gr.SelectData):
       evt.index[1]
   ]
 
-def get_gen_txt(editor, prompt, only_gen_text=False):
+def get_gen_txt(title, editor, prompt, only_gen_text=False):
   if editor.strip() == '':
-    ppm = get_new_ppm('Write a short paragraph to start a short story for me')
+    ppm = get_new_ppm(f'Write a short paragraph to start a short story titled "{title}" for me')
   else:
     ppm = get_new_ppm(f"""{prompt}
 --------------------------------
@@ -217,22 +333,9 @@ def get_gen_txt(editor, prompt, only_gen_text=False):
     print(f"something went wrong - {e}")
     return editor
 
-def gen_txt(editor):
-  return [
-      get_gen_txt(editor, "Write the next paragraph based on the following stories so far."),
-      0,
-      gr.update(interactive=True),
-      gr.update(visible=False),
-      gr.update(visible=False),
-      gr.update(visible=False),
-      gr.update(interactive=True),
-      gr.update(interactive=True),
-      gr.update(interactive=True),
-  ]
-
-def gen_txt_with_prompt(editor, prompt):
+def gen_txt(title, editor, prompt):
   return [
-      get_gen_txt(editor, prompt),
+      get_gen_txt(title, editor, "Write the next paragraph based on the following stories so far." if prompt.strip() == "" else prompt),
       0,
       gr.update(interactive=True),
       gr.update(visible=False),
@@ -240,7 +343,6 @@ def gen_txt_with_prompt(editor, prompt):
       gr.update(visible=False),
       gr.update(interactive=True),
       gr.update(interactive=True),
-      gr.update(interactive=True),
   ]
 
 def chat_gen(editor, chat_txt, chatbot, ppm, regen=False):
@@ -310,9 +412,9 @@ def replace_sel(editor, replace_type, selected_text, sel_index_from, sel_index_t
       0
   ]
 
-def gen_alt(editor, num_enabled_alts, alt_btn1, alt_btn2, alt_btn3):
+def gen_alt(title, editor, num_enabled_alts, alt_btn1, alt_btn2, alt_btn3):
   if num_enabled_alts < 3:
-    gen_txt = get_gen_txt(editor, "Write the next paragraph based on the following stories so far.", only_gen_text=True)
+    gen_txt = get_gen_txt(title, editor, "Write the next paragraph based on the following stories so far.", only_gen_text=True)
 
   return [
       min(num_enabled_alts+1, 3),
@@ -326,7 +428,6 @@ def gen_alt(editor, num_enabled_alts, alt_btn1, alt_btn2, alt_btn3):
       " ",
       gr.update(interactive=True),
       gr.update(interactive=True),
-      gr.update(interactive=True),
   ]
 
 def fill_with_gen(alt_txt, editor):
@@ -355,16 +456,32 @@ with gr.Blocks(css=STYLES) as demo:
       "be changed over time, but [meta-llama/Llama-2-70b-chat-hf](https://huggingface.co/meta-llama/Llama-2-70b-chat-hf) is selected for now.",
       elem_classes=['center', 'small-big'])
 
+  progress_bar = gr.Textbox(elem_classes=['no-label'])
+
   with gr.Row():
     with gr.Column(scale=2):
-      editor = gr.Textbox(lines=30, max_lines=30, elem_classes=['no-label', 'small-big-textarea'])
+      editor = gr.Textbox(lines=32, max_lines=32, elem_classes=['no-label', 'small-big-textarea'])
       word_counter = gr.Markdown("0 words", elem_classes=['right'])
 
     with gr.Column(scale=1):
       with gr.Tab("Control"):
+        with gr.Column(elem_classes=['group-border']):
+          gr.Markdown('### title')
+          title = gr.Textbox("pokemon training story", elem_classes=['no-label'])
+
         with gr.Column(elem_classes=['group-border']):
           with gr.Column():
-            gr.Markdown("`generate text` button generate continued text and attach it to the end. on the other hand, `generate alternatives` button generate alternate texts up to 3 and let you choose one of them.")
+            gr.Markdown("For instant generation and concatenation, use `generate text` button. "
+                        "Want to explore alternative choices? use `generate alternatives` button.")
+            
+            with gr.Accordion("longer guideline", open=False):
+              gr.Markdown("`generate text` button generate continued text and attach it to the end. " 
+                        "on the other hand, `generate alternatives` button generate alternate texts "
+                        "up to 3 and let you choose one of them. In both cases, **Write the next paragraph based on "
+                        "the following stories so far.** is the default prompt. If you want to try your own designed "
+                        "prompt, enter it in the textbox below.")
+              
+            prompt = gr.Textbox(placeholder="design your own prompt", elem_classes=['no-label'])
 
           with gr.Row():
             gen_btn = gr.Button("generate text", elem_classes=['control-label-font', 'control-button'])
@@ -383,14 +500,6 @@ with gr.Blocks(css=STYLES) as demo:
               gr.Markdown("↳", scale=1, elem_classes=['wrap'])
               alt_btn3 = gr.Button("Alternative 3", elem_classes=['alt-button'], scale=8)
 
-        with gr.Row(elem_classes=['group-border']):
-          with gr.Column():
-            gr.Markdown("'Write the next paragraph based on the following stories so far.' is the default prompt when clicking `generate text`, and the text so far will always be attached to the end. By giving your own prompt, only the default prompt will be replaced.")
-
-          with gr.Column(elem_classes=['no-gap']):
-            gen_with_prompt_btn = gr.Button("generate text with custom prompt", elem_classes=['control-label-font', 'control-button'])
-            prompt = gr.Textbox(placeholder="enter prompt: ", elem_classes=['no-label'])
-
         with gr.Column(elem_classes=['group-border']):
           with gr.Row():
             selected_text = gr.Markdown("Selected text will be displayed in this area", elem_classes=['highlighted-text'])
@@ -400,11 +509,6 @@ with gr.Blocks(css=STYLES) as demo:
               replace_sel_btn = gr.Button("replace selection", elem_classes=['control-label-font', 'control-button'])
               replace_type = gr.Dropdown(choices=['word', 'sentense', 'phrase', 'paragraph'], value='sentense', interactive=True, elem_classes=['no-label'])
 
-          # with gr.Row():
-          #   with gr.Column(elem_classes=['no-gap']):
-          #     rewrite_sel_btn = gr.Button("rewrite selection", elem_classes=['control-label-font', 'control-button'])
-          #     rewrite_prompt = gr.Textbox(placeholder="Rewrite the text: ", elem_classes=['no-label'])
-
       with gr.Tab("Chatting"):
         chatbot = gr.Chatbot([], elem_classes=['no-label-chatbot'])
         chat_txt = gr.Textbox(placeholder="enter question", elem_classes=['no-label'])
@@ -413,8 +517,33 @@ with gr.Blocks(css=STYLES) as demo:
           clear_btn = gr.Button("clear", elem_classes=['control-label-font', 'control-button'])
           regen_btn = gr.Button("regenerate", elem_classes=['control-label-font', 'control-button'])
 
-  progress_bar = gr.Textbox(elem_classes=['no-label'])
-    
+      with gr.Tab("Exporting"):
+        cover_art = gr.Image(interactive=False, elem_classes=['no-label-image'])
+        gen_cover_art_prompt = gr.Textbox(lines=5, max_lines=5, elem_classes=['no-label'])
+        # toggle between "generate prompt for cover art" and "generate cover art"
+        gen_cover_art_btn = gr.Button("generate prompt for cover art", elem_classes=['control-label-font', 'control-button'])
+        gen_pdf_btn = gr.Button("export as PDF", elem_classes=['control-label-font', 'control-button'])
+        pdf_file = gr.File(visible=False)
+
+  gen_pdf_btn.click(
+      lambda t, e, c: generate_pdf(t, e, c),
+      inputs=[title, editor, cover_art],
+      outputs=[pdf_file, progress_bar]
+  )
+
+  gen_cover_art_btn.click(
+      gen_art,
+      inputs=[editor, cover_art, gen_cover_art_prompt],
+      outputs=[cover_art, gen_cover_art_prompt]
+  )
+
+  gen_cover_art_prompt.change(
+      fn=None,
+      inputs=[gen_cover_art_prompt],
+      outputs=[gen_cover_art_btn],
+      _js="(t) => t.trim() == '' ? 'generate prompt for cover art' : 'generate cover art'"
+  )
+
   editor.change(
       fn=None,
       inputs=[editor],
@@ -437,11 +566,11 @@ with gr.Blocks(css=STYLES) as demo:
           gr.update(interactive=False),
       ),
       inputs=None,
-      outputs=[gen_btn, gen_alt_btn, gen_with_prompt_btn, replace_sel_btn]
+      outputs=[gen_btn, gen_alt_btn, replace_sel_btn]
   ).then(
       fn=gen_txt,
-      inputs=[editor],
-      outputs=[editor, num_enabled_alts, gen_alt_btn, first_alt, second_alt, third_alt, gen_btn, gen_with_prompt_btn, replace_sel_btn]
+      inputs=[title, editor, prompt],
+      outputs=[editor, num_enabled_alts, gen_alt_btn, first_alt, second_alt, third_alt, gen_btn, replace_sel_btn]
   )
 
   gen_alt_btn.click(
@@ -452,11 +581,11 @@ with gr.Blocks(css=STYLES) as demo:
           gr.update(interactive=False),
       ),
       inputs=None,
-      outputs=[gen_btn, gen_alt_btn, gen_with_prompt_btn, replace_sel_btn]
+      outputs=[gen_btn, gen_alt_btn, replace_sel_btn]
   ).then(
       fn=gen_alt,
-      inputs=[editor, num_enabled_alts, alt_btn1, alt_btn2, alt_btn3],
-      outputs=[num_enabled_alts, gen_alt_btn, first_alt, alt_btn1, second_alt, alt_btn2, third_alt, alt_btn3, progress_bar, gen_btn, gen_with_prompt_btn, replace_sel_btn],
+      inputs=[title, editor, num_enabled_alts, alt_btn1, alt_btn2, alt_btn3],
+      outputs=[num_enabled_alts, gen_alt_btn, first_alt, alt_btn1, second_alt, alt_btn2, third_alt, alt_btn3, progress_bar, gen_btn, replace_sel_btn],
   )
 
   alt_btn1.click(
@@ -475,12 +604,6 @@ with gr.Blocks(css=STYLES) as demo:
       outputs=[editor, num_enabled_alts, gen_alt_btn, first_alt, second_alt, third_alt]
   )
 
-  gen_with_prompt_btn.click(
-      gen_txt_with_prompt,
-      inputs=[editor, prompt],
-      outputs=[editor, num_enabled_alts, gen_alt_btn, first_alt, second_alt, third_alt]
-  )
-
   replace_sel_btn.click(
       fn=replace_sel,
       inputs=[editor, replace_type, selected_text, sel_index_from, sel_index_to],
@@ -499,5 +622,5 @@ with gr.Blocks(css=STYLES) as demo:
       inputs=[editor, chat_txt, chatbot, chat_history],
       outputs=[chat_txt, chatbot, chat_history]
   )
-
+    
 demo.launch()