Update app.py

app.py

@@ -1,10 +1,17 @@
 import gradio as gr
 from urllib.parse import quote
+from reportlab.lib.pagesizes import letter
+from reportlab.platypus import SimpleDocTemplate, Paragraph, Image
+from reportlab.lib.styles import getSampleStyleSheet
+import urllib.request
 import re
 import os
 import openai
 from gradio_client import Client
 
+def download_image(image_url, image_path):
+    urllib.request.urlretrieve(image_url, image_path)
+
 OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
 clipi_client = Client("https://fffiloni-clip-interrogator-2.hf.space/")
 
@@ -18,6 +25,21 @@ def generate_image_url(keywords):
     truncated_keywords = keywords[:20]
     return f"https://image.pollinations.ai/prompt/{quote(truncated_keywords)}"
 
+def save_as_pdf():
+    global chapters
+    pdf_path = "/mnt/data/chapters.pdf"
+    pdf = SimpleDocTemplate(pdf_path, pagesize=letter)
+    story = []
+    styles = getSampleStyleSheet()
+    for i, chapter in enumerate(chapters):
+        story.append(Paragraph(f"Chapter {i + 1}", styles["Heading1"]))
+        story.append(Paragraph(chapter["story"], styles["BodyText"]))
+        image_path = f"/mnt/data/image_{i}.png"
+        download_image(chapter["image"], image_path)
+        story.append(Image(image_path))
+    pdf.build(story)
+    return pdf_path
+
 def next_chapter(audience, keyword, protagonist):
     global chapter_num, current_story, current_image_url
     current_image_url = generate_image_url(current_story)
@@ -25,18 +47,16 @@ def next_chapter(audience, keyword, protagonist):
     chapter_prompt = f"{story_intro}\n\nKeyword: {keyword}\nProtagonist: {protagonist}\n\n![Chapter {chapter_num} Image]({current_image_url})\n\nChapter {chapter_num} 내용을 만들어줘."
     chat_completion = openai.ChatCompletion.create(model="gpt-3.5-turbo-16k", messages=[{"role": "user", "content": chapter_prompt}])
     current_story = chat_completion.choices[0].message.content
-    chapters.append({"story": current_story, "image": current_image_url})
     chapter_num += 1
-    chapter_table = "| 챕터 | 이야기 | 그림 |\n|------|--------|------|\n" + "\n".join([f"| {i+1} | {c['story']} | ![image]({c['image']}) |" for i, c in enumerate(chapters)])
-    return current_story, current_image_url, chapter_table
+    chapters.append({"story": current_story, "image": current_image_url})
+    return current_story, current_image_url
 
 def infer(image_input, audience, keyword, protagonist):
     global story_intro, current_story, current_image_url, chapter_num, chapters
-    chapter_num = 1
     chapters = []
+    chapter_num = 1
     gr.Info('Calling CLIP Interrogator, 이미지를 해석하고 있습니다...')
     clipi_result = clipi_client.predict(image_input, "best", 4, api_name="/clipi2")[0]
-
     story_intro = f"""
     # Illustrated Tales
     ## Created by [Sigkawat Pengnoo](https://flowgpt.com/prompt/qzv2D3OvHkzkfSE4rQCqv) at FlowGPT
@@ -46,7 +66,6 @@ def infer(image_input, audience, keyword, protagonist):
     STORY : "{{ {clipi_result} }}"
     Let's begin with Chapter 1!
     """
-
     current_story = clipi_result
     return next_chapter(audience, keyword, protagonist)
 
@@ -71,14 +90,13 @@ with gr.Blocks(css=css) as demo:
                 protagonist_in = gr.Textbox(label="주인공")
                 submit_btn = gr.Button('이야기와 그림을 만들어 주세요')
                 next_chapter_btn = gr.Button('다음 이야기')
+                save_pdf_btn = gr.Button('PDF로 저장하기')
             with gr.Column():
                 chapter_story = gr.Markdown(label="이야기", elem_id="chapter_story")
                 chapter_image = gr.Image(label="그림", elem_id="chapter_image")
-                chapter_table_md = gr.Markdown(label="모든 챕터")
 
-    submit_btn.click(fn=infer, inputs=[image_in, audience, keyword_in, protagonist_in], 
-                     outputs=[chapter_story, chapter_image, chapter_table_md])
-    next_chapter_btn.click(fn=lambda: next_chapter(audience=audience.value, keyword=keyword_in.value, protagonist=protagonist_in.value), 
-                           outputs=[chapter_story, chapter_image, chapter_table_md])
+    submit_btn.click(fn=infer, inputs=[image_in, audience, keyword_in, protagonist_in], outputs=[chapter_story, chapter_image])
+    next_chapter_btn.click(fn=lambda: next_chapter(audience=audience.value, keyword=keyword_in.value, protagonist=protagonist_in.value), outputs=[chapter_story, chapter_image])
+    save_pdf_btn.click(fn=save_as_pdf, outputs=gr.File(label="PDF 다운로드"))
 
     demo.queue(max_size=12).launch()