Update app.py

app.py

@@ -1,16 +1,31 @@
 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 openai
+from fpdf import FPDF
 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)
+# FPDF 클래스 확장
+class PDF(FPDF):
+    def chapter_title(self, num, label):
+        self.set_font('Arial', 'B', 12)
+        self.cell(0, 10, 'Chapter %d : %s' % (num, label), 0, 1, 'C')
+        self.ln(10)
+
+    def chapter_body(self, body, image_path):
+        self.set_font('Arial', '', 12)
+        self.multi_cell(0, 10, body)
+        self.image(image_path, x=10, w=190)
+        self.ln()
+
+    def add_chapter(self, num, title, body, image_url):
+        image_path = f"/mnt/data/image_{num}.png"
+        urllib.request.urlretrieve(image_url, image_path)
+        self.add_page()
+        self.chapter_title(num, title)
+        self.chapter_body(body, image_path)
+        os.remove(image_path)  # 이미지 파일 삭제
 
 OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
 clipi_client = Client("https://fffiloni-clip-interrogator-2.hf.space/")
@@ -27,17 +42,13 @@ def generate_image_url(keywords):
 
 def save_as_pdf():
     global chapters
+    pdf = PDF()
+    pdf.set_auto_page_break(auto=True, margin=15)
+    pdf.set_title("Chapters")
+    for i, chapter in enumerate(chapters, 1):
+        pdf.add_chapter(i, f"Chapter {i}", chapter["story"], chapter["image"])
     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)
+    pdf.output(pdf_path)
     return pdf_path
 
 def next_chapter(audience, keyword, protagonist):
@@ -97,6 +108,6 @@ with gr.Blocks(css=css) as demo:
 
     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 다운로드"))
+    save_pdf_btn.click(fn=save_as_pdf, outputs=gr.Download(label="PDF 다운로드"))
 
     demo.queue(max_size=12).launch()