Update app.py

app.py

@@ -12,9 +12,12 @@ import time
 import base64
 import json
 from typing import Dict, List, Tuple, Optional
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, BackgroundTasks
 from pydantic import BaseModel
 import random
+import gc
+import psutil
+import threading
 
 # External OCI API URL
 OCI_API_BASE_URL = "https://yukee1992-oci-story-book.hf.space"
@@ -66,6 +69,26 @@ current_pipe = None
 character_descriptions = {}
 character_seeds = {}  # Store seeds for consistent character generation
 
+# Memory monitoring function
+def monitor_memory():
+    """Monitor current memory usage"""
+    try:
+        process = psutil.Process()
+        memory_usage = process.memory_info().rss / 1024 / 1024  # MB
+        print(f"📊 Memory usage: {memory_usage:.2f} MB")
+        return memory_usage
+    except:
+        print("⚠️  Could not monitor memory (psutil not available)")
+        return 0
+
+# Memory cleanup function
+def cleanup_memory():
+    """Clean up memory and cache"""
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    print("🧹 Memory cleaned up")
+
 def load_model(model_name="dreamshaper-8"):
     """Load model into global cache - runs only once per model"""
     global model_cache, current_model_name, current_pipe
@@ -98,6 +121,7 @@ def load_model(model_name="dreamshaper-8"):
         current_model_name = model_name
         
         print(f"✅ Model loaded and cached: {model_name}")
+        monitor_memory()
         return pipe
         
     except Exception as e:
@@ -117,56 +141,54 @@ def load_model(model_name="dreamshaper-8"):
 print("🚀 Initializing Storybook Generator...")
 current_pipe = load_model("dreamshaper-8")
 print("✅ Default model loaded and ready!")
+monitor_memory()
 
-# PROFESSIONAL PROMPT ENGINEERING
-def enhance_prompt(prompt, style="childrens_book"):
-    """Transform basic prompts into professional-grade prompts"""
+# IMPROVED PROMPT ENGINEERING FOR CONSISTENCY
+def enhance_prompt(scene_visual, characters, style="childrens_book", page_number=1):
+    """Transform basic prompts into professional-grade prompts with character consistency"""
     
-    style_templates = {
-        "childrens_book": [
-            "masterpiece, best quality, 4K, ultra detailed, children's book illustration",
-            "watercolor painting, whimsical, cute, charming, storybook style",
-            "vibrant colors, soft lighting, magical, enchanting, dreamlike",
-            "Pixar style, Disney animation, high detail, professional artwork"
-        ],
-        "realistic": [
-            "photorealistic, 8K, ultra detailed, professional photography",
-            "sharp focus, studio lighting, high resolution, intricate details",
-            "realistic textures, natural lighting, cinematic quality"
-        ],
-        "fantasy": [
-            "epic fantasy art, digital painting, concept art, trending on artstation",
-            "magical, mystical, ethereal, otherworldly, fantasy illustration",
-            "dynamic composition, dramatic lighting, highly detailed"
-        ],
-        "anime": [
-            "anime style, Japanese animation, high quality, detailed artwork",
-            "beautiful anime illustration, vibrant colors, clean lines",
-            "studio ghibli style, makoto shinkai, professional anime art"
-        ]
+    # Character context - CRITICAL for consistency
+    character_context = ""
+    if characters:
+        character_context = " ".join([f"{char.description}" for char in characters])
+        character_context = f"Features: {character_context}. "
+    
+    # Scene continuity context
+    continuity_context = f"Scene {page_number}, " if page_number > 1 else ""
+    
+    # Style templates (less overwhelming)
+    style_presets = {
+        "childrens_book": "children's book illustration, watercolor style, whimsical, charming, vibrant colors",
+        "realistic": "photorealistic, professional photography, natural lighting, detailed",
+        "fantasy": "fantasy art, digital painting, magical, epic, concept art",
+        "anime": "anime style, Japanese animation, clean lines, vibrant colors"
     }
     
-    templates = style_templates.get(style, style_templates["childrens_book"])
-    style_prompt = templates[0]
+    style_prompt = style_presets.get(style, style_presets["childrens_book"])
     
-    enhanced = f"{style_prompt}, {prompt}"
+    # Build the main prompt - SCENE VISUAL COMES FIRST for importance
+    main_prompt = f"{continuity_context}{scene_visual}. {character_context}{style_prompt}"
     
-    quality_boosters = [
-        "intricate details", "beautiful composition", "perfect lighting",
-        "professional artwork", "award winning", "trending on artstation"
+    # Quality boosters (subtle)
+    quality_enhancers = [
+        "high quality", "detailed", "well composed", "good lighting"
     ]
     
-    boosters = random.sample(quality_boosters, 2)
-    enhanced += ", " + ", ".join(boosters)
+    # Add 1-2 quality enhancers randomly
+    enhancers = random.sample(quality_enhancers, min(2, len(quality_enhancers)))
+    main_prompt += ", " + ", ".join(enhancers)
     
+    # STRONG negative prompt to prevent deviations
     negative_prompt = (
         "blurry, low quality, low resolution, ugly, deformed, poorly drawn, "
         "bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, "
-        "disconnected limbs, mutation, mutated, ugly, disgusting, bad art, "
-        "beginner, amateur, distorted, watermark, signature, text, username"
+        "disconnected limbs, mutation, mutated, disgusting, bad art, "
+        "beginner, amateur, distorted, watermark, signature, text, username, "
+        "multiple people, crowd, group, different characters, inconsistent features, "
+        "changed appearance, different face, altered features"
     )
     
-    return enhanced, negative_prompt
+    return main_prompt, negative_prompt
 
 def save_complete_storybook_page(image, story_title, sequence_number, scene_text):
     """Save image AND text to OCI with organized structure"""
@@ -204,67 +226,64 @@ def save_complete_storybook_page(image, story_title, sequence_number, scene_text
     except Exception as e:
         return f"❌ Save failed: {str(e)}"
 
-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"""
+def get_character_seed(story_title, character_name, page_number):
+    """Get consistent but varied seed for character generation across pages"""
     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}")
+    seed_key = f"{character_name}_{page_number}"
+    
+    if seed_key not in character_seeds[story_title]:
+        # Generate stable but varied seed based on character, story, and page
+        base_seed = hash(f"{story_title}_{character_name}") % 1000000
+        # Add page variation but keep it consistent for same character+page
+        page_variation = (page_number * 13) % 1000  # Small variation per page
+        seed_value = (base_seed + page_variation) % 1000000
+        character_seeds[story_title][seed_key] = seed_value
+        print(f"🌱 Seed for {character_name} page {page_number}: {seed_value}")
     
-    return character_seeds[story_title][character_name]
+    return character_seeds[story_title][seed_key]
 
 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"""
+    """Generate a storybook page with STRONG character consistency"""
     global current_pipe, current_model_name
     
     try:
-        # Switch model if different from current - BUT DON'T RELOAD UNLESS NECESSARY
+        # Switch model if different from current
         if model_choice != current_model_name:
             print(f"🔄 Switching to model: {model_choice}")
-            current_pipe = load_model(model_choice)  # This uses cached version if available
-        
-        # ENHANCE PROMPT WITH CHARACTER CONTEXT
-        enhanced_visual = enhance_with_character_context(scene_visual, story_title, characters)
+            current_pipe = load_model(model_choice)
         
-        # 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)
+        # ENHANCE PROMPT WITH PROPER WEIGHTING
+        enhanced_prompt, negative_prompt = enhance_prompt(
+            scene_visual, characters, style, sequence_number
+        )
         
         print(f"📖 Generating page {sequence_number} for: {story_title}")
+        print(f"📝 Prompt: {enhanced_prompt[:100]}...")
         if characters:
             print(f"👤 Characters: {[char.name for char in characters]}")
         
-        # Use consistent seed for character generation
+        # Use consistent seed strategy
         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)
+            # Use seed from main character for this specific page
+            main_char_seed = get_character_seed(story_title, characters[0].name, sequence_number)
             generator.manual_seed(main_char_seed)
-            print(f"🌱 Using seed {main_char_seed} for character consistency")
+            print(f"🌱 Using seed {main_char_seed} for {characters[0].name} on page {sequence_number}")
         else:
-            seed = int(time.time())
-            generator.manual_seed(seed)
-            print(f"🌱 Using timestamp seed {seed}")
+            # Still use consistent seed for non-character scenes
+            scene_seed = hash(f"{story_title}_{sequence_number}") % 1000000
+            generator.manual_seed(scene_seed)
+            print(f"🌱 Using scene seed {scene_seed} for page {sequence_number}")
         
-        # Generate high-quality image - USE THE GLOBAL current_pipe
+        # Generate with higher steps for better consistency
         image = current_pipe(
             prompt=enhanced_prompt,
             negative_prompt=negative_prompt,
-            num_inference_steps=30,
-            guidance_scale=8.5,
+            num_inference_steps=35,  # Increased for better detail
+            guidance_scale=7.5,      # Slightly lower for more prompt adherence
             width=768,
             height=768,
             generator=generator
@@ -279,7 +298,7 @@ def generate_storybook_page(scene_visual, story_title, sequence_number, scene_te
         return None, f"❌ Generation failed: {str(e)}"
 
 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"""
+    """Generate complete storybook with memory management and consistency"""
     global character_descriptions, current_pipe
     
     results = []
@@ -290,40 +309,63 @@ def batch_generate_complete_storybook(story_title, scenes_data, characters, mode
     print(f"👤 Characters: {len(characters)}")
     print(f"🎨 Using model: {model_choice}")
     
+    # Initial memory check
+    initial_memory = monitor_memory()
+    
     # 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 - THIS IS THE KEY FIX
+    # Load model once at the beginning
     print(f"🔧 Loading model for this storybook...")
     current_pipe = load_model(model_choice)
     
     start_time = time.time()
     
     for i, scene_data in enumerate(scenes_data, 1):
-        scene_visual = scene_data.get('visual', '')
-        scene_text = scene_data.get('text', '')
-        
-        print(f"🔄 Generating page {i}/{len(scenes_data)}...")
-        image, status = generate_storybook_page(
-            scene_visual, story_title, i, scene_text, characters, model_choice, style
-        )
-        
-        if image:
-            results.append((f"Page {i}", image, scene_text))
-        status_messages.append(f"Page {i}: {status}")
+        try:
+            # Clean memory every 2 pages
+            if i % 2 == 0:
+                cleanup_memory()
+                monitor_memory()
+            
+            scene_visual = scene_data.get('visual', '')
+            scene_text = scene_data.get('text', '')
+            
+            print(f"🔄 Generating page {i}/{len(scenes_data)}...")
+            print(f"🎬 Scene: {scene_visual[:80]}...")
+            
+            image, status = generate_storybook_page(
+                scene_visual, story_title, i, scene_text, characters, model_choice, style
+            )
+            
+            if image:
+                results.append((f"Page {i}", image, scene_text))
+            status_messages.append(f"Page {i}: {status}")
+            
+            # Brief pause between generations
+            if i < len(scenes_data):
+                time.sleep(2)
+            
+        except Exception as e:
+            error_msg = f"❌ Failed page {i}: {str(e)}"
+            print(error_msg)
+            status_messages.append(error_msg)
+            # Continue with next page instead of stopping
     
     total_time = time.time() - start_time
+    final_memory = monitor_memory()
+    
     print(f"✅ Batch generation completed in {total_time:.2f} seconds")
-    print(f"📊 Average: {total_time/len(scenes_data):.2f} seconds per page")
+    print(f"📊 Memory delta: {final_memory - initial_memory:.2f} MB")
     
     return results, "\n".join(status_messages)
 
 # FastAPI endpoint for n8n
 @app.post("/api/generate-storybook")
 async def api_generate_storybook(request: StorybookRequest):
-    """API endpoint for n8n automation - OPTIMIZED with character consistency"""
+    """API endpoint for n8n automation - OPTIMIZED for consistency"""
     try:
         print(f"📚 Received storybook request: {request.story_title}")
         print(f"📖 Pages to generate: {len(request.scenes)}")
@@ -331,14 +373,14 @@ async def api_generate_storybook(request: StorybookRequest):
         
         if request.characters:
             for char in request.characters:
-                print(f"   - {char.name}: {char.description[:50]}...")
+                print(f"   - {char.name}: {char.description}")
         
         start_time = time.time()
         
         # Convert to scene data format
         scenes_data = [{"visual": scene.visual, "text": scene.text} for scene in request.scenes]
         
-        # Generate storybook (model loads only once)
+        # Generate storybook
         results, status = batch_generate_complete_storybook(
             request.story_title,
             scenes_data,
@@ -374,20 +416,23 @@ async def api_generate_storybook(request: StorybookRequest):
         traceback.print_exc()
         raise HTTPException(status_code=500, detail=error_msg)
 
-# Health check endpoint
+# Health check endpoint with memory info
 @app.get("/api/health")
 async def health_check():
+    memory_info = monitor_memory()
     return {
         "status": "healthy",
         "service": "Storybook Generator API",
         "timestamp": datetime.now().isoformat(),
+        "memory_usage_mb": round(memory_info, 2),
         "models_loaded": list(model_cache.keys()),
         "current_model": current_model_name,
         "cached_models_count": len(model_cache),
-        "stories_tracked": len(character_descriptions)
+        "stories_tracked": len(character_descriptions),
+        "character_seeds_stored": sum(len(seeds) for seeds in character_seeds.values())
     }
 
-# Gradio Interface Functions (simplified)
+# Gradio Interface Functions
 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: