Update app.py

app.py

@@ -4,6 +4,7 @@ import re
 import os
 import openai
 from gradio_client import Client
+from fpdf import FPDF
 
 OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
 clipi_client = Client("https://fffiloni-clip-interrogator-2.hf.space/")
@@ -12,22 +13,40 @@ chapter_num = 1
 story_intro = ""
 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():
+    pdf = FPDF()
+    pdf.set_auto_page_break(auto=True, margin=15)
+    pdf.add_page()
+
+    for chapter in chapters:
+        pdf.set_font("Arial", size=12)
+        pdf.multi_cell(0, 10, chapter['story'])
+        pdf.image(chapter['image'], x=10, y=None, w=90)
+    
+    pdf.output("illustrated_tales.pdf")
+    return "PDF 저장 완료!"
+
 def next_chapter(audience, keyword, protagonist):
     global chapter_num, current_story, current_image_url
-    current_image_url = generate_image_url(current_story)  # 다음 장의 이미지 URL을 생성합니다.
+    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
+    chapters.append({"story": current_story, "image": current_image_url}) # 챕터 정보 저장
     chapter_num += 1
-    return current_story, current_image_url
+    chapter_table = [[c['story'], c['image']] for c in chapters]
+    return current_story, current_image_url, chapter_table
 
 def infer(image_input, audience, keyword, protagonist):
     global story_intro, current_story, current_image_url, chapter_num
+    chapters.clear()  # 새로운 이야기 시작 시 초기화
     chapter_num = 1
     gr.Info('Calling CLIP Interrogator, 이미지를 해석하고 있습니다...')
     clipi_result = clipi_client.predict(image_input, "best", 4, api_name="/clipi2")[0]
@@ -66,15 +85,16 @@ 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 = gr.Table(label="모든 챕터")
 
     submit_btn.click(fn=infer, inputs=[image_in, audience, keyword_in, protagonist_in], 
-                     outputs=[chapter_story, chapter_image])
+                     outputs=[chapter_story, chapter_image, chapter_table])
     next_chapter_btn.click(fn=lambda: next_chapter(audience=audience.value, keyword=keyword_in.value, protagonist=protagonist_in.value), 
-                       outputs=[chapter_story, chapter_image])
-    # next_chapter_btn.click(fn=next_chapter, 
-    #                        outputs=[chapter_story, chapter_image])
+                           outputs=[chapter_story, chapter_image, chapter_table])
+    save_pdf_btn.click(fn=save_as_pdf, outputs="text")
 
     demo.queue(max_size=12).launch()