Update app.py

app.py

@@ -8,7 +8,8 @@ import os
 from datetime import datetime
 import re
 import time
-from typing import List, Optional
+import json
+from typing import List, Optional, Dict
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 import gc
@@ -57,6 +58,7 @@ class StorybookResponse(BaseModel):
     message: str
     folder_path: str
     pages: List[dict]
+    request_id: str
 
 # MODEL SELECTION
 MODEL_CHOICES = {
@@ -72,6 +74,7 @@ model_lock = threading.Lock()
 
 # Character consistency tracking
 character_seeds = {}
+active_requests = {}
 
 def monitor_memory():
     try:
@@ -125,70 +128,62 @@ print("🚀 Initializing Storybook Generator...")
 load_model("dreamshaper-8")
 print("✅ Model loaded and ready!")
 
-# PROMPT OPTIMIZATION
+# PROMPT OPTIMIZATION - PRESERVE FULL DESCRIPTIONS
 def optimize_prompt(scene_visual, characters, style="childrens_book", page_number=1):
     """
-    Create a prompt that FITS within 77 tokens
+    Create a prompt that PRESERVES all visual descriptions while fitting 77 tokens
     """
-    # Extract character essence
+    # 1. PRESERVE THE ENTIRE SCENE VISUAL DESCRIPTION (most important)
+    scene_prompt = scene_visual
+    
+    # 2. Extract only ESSENTIAL character features (not full descriptions)
     character_essence = ""
     if characters:
-        char_descriptors = []
+        char_names = []
         for char in characters:
-            desc = char.get('description', '') if isinstance(char, dict) else getattr(char, 'description', '')
-            
-            import re
-            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"
-            
-            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 ""
-            
-            key_feature = ""
-            if 'glasses' in desc.lower(): key_feature = "with glasses"
-            elif 'dress' in desc.lower(): key_feature = "in dress"
-            elif 'hat' in desc.lower(): key_feature = "with hat"
-            
-            char_descriptors.append(f"{color} {species} {key_feature}".strip())
+            char_name = char.get('name', '') if isinstance(char, dict) else getattr(char, 'name', '')
+            char_names.append(char_name.split()[0])  # Just first name
         
-        character_essence = f"Characters: {', '.join(char_descriptors)}. "
-    
-    # Compress scene description
-    scene_words = scene_visual.split()
-    if len(scene_words) > 30:
-        scene_compressed = ' '.join(scene_words[:30])
-    else:
-        scene_compressed = scene_visual
+        character_essence = f" featuring {', '.join(char_names)}"
     
-    # Style context
+    # 3. Add style context briefly
     style_context = {
-        "childrens_book": "children's book illustration",
-        "realistic": "photorealistic",
-        "fantasy": "fantasy art",
+        "childrens_book": "children's book illustration style",
+        "realistic": "photorealistic style",
+        "fantasy": "fantasy art style",
         "anime": "anime style"
-    }.get(style, "children's book illustration")
+    }.get(style, "children's book illustration style")
     
-    # Build final prompt
-    continuity = f"Scene {page_number}: " if page_number > 1 else ""
-    final_prompt = f"{continuity}{scene_compressed}. {character_essence}{style_context}. masterpiece, best quality"
+    # 4. Build the final prompt - SCENE DESCRIPTION COMES FIRST
+    continuity = f"Scene {page_number}, " if page_number > 1 else ""
+    final_prompt = f"{continuity}{scene_prompt}{character_essence}. {style_context}. high quality, detailed"
     
-    # Ensure it's under 55 words for safety
+    # 5. If still too long, prioritize scene description over style
     words = final_prompt.split()
-    if len(words) > 55:
-        final_prompt = ' '.join(words[:55])
+    if len(words) > 60:
+        # Keep the scene description intact, trim the end
+        scene_words = scene_visual.split()
+        if len(scene_words) > 45:
+            # If scene itself is too long, keep first 40 words of scene
+            scene_part = ' '.join(scene_words[:40])
+            final_prompt = f"{continuity}{scene_part}...{character_essence}. {style_context}"
+        else:
+            # Keep entire scene, trim style part
+            final_prompt = f"{continuity}{scene_visual}{character_essence}. high quality"
     
-    print(f"📝 Optimized prompt: {final_prompt}")
+    print(f"📝 Final prompt: {final_prompt}")
     print(f"📏 Length: {len(final_prompt.split())} words")
     
     return final_prompt
 
 def enhance_prompt(scene_visual, characters, style="childrens_book", page_number=1):
-    """Create optimized prompt"""
+    """Create optimized prompt that preserves visual descriptions"""
     main_prompt = optimize_prompt(scene_visual, characters, style, page_number)
     
     negative_prompt = (
         "blurry, low quality, ugly, deformed, bad anatomy, "
-        "watermark, text, username, multiple people, inconsistent"
+        "watermark, text, username, multiple people, inconsistent, "
+        "missing limbs, extra limbs, disfigured, malformed"
     )
     
     return main_prompt, negative_prompt
@@ -237,120 +232,146 @@ def get_character_seed(story_title, character_name, page_number):
     
     return character_seeds[story_title][seed_key]
 
-def generate_single_page(scene_visual, scene_text, story_title, sequence_number, characters, model_choice, style):
-    """Generate a single page"""
-    try:
-        print(f"🔄 Generating page {sequence_number}...")
-        
-        enhanced_prompt, negative_prompt = enhance_prompt(
-            scene_visual, characters, style, sequence_number
-        )
-        
-        # Get character name for seed
-        main_char_name = "default"
-        if characters:
-            first_char = characters[0]
-            main_char_name = first_char.get('name', 'default') if isinstance(first_char, dict) else getattr(first_char, 'name', 'default')
-        
-        # Use consistent seed
-        generator = torch.Generator(device="cpu")
-        main_char_seed = get_character_seed(story_title, main_char_name, sequence_number)
-        generator.manual_seed(main_char_seed)
-        
-        # Generate image
-        global current_pipe
-        image = current_pipe(
-            prompt=enhanced_prompt,
-            negative_prompt=negative_prompt,
-            num_inference_steps=20,
-            guidance_scale=7.0,
-            width=512,
-            height=512,
-            generator=generator
-        ).images[0]
-        
-        # Save to OCI
-        success, save_status = save_complete_storybook_page(image, story_title, sequence_number, scene_text)
-        
-        if success:
-            print(f"✅ Page {sequence_number} completed successfully")
-            return True, save_status
-        else:
-            print(f"❌ Page {sequence_number} save failed: {save_status}")
-            return False, save_status
-        
-    except Exception as e:
-        error_msg = f"❌ Page {sequence_number} generation failed: {str(e)}"
-        print(error_msg)
-        return False, error_msg
-
-# FastAPI endpoint - SYNCHRONOUS VERSION
-@app.post("/api/generate-storybook", response_model=StorybookResponse)
-async def api_generate_storybook(request: StorybookRequest):
-    """Synchronous API endpoint that actually works on Hugging Face"""
+def process_storybook_generation(request_id, request_data):
+    """Process generation in background and store results"""
     try:
-        print(f"📚 Received request: {request.story_title}")
-        print(f"📖 Pages: {len(request.scenes)}")
-        print(f"👤 Characters: {len(request.characters)}")
-        
-        start_time = time.time()
+        print(f"🔧 Processing request {request_id} in background...")
         
-        # Load model ONCE
-        load_model(request.model_choice)
+        # Load model
+        load_model(request_data["model_choice"])
         
         # Convert characters to dict
         characters_dict = []
-        for char in request.characters:
+        for char in request_data["characters"]:
             characters_dict.append({
-                "name": char.name,
-                "description": char.description
+                "name": char["name"],
+                "description": char["description"]
             })
         
-        generated_count = 0
+        results = []
         status_messages = []
+        start_time = time.time()
         
-        # Process each page SEQUENTIALLY
-        for i, scene in enumerate(request.scenes, 1):
+        # Process each page
+        for i, scene in enumerate(request_data["scenes"], 1):
             try:
-                success, message = generate_single_page(
-                    scene.visual,
-                    scene.text,
-                    request.story_title,
-                    i,
-                    characters_dict,
-                    request.model_choice,
-                    request.style
+                print(f"🔄 Generating page {i}...")
+                
+                enhanced_prompt, negative_prompt = enhance_prompt(
+                    scene["visual"], characters_dict, request_data["style"], i
+                )
+                
+                # Get character name for seed
+                main_char_name = "default"
+                if characters_dict:
+                    main_char_name = characters_dict[0]["name"]
+                
+                # Use consistent seed
+                generator = torch.Generator(device="cpu")
+                main_char_seed = get_character_seed(request_data["story_title"], main_char_name, i)
+                generator.manual_seed(main_char_seed)
+                
+                # Generate image
+                global current_pipe
+                image = current_pipe(
+                    prompt=enhanced_prompt,
+                    negative_prompt=negative_prompt,
+                    num_inference_steps=25,
+                    guidance_scale=7.0,
+                    width=512,
+                    height=512,
+                    generator=generator
+                ).images[0]
+                
+                # Save to OCI
+                success, save_status = save_complete_storybook_page(
+                    image, request_data["story_title"], i, scene["text"]
                 )
                 
                 if success:
-                    generated_count += 1
-                    status_messages.append(f"Page {i}: {message}")
+                    results.append({"page_number": i, "status": "success"})
+                    status_messages.append(f"Page {i}: {save_status}")
+                    print(f"✅ Page {i} completed")
                 else:
-                    status_messages.append(f"Page {i}: {message}")
+                    results.append({"page_number": i, "status": "error", "message": save_status})
+                    status_messages.append(f"Page {i}: {save_status}")
                 
-                # Clean memory after each page
                 cleanup_memory()
                 
-                # Add small delay between pages
-                if i < len(request.scenes):
+                if i < len(request_data["scenes"]):
                     time.sleep(1)
                     
             except Exception as e:
                 error_msg = f"Page {i} failed: {str(e)}"
+                results.append({"page_number": i, "status": "error", "message": error_msg})
                 status_messages.append(error_msg)
                 print(f"❌ {error_msg}")
         
         total_time = time.time() - start_time
         
-        # Create response
+        # Store results
+        active_requests[request_id] = {
+            "status": "completed",
+            "results": results,
+            "message": "\n".join(status_messages),
+            "generation_time": total_time,
+            "completed_at": datetime.now().isoformat()
+        }
+        
+        print(f"✅ Request {request_id} completed in {total_time:.2f} seconds")
+        
+    except Exception as e:
+        active_requests[request_id] = {
+            "status": "error",
+            "message": f"Processing failed: {str(e)}"
+        }
+        print(f"❌ Request {request_id} failed: {e}")
+
+# FastAPI endpoint - IMMEDIATE RESPONSE
+@app.post("/api/generate-storybook", response_model=StorybookResponse)
+async def api_generate_storybook(request: StorybookRequest):
+    """API endpoint that returns immediately"""
+    try:
+        print(f"📚 Received request: {request.story_title}")
+        print(f"📖 Pages: {len(request.scenes)}")
+        
+        # Create request ID
+        request_id = f"{request.story_title}_{int(time.time())}"
+        
+        # Convert to dict for background processing
+        request_data = {
+            "story_title": request.story_title,
+            "scenes": [{"visual": scene.visual, "text": scene.text} for scene in request.scenes],
+            "characters": [{"name": char.name, "description": char.description} for char in request.characters],
+            "model_choice": request.model_choice,
+            "style": request.style
+        }
+        
+        # Store initial request state
+        active_requests[request_id] = {
+            "status": "processing",
+            "started_at": datetime.now().isoformat(),
+            "total_pages": len(request.scenes)
+        }
+        
+        # Start background processing in a thread
+        import threading
+        thread = threading.Thread(
+            target=process_storybook_generation,
+            args=(request_id, request_data)
+        )
+        thread.daemon = True
+        thread.start()
+        
+        # IMMEDIATE RESPONSE to n8n
         response_data = {
-            "status": "success" if generated_count > 0 else "partial",
+            "status": "processing",
             "story_title": request.story_title,
             "total_pages": len(request.scenes),
             "characters_used": len(request.characters),
-            "generated_pages": generated_count,
-            "generation_time": round(total_time, 2),
-            "message": "\n".join(status_messages),
+            "generated_pages": 0,
+            "generation_time": 0,
+            "message": f"Generation started for {len(request.scenes)} pages. Request ID: {request_id}",
             "folder_path": f"storybook-library/stories/{request.story_title.replace(' ', '_')}/",
             "pages": [
                 {
@@ -358,19 +379,33 @@ async def api_generate_storybook(request: StorybookRequest):
                     "image_file": f"page_{i+1:03d}_{request.story_title.replace(' ', '_')}.png",
                     "text_file": f"page_{i+1:03d}_{request.story_title.replace(' ', '_')}.txt"
                 } for i in range(len(request.scenes))
-            ]
+            ],
+            "request_id": request_id
         }
         
-        print(f"✅ Generation completed in {total_time:.2f} seconds")
-        print(f"📊 Generated {generated_count}/{len(request.scenes)} pages")
-        
         return response_data
         
     except Exception as e:
-        error_msg = f"Storybook generation failed: {str(e)}"
+        error_msg = f"Request failed: {str(e)}"
         print(f"❌ {error_msg}")
         raise HTTPException(status_code=500, detail=error_msg)
 
+# Status check endpoint for n8n
+@app.get("/api/status/{request_id}")
+async def check_status(request_id: str):
+    """Check status of a generation request"""
+    if request_id not in active_requests:
+        return {"status": "not_found", "message": "Request ID not found"}
+    
+    request_data = active_requests[request_id]
+    return {
+        "status": request_data["status"],
+        "message": request_data.get("message", ""),
+        "generation_time": request_data.get("generation_time", 0),
+        "completed_at": request_data.get("completed_at", ""),
+        "total_pages": request_data.get("total_pages", 0)
+    }
+
 @app.get("/api/health")
 async def health_check():
     return {
@@ -379,7 +414,8 @@ async def health_check():
         "timestamp": datetime.now().isoformat(),
         "memory_usage_mb": monitor_memory(),
         "models_loaded": list(model_cache.keys()),
-        "current_model": current_model_name
+        "current_model": current_model_name,
+        "active_requests": len(active_requests)
     }
 
 # Simple Gradio interface
@@ -388,21 +424,32 @@ with gr.Blocks(title="Storybook Generator", theme="soft") as demo:
     
     with gr.Row():
         story_title = gr.Textbox(label="Story Title", value="Test Story")
-        prompt_input = gr.Textbox(label="Scene Description", lines=3, value="A beautiful sunset over mountains")
+        prompt_input = gr.Textbox(label="Scene Description", lines=3, value="A beautiful sunset over mountains with vibrant colors")
         generate_btn = gr.Button("Generate Test Page")
         output_image = gr.Image()
         status = gr.Textbox()
     
     def generate_test_page(prompt, title):
         try:
-            success, message = generate_single_page(
-                prompt, "", title, 1, [], "dreamshaper-8", "childrens_book"
-            )
-            if success:
-                # For demo, return a placeholder since we can't easily get the image back
-                return None, message
-            else:
-                return None, message
+            # Test with a simple generation
+            enhanced_prompt, negative_prompt = enhance_prompt(prompt, [], "childrens_book", 1)
+            
+            generator = torch.Generator(device="cpu")
+            generator.manual_seed(123)
+            
+            global current_pipe
+            image = current_pipe(
+                prompt=enhanced_prompt,
+                negative_prompt=negative_prompt,
+                num_inference_steps=20,
+                guidance_scale=7.0,
+                width=512,
+                height=512,
+                generator=generator
+            ).images[0]
+            
+            return image, f"✅ Generated: {enhanced_prompt}"
+            
         except Exception as e:
             return None, f"Error: {str(e)}"
     
@@ -412,7 +459,7 @@ with gr.Blocks(title="Storybook Generator", theme="soft") as demo:
         outputs=[output_image, status]
     )
 
-app = gr.mount_gradio_app(app, demo, path="/")
+app = gr.mount_grado_app(app, demo, path="/")
 
 if __name__ == "__main__":
     print("🚀 Starting Storybook Generator API...")