Update app.py

app.py

@@ -1,31 +1,40 @@
+from fpdf import FPDF
 import gradio as gr
 from urllib.parse import quote
-import openai
-from fpdf import FPDF
-import urllib.request
+import re
 import os
+import openai
 from gradio_client import Client
 
-# FPDF 클래스 확장
+def generate_image_url(keywords):
+    truncated_keywords = keywords[:20]
+    return f"https://image.pollinations.ai/prompt/{quote(truncated_keywords)}"
+
 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.cell(0, 10, f"Chapter {num} : {label}", 0, 1, 'C')
         self.ln(10)
 
-    def chapter_body(self, body, image_path):
+    def chapter_body(self, body, image_url):
         self.set_font('Arial', '', 12)
         self.multi_cell(0, 10, body)
-        self.image(image_path, x=10, w=190)
-        self.ln()
+        self.image(image_url, x=10, w=100)
+        self.ln(60)
 
     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)  # 이미지 파일 삭제
+        self.chapter_body(body, image_url)
+
+def save_as_pdf():
+    global chapters
+    pdf_path = "/mnt/data/chapters.pdf"
+    pdf = PDF()
+    for i, chapter in enumerate(chapters, start=1):
+        pdf.add_chapter(i, f"Chapter {i}", chapter["story"], chapter["image"])
+    pdf.output(pdf_path)
+    return pdf_path
 
 OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
 clipi_client = Client("https://fffiloni-clip-interrogator-2.hf.space/")
@@ -36,36 +45,21 @@ current_story = ""
 current_image_url = ""
 chapters = []
 
-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 = 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.output(pdf_path)
-    return pdf_path
-
 def next_chapter(audience, keyword, protagonist):
-    global chapter_num, current_story, current_image_url
+    global chapter_num, current_story, current_image_url, chapters
     current_image_url = generate_image_url(current_story)
     gr.Info(f'Chapter {chapter_num}를 생성하고 있습니다...')
     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
-    chapter_num += 1
     chapters.append({"story": current_story, "image": current_image_url})
+    chapter_num += 1
     return current_story, current_image_url
 
 def infer(image_input, audience, keyword, protagonist):
     global story_intro, current_story, current_image_url, chapter_num, chapters
-    chapters = []
     chapter_num = 1
+    chapters = []
     gr.Info('Calling CLIP Interrogator, 이미지를 해석하고 있습니다...')
     clipi_result = clipi_client.predict(image_input, "best", 4, api_name="/clipi2")[0]
     story_intro = f"""
@@ -80,11 +74,13 @@ def infer(image_input, audience, keyword, protagonist):
     current_story = clipi_result
     return next_chapter(audience, keyword, protagonist)
 
+# CSS for layout
 css = """
 #col-container {max-width: 910px; margin-left: auto; margin-right: auto;}
 a {text-decoration-line: underline; font-weight: 600;}
 """
 
+# Gradio UI
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown(