Update app.py

app.py

@@ -1825,8 +1825,42 @@ def stitch_story(raw_story: List[dict], theme_id: str, level: str) -> Dict[str,
         logging.error(f"Error in story stitching: {str(e)}")
         raise
 
+def generate_story_illustration(story_summary: str) -> str:
+    """Generate illustration for the story using DALL-E 3"""
+    try:
+        # Prepare the prompt for DALL-E
+        illustration_prompt = f"""
+        Create a children's book illustration for this story:
+        {story_summary}
+        
+        Style requirements:
+        - Cute and child-friendly illustration
+        - Soft, warm colors
+        - Storybook art style
+        - Clear and simple composition
+        - Safe for children
+        - No text in the image
+        """
+        
+        # Generate image
+        response = client.images.generate(
+            model="dall-e-3",
+            prompt=illustration_prompt,
+            size="1024x1024",
+            quality="standard",
+            n=1,
+        )
+        
+        image_url = response.data[0].url
+        logging.info("Story illustration generated successfully")
+        return image_url
+        
+    except Exception as e:
+        logging.error(f"Error generating story illustration: {str(e)}")
+        return None
+
 def show_stitched_story(story_data: Dict[str, str]):
-    """Display the stitched story with translation"""
+    """Display the stitched story with translation and illustration"""
     try:
         st.markdown("# 📖 Your Polished Story | เรื่องราวฉบับสมบูรณ์")
         
@@ -1842,6 +1876,27 @@ def show_stitched_story(story_data: Dict[str, str]):
             unsafe_allow_html=True
         )
         
+        # Generate and display illustration
+        with st.spinner("🎨 กำลังวาดภาพประกอบ..."):
+            illustration_url = generate_story_illustration(story_data['polished_english'])
+            if illustration_url:
+                st.markdown("### 🎨 Story Illustration | ภาพประกอบเรื่องราว")
+                st.image(
+                    illustration_url,
+                    caption="AI-generated illustration for your story",
+                    use_column_width=True
+                )
+                
+                # Add download button for illustration
+                st.markdown(f"""
+                    <div style="text-align: center; margin: 10px 0;">
+                        <a href="{illustration_url}" download="story_illustration.png" 
+                           target="_blank" style="text-decoration: none;">
+                            📥 ดาวน์โหลดภาพประกอบ
+                        </a>
+                    </div>
+                """, unsafe_allow_html=True)
+        
         # Thai Version
         st.markdown("### 🇹🇭 ฉบับภาษาไทย")
         st.markdown(
@@ -1854,13 +1909,14 @@ def show_stitched_story(story_data: Dict[str, str]):
             unsafe_allow_html=True
         )
 
-        # Save Options
+        # Save Options with illustration
         st.markdown("### 💾 บันทึกเรื่องราวฉบับสมบูรณ์")
         save_col1, save_col2 = st.columns(2)
         
         with save_col1:
             if st.button("📥 บันทึกเป็น PDF", key="save_stitched_pdf", use_container_width=True):
-                pdf_data = create_bilingual_story_pdf(story_data)
+                # Add illustration to PDF if available
+                pdf_data = create_bilingual_story_pdf(story_data, illustration_url)
                 st.download_button(
                     "ดาวน์โหลด PDF",
                     data=pdf_data,
@@ -1871,10 +1927,14 @@ def show_stitched_story(story_data: Dict[str, str]):
         
         with save_col2:
             if st.button("💾 บันทึกข้อความ", key="save_stitched_text", use_container_width=True):
-                text_data = json.dumps(story_data, ensure_ascii=False, indent=2)
+                text_data = {
+                    **story_data,
+                    'illustration_url': illustration_url
+                }
+                json_str = json.dumps(text_data, ensure_ascii=False, indent=2)
                 st.download_button(
                     "ดาวน์โหลดข้อความ",
-                    data=text_data,
+                    data=json_str,
                     file_name=f"my_story_{datetime.now().strftime('%Y%m%d_%H%M')}.json",
                     mime="application/json",
                     key="download_stitched_text"
@@ -1884,8 +1944,8 @@ def show_stitched_story(story_data: Dict[str, str]):
         logging.error(f"Error showing stitched story: {str(e)}")
         st.error("เกิดข้อผิดพลาดในการแสดงเรื่องราว กรุณาลองใหม่อีกครั้ง")
 
-def create_bilingual_story_pdf(story_data: Dict[str, str]) -> bytes:
-    """Create a PDF with both English and Thai versions of the story"""
+def create_bilingual_story_pdf(story_data: Dict[str, str], illustration_url: str = None) -> bytes:
+    """Create a PDF with both English and Thai versions of the story and illustration"""
     try:
         buffer = io.BytesIO()
         doc = SimpleDocTemplate(
@@ -1896,7 +1956,7 @@ def create_bilingual_story_pdf(story_data: Dict[str, str]) -> bytes:
             topMargin=72,
             bottomMargin=72
         )
-        
+
         # Create styles
         styles = getSampleStyleSheet()
         title_style = ParagraphStyle(
@@ -1929,6 +1989,18 @@ def create_bilingual_story_pdf(story_data: Dict[str, str]) -> bytes:
         elements.append(Paragraph("My Story - JoyStory", title_style))
         elements.append(Spacer(1, 20))
         
+        # Add illustration if available
+        if illustration_url:
+            try:
+                # Download and add illustration
+                response = requests.get(illustration_url)
+                img_data = io.BytesIO(response.content)
+                img = Image(img_data, width=400, height=400)
+                elements.append(img)
+                elements.append(Spacer(1, 20))
+            except Exception as e:
+                logging.error(f"Error adding illustration to PDF: {str(e)}")
+        
         # Add English version
         elements.append(Paragraph("English Version", heading_style))
         elements.append(Paragraph(story_data['polished_english'], body_style))