Update app.py

app.py

@@ -18,6 +18,7 @@ import random
 import gc
 import psutil
 import threading
+from functools import lru_cache
 
 # External OCI API URL
 OCI_API_BASE_URL = "https://yukee1992-oci-story-book.hf.space"
@@ -51,7 +52,7 @@ class StorybookRequest(BaseModel):
     model_choice: str = "sdxl"
     style: str = "childrens_book"
 
-# MODEL SELECTION - SDXL handles longer prompts better
+# MODEL SELECTION
 MODEL_CHOICES = {
     "sdxl": "stabilityai/stable-diffusion-xl-base-1.0",
     "sdxl-turbo": "stabilityai/sdxl-turbo",
@@ -59,10 +60,11 @@ MODEL_CHOICES = {
     "realistic-vision": "SG161222/Realistic_Vision_V5.1", 
 }
 
-# GLOBAL MODEL CACHE
+# GLOBAL MODEL CACHE with proper locking
 model_cache = {}
 current_model_name = None
 current_pipe = None
+model_lock = threading.Lock()
 
 # Character consistency tracking
 character_descriptions = {}
@@ -80,151 +82,139 @@ def cleanup_memory():
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
-    print("🧹 Memory cleaned up")
 
 def load_model(model_name="sdxl"):
+    """Thread-safe model loading with proper caching"""
     global model_cache, current_model_name, current_pipe
     
-    if model_name in model_cache:
-        current_pipe = model_cache[model_name]
-        current_model_name = model_name
-        return current_pipe
-    
-    print(f"🔄 Loading model: {model_name}")
-    try:
-        if model_name in ["sdxl", "sdxl-turbo"]:
-            model_id = MODEL_CHOICES[model_name]
-            pipe = StableDiffusionXLPipeline.from_pretrained(
-                model_id,
-                torch_dtype=torch.float32,
-                use_safetensors=True,
-                safety_checker=None,
-                requires_safety_checker=False
-            )
-            pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-        else:
-            model_id = MODEL_CHOICES.get(model_name, "lykon/dreamshaper-8")
-            pipe = StableDiffusionPipeline.from_pretrained(
-                model_id, 
-                torch_dtype=torch.float32,
-                safety_checker=None,
-                requires_safety_checker=False
-            )
-            pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+    with model_lock:
+        if model_name in model_cache:
+            print(f"✅ Using cached model: {model_name}")
+            current_pipe = model_cache[model_name]
+            current_model_name = model_name
+            return current_pipe
         
-        pipe = pipe.to("cpu")
-        model_cache[model_name] = pipe
-        current_pipe = pipe
-        current_model_name = model_name
-        
-        print(f"✅ Model loaded: {model_name}")
-        return pipe
-        
-    except Exception as e:
-        print(f"❌ Model loading failed: {e}")
-        pipe = StableDiffusionPipeline.from_pretrained(
-            "runwayml/stable-diffusion-v1-5", 
-            torch_dtype=torch.float32
-        ).to("cpu")
-        model_cache[model_name] = pipe
-        return pipe
+        print(f"🔄 Loading model: {model_name}")
+        try:
+            if model_name in ["sdxl", "sdxl-turbo"]:
+                model_id = MODEL_CHOICES[model_name]
+                pipe = StableDiffusionXLPipeline.from_pretrained(
+                    model_id,
+                    torch_dtype=torch.float32,
+                    use_safetensors=True,
+                    safety_checker=None,
+                    requires_safety_checker=False
+                )
+                pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+            else:
+                model_id = MODEL_CHOICES.get(model_name, "lykon/dreamshaper-8")
+                pipe = StableDiffusionPipeline.from_pretrained(
+                    model_id, 
+                    torch_dtype=torch.float32,
+                    safety_checker=None,
+                    requires_safety_checker=False
+                )
+                pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+            
+            pipe = pipe.to("cpu")
+            model_cache[model_name] = pipe
+            current_pipe = pipe
+            current_model_name = model_name
+            
+            print(f"✅ Model loaded and cached: {model_name}")
+            return pipe
+            
+        except Exception as e:
+            print(f"❌ Model loading failed: {e}")
+            pipe = StableDiffusionPipeline.from_pretrained(
+                "runwayml/stable-diffusion-v1-5", 
+                torch_dtype=torch.float32
+            ).to("cpu")
+            model_cache[model_name] = pipe
+            return pipe
 
 # Initialize default model
 print("🚀 Initializing Storybook Generator...")
 current_pipe = load_model("sdxl")
 print("✅ Model loaded and ready!")
 
-# TRUE UNLIMITED PROMPT SOLUTION
-def create_compressed_prompt(scene_visual, characters, style="childrens_book", page_number=1):
-    """
-    Create a compressed but comprehensive prompt that fits within token limits
-    while preserving ALL important information
-    """
-    # Extract ONLY the most critical character features
-    character_features = []
+# OPTIMIZED PROMPT COMPRESSION
+@lru_cache(maxsize=100)
+def compress_prompt(text, style="childrens_book"):
+    """Cache compressed prompts to avoid recomputation"""
+    # Simple compression: remove redundant words and shorten
+    words = text.split()
+    if len(words) <= 50:
+        return text
+    
+    # Keep first 40 words (most important part) and key descriptors
+    compressed = ' '.join(words[:40])
+    
+    # Add style context
+    style_context = {
+        "childrens_book": "children's book style",
+        "realistic": "realistic style", 
+        "fantasy": "fantasy style",
+        "anime": "anime style"
+    }
+    
+    return f"{compressed}... {style_context.get(style, '')} masterpiece 4K"
+
+def create_optimized_prompt(scene_visual, characters, style="childrens_book", page_number=1):
+    """Create optimized prompt within token limits"""
+    # Compress the scene visual
+    scene_compressed = compress_prompt(scene_visual, style)
+    
+    # Extract character essentials
+    char_descriptors = []
     if characters:
         for char in characters:
-            if hasattr(char, 'description'):
+            if hasattr(char, 'name'):
+                name = char.name
                 desc = char.description
             elif isinstance(char, dict):
+                name = char.get('name', 'Unknown')
                 desc = char.get('description', '')
             else:
-                desc = str(char)
+                continue
             
-            # Extract key features: age, appearance, clothing
+            # Extract key features
             import re
-            # Get age if mentioned
-            age_match = re.search(r'(\d+)[\- ]?year[\- ]?old', desc, re.IGNORECASE)
-            age = f"{age_match.group(1)} year old" if age_match else ""
-            
             # Get species/type
-            species_match = re.search(r'(rabbit|hedgehog|bird|dog|cat|fox|bear|dragon|unicorn|human|girl|boy)', desc, re.IGNORECASE)
+            species_match = re.search(r'(rabbit|hedgehog|bird|dog|cat|fox|bear|dragon|human|girl|boy)', desc, re.IGNORECASE)
             species = species_match.group(1) if species_match else "character"
             
-            # Get color/main features
-            color_match = re.search(r'(blonde|brown|black|white|blue|red|green|yellow|golden|silver)', desc, re.IGNORECASE)
+            # Get color
+            color_match = re.search(r'(white|black|brown|blue|red|green|yellow|golden|pink)', desc, re.IGNORECASE)
             color = color_match.group(1) if color_match else ""
             
-            # Get key accessories
-            accessories = []
-            if 'glasses' in desc.lower(): accessories.append('glasses')
-            if 'dress' in desc.lower(): accessories.append('dress')
-            if 'hat' in desc.lower(): accessories.append('hat')
-            if 'satchel' in desc.lower(): accessories.append('satchel')
-            
-            # Build compressed description
-            compressed_desc = f"{age} {color} {species}".strip()
-            if accessories:
-                compressed_desc += f" with {', '.join(accessories)}"
-            
-            character_features.append(compressed_desc)
-    
-    # Build scene context
-    continuity_context = f"scene {page_number}" if page_number > 1 else ""
-    
-    # Style templates (compressed)
-    style_presets = {
-        "childrens_book": "children's book illustration, watercolor, whimsical",
-        "realistic": "photorealistic, professional photography",
-        "fantasy": "fantasy art, digital painting, magical",
-        "anime": "anime style, clean lines, vibrant colors"
-    }
+            char_descriptors.append(f"{color} {species}".strip())
     
-    style_prompt = style_presets.get(style, style_presets["childrens_book"])
+    # Build the final prompt
+    continuity = f"scene {page_number} " if page_number > 1 else ""
+    chars_text = f"Characters: {', '.join(char_descriptors)}. " if char_descriptors else ""
     
-    # Build the final compressed prompt
-    compressed_prompt = f"{continuity_context} {scene_visual}"
+    final_prompt = f"{continuity}{scene_compressed}. {chars_text}masterpiece best quality 4K"
     
-    if character_features:
-        compressed_prompt += f". Characters: {', '.join(character_features)}"
-    
-    compressed_prompt += f". Style: {style_prompt}. masterpiece, best quality, 4K"
-    
-    # Ensure it's within reasonable length
-    words = compressed_prompt.split()
+    # Ensure it's under 60 words
+    words = final_prompt.split()
     if len(words) > 60:
-        compressed_prompt = ' '.join(words[:60]) + '...'
+        final_prompt = ' '.join(words[:60])
     
-    return compressed_prompt
+    return final_prompt
 
 def enhance_prompt(scene_visual, characters, style="childrens_book", page_number=1):
-    """
-    Create optimized prompt that preserves essence while fitting token limits
-    """
-    # Use compressed prompt for the actual generation
-    main_prompt = create_compressed_prompt(scene_visual, characters, style, page_number)
+    """Create optimized prompt"""
+    main_prompt = create_optimized_prompt(scene_visual, characters, style, page_number)
     
-    print(f"📝 Compressed prompt: {main_prompt}")
+    print(f"📝 Optimized prompt: {main_prompt}")
     print(f"📏 Length: {len(main_prompt.split())} words")
     
     # Negative prompt
     negative_prompt = (
-        "blurry, low quality, ugly, deformed, poorly drawn, bad anatomy, "
-        "wrong anatomy, extra limb, missing limb, floating limbs, "
-        "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, low resolution"
+        "blurry, low quality, ugly, deformed, bad anatomy, "
+        "watermark, signature, text, username, multiple people, "
+        "inconsistent features, low resolution"
     )
     
     return main_prompt, negative_prompt
@@ -274,18 +264,22 @@ def get_character_seed(story_title, character_name, page_number):
     return character_seeds[story_title][seed_key]
 
 def generate_storybook_page(scene_visual, story_title, sequence_number, scene_text, characters, model_choice="sdxl", style="childrens_book"):
+    """Generate a single page - OPTIMIZED VERSION"""
     global current_pipe, current_model_name
     
     try:
+        # ONLY load model if different from current
         if model_choice != current_model_name:
+            print(f"🔄 Switching to model: {model_choice}")
             current_pipe = load_model(model_choice)
+        else:
+            print(f"✅ Using already loaded model: {model_choice}")
         
         enhanced_prompt, negative_prompt = enhance_prompt(
             scene_visual, characters, style, sequence_number
         )
         
         print(f"📖 Generating page {sequence_number}")
-        print(f"📝 Using prompt: {enhanced_prompt}")
         
         if characters:
             char_names = []
@@ -308,17 +302,23 @@ def generate_storybook_page(scene_visual, story_title, sequence_number, scene_te
             scene_seed = hash(f"{story_title}_{sequence_number}") % 1000000
             generator.manual_seed(scene_seed)
         
-        # Generate image
+        # Generate image with optimized parameters
+        print("⏳ Starting image generation...")
+        start_time = time.time()
+        
         image = current_pipe(
             prompt=enhanced_prompt,
             negative_prompt=negative_prompt,
-            num_inference_steps=35,
-            guidance_scale=7.5,
-            width=768,
-            height=768,
+            num_inference_steps=25,  # Reduced from 35 for speed
+            guidance_scale=7.0,
+            width=512,  # Reduced from 768 for speed
+            height=512,
             generator=generator
         ).images[0]
         
+        gen_time = time.time() - start_time
+        print(f"✅ Image generated in {gen_time:.1f} seconds")
+        
         save_status = save_complete_storybook_page(image, story_title, sequence_number, scene_text)
         return image, save_status
         
@@ -326,48 +326,52 @@ 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="sdxl", style="childrens_book"):
-    global character_descriptions, current_pipe
+    """Batch generation with significant optimizations"""
+    global current_pipe
     
     results = []
     status_messages = []
     
-    print(f"📚 Starting batch generation: {story_title}")
+    print(f"📚 Starting OPTIMIZED batch generation: {story_title}")
     print(f"📖 Pages: {len(scenes_data)}")
     print(f"👤 Characters: {len(characters)}")
     
-    if characters:
-        character_descriptions[story_title] = characters
-    
+    # Load model ONCE at the beginning
+    print(f"🔧 Loading model once for entire batch...")
     current_pipe = load_model(model_choice)
-    start_time = time.time()
+    batch_start_time = time.time()
     
     for i, scene_data in enumerate(scenes_data, 1):
         try:
-            if i % 2 == 0:
-                cleanup_memory()
-            
             scene_visual = scene_data.get('visual', '')
             scene_text = scene_data.get('text', '')
             
             print(f"🔄 Generating page {i}/{len(scenes_data)}...")
+            page_start_time = time.time()
+            
             image, status = generate_storybook_page(
                 scene_visual, story_title, i, scene_text, characters, model_choice, style
             )
             
+            page_time = time.time() - page_start_time
+            print(f"⏰ Page {i} completed in {page_time:.1f} seconds")
+            
             if image:
                 results.append((f"Page {i}", image, scene_text))
             status_messages.append(f"Page {i}: {status}")
             
-            if i < len(scenes_data):
-                time.sleep(2)
+            # Clean memory every 3 pages
+            if i % 3 == 0:
+                cleanup_memory()
             
         except Exception as e:
             error_msg = f"❌ Failed page {i}: {str(e)}"
             print(error_msg)
             status_messages.append(error_msg)
     
-    total_time = time.time() - start_time
+    total_time = time.time() - batch_start_time
     print(f"✅ Batch completed in {total_time:.2f} seconds")
+    print(f"📊 Average: {total_time/len(scenes_data):.1f} seconds per page")
     
     return results, "\n".join(status_messages)
 
@@ -380,7 +384,15 @@ async def api_generate_storybook(request: StorybookRequest):
         
         start_time = time.time()
         scenes_data = [{"visual": scene.visual, "text": scene.text} for scene in request.scenes]
-        characters_dict = [char.dict() for char in request.characters]
+        
+        # Convert characters to dict ONCE
+        characters_dict = []
+        for char in request.characters:
+            if hasattr(char, 'dict'):
+                characters_dict.append(char.dict())
+            else:
+                characters_dict.append({"name": getattr(char, 'name', 'Unknown'), 
+                                      "description": getattr(char, 'description', '')})
         
         results, status = batch_generate_complete_storybook(
             request.story_title,