Update app.py

app.py

@@ -11,7 +11,7 @@ import tempfile
 import time
 import base64
 import json
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Optional
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 import random
@@ -37,9 +37,14 @@ class StoryScene(BaseModel):
     visual: str
     text: str
 
+class CharacterDescription(BaseModel):
+    name: str
+    description: str
+
 class StorybookRequest(BaseModel):
     story_title: str
     scenes: List[StoryScene]
+    characters: List[CharacterDescription] = []
     model_choice: str = "dreamshaper-8"
     style: str = "childrens_book"
 
@@ -57,6 +62,10 @@ model_cache = {}
 current_model_name = None
 pipe = None
 
+# Character consistency tracking
+character_descriptions = {}
+character_seeds = {}  # Store seeds for consistent character generation
+
 def load_model(model_name="dreamshaper-8"):
     """Load model into global cache - runs only once per model"""
     global model_cache, current_model_name, pipe
@@ -189,8 +198,28 @@ def save_complete_storybook_page(image, story_title, sequence_number, scene_text
     except Exception as e:
         return f"❌ Save failed: {str(e)}"
 
-def generate_storybook_page(scene_visual, story_title, sequence_number, scene_text, model_choice="dreamshaper-8", style="childrens_book"):
-    """Generate a storybook page with both image and text"""
+def enhance_with_character_context(scene_visual, story_title, characters):
+    """Add character descriptions to maintain consistency"""
+    if characters:
+        character_context = " ".join([f"{char.name}: {char.description}" for char in characters])
+        return f"Character descriptions: {character_context}. {scene_visual}"
+    return scene_visual
+
+def get_character_seed(story_title, character_name):
+    """Get consistent seed for character generation"""
+    if story_title not in character_seeds:
+        character_seeds[story_title] = {}
+    
+    if character_name not in character_seeds[story_title]:
+        # Generate a stable seed based on character name and story title
+        seed_value = hash(f"{story_title}_{character_name}") % 1000000
+        character_seeds[story_title][character_name] = seed_value
+        print(f"🌱 Seed for {character_name}: {seed_value}")
+    
+    return character_seeds[story_title][character_name]
+
+def generate_storybook_page(scene_visual, story_title, sequence_number, scene_text, characters, model_choice="dreamshaper-8", style="childrens_book"):
+    """Generate a storybook page with character consistency"""
     global pipe, current_model_name
     
     try:
@@ -199,10 +228,27 @@ def generate_storybook_page(scene_visual, story_title, sequence_number, scene_te
             print(f"🔄 Switching to model: {model_choice}")
             pipe = load_model(model_choice)
         
-        # Enhanced prompt for better quality
-        enhanced_prompt, negative_prompt = enhance_prompt(scene_visual, style)
+        # ENHANCE PROMPT WITH CHARACTER CONTEXT
+        enhanced_visual = enhance_with_character_context(scene_visual, story_title, characters)
+        
+        # Add scene continuity context
+        if sequence_number > 1:
+            enhanced_visual = f"Scene {sequence_number}, maintain character consistency from previous scenes. {enhanced_visual}"
+        
+        enhanced_prompt, negative_prompt = enhance_prompt(enhanced_visual, style)
         
         print(f"📖 Generating page {sequence_number} for: {story_title}")
+        if characters:
+            print(f"👤 Characters: {[char.name for char in characters]}")
+        
+        # Use consistent seed for character generation
+        generator = torch.Generator(device="cpu")
+        if characters:
+            # Use seed from main character for consistency
+            main_char_seed = get_character_seed(story_title, characters[0].name)
+            generator.manual_seed(main_char_seed)
+        else:
+            generator.manual_seed(int(time.time()))
         
         # Generate high-quality image
         image = pipe(
@@ -212,7 +258,7 @@ def generate_storybook_page(scene_visual, story_title, sequence_number, scene_te
             guidance_scale=8.5,
             width=768,
             height=768,
-            generator=torch.Generator(device="cpu").manual_seed(int(time.time()))
+            generator=generator
         ).images[0]
         
         # Save both image and text
@@ -223,15 +269,23 @@ def generate_storybook_page(scene_visual, story_title, sequence_number, scene_te
     except Exception as e:
         return None, f"❌ Generation failed: {str(e)}"
 
-def batch_generate_complete_storybook(story_title, scenes_data, model_choice="dreamshaper-8", style="childrens_book"):
+def batch_generate_complete_storybook(story_title, scenes_data, characters, model_choice="dreamshaper-8", style="childrens_book"):
     """Generate complete storybook with images and text - MODEL LOADS ONLY ONCE"""
+    global character_descriptions
+    
     results = []
     status_messages = []
     
     print(f"📚 Starting batch generation for: {story_title}")
     print(f"📖 Total pages: {len(scenes_data)}")
+    print(f"👤 Characters: {len(characters)}")
     print(f"🎨 Using model: {model_choice}")
     
+    # Store character descriptions for this story
+    if characters:
+        character_descriptions[story_title] = characters
+        print(f"✅ Character context stored for {story_title}")
+    
     # Load model once at the beginning
     global pipe
     pipe = load_model(model_choice)
@@ -244,7 +298,7 @@ def batch_generate_complete_storybook(story_title, scenes_data, model_choice="dr
         
         print(f"🔄 Generating page {i}/{len(scenes_data)}...")
         image, status = generate_storybook_page(
-            scene_visual, story_title, i, scene_text, model_choice, style
+            scene_visual, story_title, i, scene_text, characters, model_choice, style
         )
         
         if image:
@@ -260,12 +314,15 @@ def batch_generate_complete_storybook(story_title, scenes_data, model_choice="dr
 # FastAPI endpoint for n8n
 @app.post("/api/generate-storybook")
 async def api_generate_storybook(request: StorybookRequest):
-    """API endpoint for n8n automation - OPTIMIZED"""
+    """API endpoint for n8n automation - OPTIMIZED with character consistency"""
     try:
         print(f"📚 Received storybook request: {request.story_title}")
         print(f"📖 Pages to generate: {len(request.scenes)}")
-        print(f"🎨 Model: {request.model_choice}")
-        print(f"🎨 Style: {request.style}")
+        print(f"👤 Characters received: {len(request.characters)}")
+        
+        if request.characters:
+            for char in request.characters:
+                print(f"   - {char.name}: {char.description[:50]}...")
         
         # Convert to scene data format
         scenes_data = [{"visual": scene.visual, "text": scene.text} for scene in request.scenes]
@@ -274,6 +331,7 @@ async def api_generate_storybook(request: StorybookRequest):
         results, status = batch_generate_complete_storybook(
             request.story_title,
             scenes_data,
+            request.characters,
             request.model_choice,
             request.style
         )
@@ -282,6 +340,7 @@ async def api_generate_storybook(request: StorybookRequest):
             "status": "success",
             "story_title": request.story_title,
             "total_pages": len(request.scenes),
+            "characters_used": len(request.characters),
             "generated_pages": len(results),
             "message": status,
             "folder_path": f"storybook-library/stories/{request.story_title.replace(' ', '_')}/",
@@ -297,6 +356,8 @@ async def api_generate_storybook(request: StorybookRequest):
     except Exception as e:
         error_msg = f"Storybook generation failed: {str(e)}"
         print(f"❌ {error_msg}")
+        import traceback
+        traceback.print_exc()
         raise HTTPException(status_code=500, detail=error_msg)
 
 # Health check endpoint
@@ -308,58 +369,34 @@ async def health_check():
         "timestamp": datetime.now().isoformat(),
         "models_loaded": list(model_cache.keys()),
         "current_model": current_model_name,
-        "cached_models_count": len(model_cache)
+        "cached_models_count": len(model_cache),
+        "stories_tracked": len(character_descriptions)
     }
 
-# ... (keep the rest of your Gradio interface code the same) ...
-
-# Gradio Interface Functions
+# Gradio Interface Functions (simplified)
 def generate_single_page(prompt, story_title, scene_text, model_choice, style):
     """Generate a single page for Gradio interface"""
     if not prompt or not story_title:
         return None, "❌ Please enter both scene description and story title"
     
     image, status = generate_storybook_page(
-        prompt, story_title, 1, scene_text or "", model_choice, style
+        prompt, story_title, 1, scene_text or "", [], model_choice, style
     )
     return image, status
 
-def generate_full_storybook(story_title, scenes_text, model_choice, style):
-    """Generate full storybook for Gradio interface"""
-    if not story_title or not scenes_text:
-        return [], "❌ Please provide story title and scenes"
-    
-    # Parse scenes from text input
-    scenes = []
-    for i, line in enumerate(scenes_text.split('\n')):
-        if line.strip():
-            scenes.append({
-                "visual": line.strip(),
-                "text": f"Page {i+1} text description"
-            })
-    
-    results, status = batch_generate_complete_storybook(
-        story_title, scenes, model_choice, style
-    )
-    return results, status
-
 # Create the Gradio interface
 with gr.Blocks(title="Storybook Generator", theme="soft") as demo:
     gr.Markdown("# 📚 Storybook Generator")
-    gr.Markdown("Create beautiful storybooks with images and text")
+    gr.Markdown("Create beautiful storybooks with consistent characters")
     
     with gr.Row():
         with gr.Column(scale=1):
-            gr.Markdown("### 📖 Story Information")
-            
             story_title_input = gr.Textbox(
                 label="Story Title",
                 placeholder="Enter your story title...",
                 lines=1
             )
             
-            gr.Markdown("### 🎯 Quality Settings")
-            
             model_choice = gr.Dropdown(
                 label="AI Model",
                 choices=list(MODEL_CHOICES.keys()),
@@ -373,8 +410,6 @@ with gr.Blocks(title="Storybook Generator", theme="soft") as demo:
             )
         
         with gr.Column(scale=2):
-            gr.Markdown("### 🎨 Single Page Generation")
-            
             prompt_input = gr.Textbox(
                 label="Visual Description",
                 placeholder="Describe the scene for image generation...",
@@ -391,35 +426,11 @@ with gr.Blocks(title="Storybook Generator", theme="soft") as demo:
             image_output = gr.Image(label="Generated Page", height=400)
             status_output = gr.Textbox(label="Status", interactive=False)
     
-    with gr.Row():
-        gr.Markdown("### 📚 Complete Storybook Generation")
-    
-    with gr.Row():
-        with gr.Column():
-            scenes_input = gr.Textbox(
-                label="Visual Descriptions (One per line)",
-                placeholder="Enter each page's visual description on separate lines...",
-                lines=6
-            )
-            
-            batch_btn = gr.Button("📖 Generate Complete Storybook", variant="primary")
-        
-        with gr.Column():
-            batch_status = gr.Textbox(label="Generation Status", interactive=False, lines=6)
-            batch_gallery = gr.Gallery(label="Storybook Pages", columns=2, height=600)
-
-    # Connect buttons to functions
     generate_btn.click(
         fn=generate_single_page,
         inputs=[prompt_input, story_title_input, text_input, model_choice, style_choice],
         outputs=[image_output, status_output]
     )
-    
-    batch_btn.click(
-        fn=generate_full_storybook,
-        inputs=[story_title_input, scenes_input, model_choice, style_choice],
-        outputs=[batch_gallery, batch_status]
-    )
 
 # Mount Gradio app to FastAPI
 app = gr.mount_gradio_app(app, demo, path="/")