Update app.py

app.py

@@ -9,12 +9,11 @@ from datetime import datetime
 import re
 import time
 from typing import List, Optional
-from fastapi import FastAPI, HTTPException, BackgroundTasks
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 import gc
 import psutil
 import threading
-from concurrent.futures import ThreadPoolExecutor, as_completed
 
 # External OCI API URL
 OCI_API_BASE_URL = "https://yukee1992-oci-story-book.hf.space"
@@ -74,9 +73,6 @@ model_lock = threading.Lock()
 # Character consistency tracking
 character_seeds = {}
 
-# Thread pool for parallel processing
-executor = ThreadPoolExecutor(max_workers=2)
-
 def monitor_memory():
     try:
         process = psutil.Process()
@@ -129,19 +125,18 @@ print("🚀 Initializing Storybook Generator...")
 load_model("dreamshaper-8")
 print("✅ Model loaded and ready!")
 
-# CRITICAL: PROMPT OPTIMIZATION THAT ACTUALLY WORKS
+# PROMPT OPTIMIZATION
 def optimize_prompt(scene_visual, characters, style="childrens_book", page_number=1):
     """
-    Create a prompt that FITS within 77 tokens while preserving the ESSENCE
+    Create a prompt that FITS within 77 tokens
     """
-    # Extract ONLY the most critical information
+    # Extract character essence
     character_essence = ""
     if characters:
         char_descriptors = []
         for char in characters:
             desc = char.get('description', '') if isinstance(char, dict) else getattr(char, 'description', '')
             
-            # Extract ONLY: species + color + 1 key feature
             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"
@@ -149,7 +144,6 @@ def optimize_prompt(scene_visual, characters, style="childrens_book", page_numbe
             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 ""
             
-            # Find one key feature
             key_feature = ""
             if 'glasses' in desc.lower(): key_feature = "with glasses"
             elif 'dress' in desc.lower(): key_feature = "in dress"
@@ -157,16 +151,16 @@ def optimize_prompt(scene_visual, characters, style="childrens_book", page_numbe
             
             char_descriptors.append(f"{color} {species} {key_feature}".strip())
         
-        character_essence = f"Features: {', '.join(char_descriptors)}. "
+        character_essence = f"Characters: {', '.join(char_descriptors)}. "
     
-    # Compress scene description to MAX 40 words
+    # Compress scene description
     scene_words = scene_visual.split()
-    if len(scene_words) > 40:
-        scene_compressed = ' '.join(scene_words[:40])
+    if len(scene_words) > 30:
+        scene_compressed = ' '.join(scene_words[:30])
     else:
         scene_compressed = scene_visual
     
-    # Style context (very brief)
+    # Style context
     style_context = {
         "childrens_book": "children's book illustration",
         "realistic": "photorealistic",
@@ -174,14 +168,14 @@ def optimize_prompt(scene_visual, characters, style="childrens_book", page_numbe
         "anime": "anime style"
     }.get(style, "children's book illustration")
     
-    # Build the final prompt (GUARANTEED to fit 77 tokens)
+    # 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"
     
-    # Ensure it's under 60 words
+    # Ensure it's under 55 words for safety
     words = final_prompt.split()
-    if len(words) > 60:
-        final_prompt = ' '.join(words[:60])
+    if len(words) > 55:
+        final_prompt = ' '.join(words[:55])
     
     print(f"📝 Optimized prompt: {final_prompt}")
     print(f"📏 Length: {len(final_prompt.split())} words")
@@ -189,7 +183,7 @@ def optimize_prompt(scene_visual, characters, style="childrens_book", page_numbe
     return final_prompt
 
 def enhance_prompt(scene_visual, characters, style="childrens_book", page_number=1):
-    """Create optimized prompt that WILL work"""
+    """Create optimized prompt"""
     main_prompt = optimize_prompt(scene_visual, characters, style, page_number)
     
     negative_prompt = (
@@ -225,10 +219,10 @@ def save_complete_storybook_page(image, story_title, sequence_number, scene_text
         }
         text_response = requests.post(text_api_url, data=text_data, timeout=30)
         
-        return f"✅ Page {sequence_number}: Image & Text saved"
+        return True, f"✅ Page {sequence_number}: Image & Text saved"
             
     except Exception as e:
-        return f"❌ Save failed: {str(e)}"
+        return False, f"❌ Save failed: {str(e)}"
 
 def get_character_seed(story_title, character_name, page_number):
     if story_title not in character_seeds:
@@ -243,12 +237,9 @@ def get_character_seed(story_title, character_name, page_number):
     
     return character_seeds[story_title][seed_key]
 
-def generate_single_page(scene_data, story_title, sequence_number, characters, model_choice, style):
-    """Generate a single page - isolated for error handling"""
+def generate_single_page(scene_visual, scene_text, story_title, sequence_number, characters, model_choice, style):
+    """Generate a single page"""
     try:
-        scene_visual = scene_data.get('visual', '')
-        scene_text = scene_data.get('text', '')
-        
         print(f"🔄 Generating page {sequence_number}...")
         
         enhanced_prompt, negative_prompt = enhance_prompt(
@@ -266,67 +257,42 @@ def generate_single_page(scene_data, story_title, sequence_number, characters, m
         main_char_seed = get_character_seed(story_title, main_char_name, sequence_number)
         generator.manual_seed(main_char_seed)
         
-        # Generate with current pipe (already loaded)
+        # Generate image
         global current_pipe
         image = current_pipe(
             prompt=enhanced_prompt,
             negative_prompt=negative_prompt,
-            num_inference_steps=20,  # Faster generation
+            num_inference_steps=20,
             guidance_scale=7.0,
-            width=512,  # Smaller for speed
+            width=512,
             height=512,
             generator=generator
         ).images[0]
         
-        save_status = save_complete_storybook_page(image, story_title, sequence_number, scene_text)
+        # Save to OCI
+        success, save_status = save_complete_storybook_page(image, story_title, sequence_number, scene_text)
         
-        return {
-            "success": True,
-            "page_number": sequence_number,
-            "image": image,
-            "status": save_status
-        }
+        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:
-        return {
-            "success": False,
-            "page_number": sequence_number,
-            "error": f"Generation failed: {str(e)}"
-        }
+        error_msg = f"❌ Page {sequence_number} generation failed: {str(e)}"
+        print(error_msg)
+        return False, error_msg
 
-# FastAPI endpoint - OPTIMIZED
+# FastAPI endpoint - SYNCHRONOUS VERSION
 @app.post("/api/generate-storybook", response_model=StorybookResponse)
 async def api_generate_storybook(request: StorybookRequest):
-    """API endpoint that WON'T timeout"""
+    """Synchronous API endpoint that actually works on Hugging Face"""
     try:
         print(f"📚 Received request: {request.story_title}")
         print(f"📖 Pages: {len(request.scenes)}")
+        print(f"👤 Characters: {len(request.characters)}")
         
-        # IMMEDIATE response to n8n to prevent timeout
-        response_data = {
-            "status": "processing",
-            "story_title": request.story_title,
-            "total_pages": len(request.scenes),
-            "characters_used": len(request.characters),
-            "generated_pages": 0,
-            "generation_time": 0,
-            "message": "Processing started in background",
-            "folder_path": f"storybook-library/stories/{request.story_title.replace(' ', '_')}/",
-            "pages": []
-        }
-        
-        # Start background processing
-        background_tasks = BackgroundTasks()
-        background_tasks.add_task(process_storybook_background, request)
-        
-        return response_data
-        
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Request failed: {str(e)}")
-
-def process_storybook_background(request):
-    """Background processing to avoid timeouts"""
-    try:
         start_time = time.time()
         
         # Load model ONCE
@@ -340,14 +306,15 @@ def process_storybook_background(request):
                 "description": char.description
             })
         
-        results = []
+        generated_count = 0
         status_messages = []
         
-        # Process each page SEQUENTIALLY (better for memory)
+        # Process each page SEQUENTIALLY
         for i, scene in enumerate(request.scenes, 1):
             try:
-                result = generate_single_page(
-                    {"visual": scene.visual, "text": scene.text},
+                success, message = generate_single_page(
+                    scene.visual,
+                    scene.text,
                     request.story_title,
                     i,
                     characters_dict,
@@ -355,20 +322,18 @@ def process_storybook_background(request):
                     request.style
                 )
                 
-                if result["success"]:
-                    results.append(result)
-                    status_messages.append(f"Page {i}: {result['status']}")
-                    print(f"✅ Page {i} completed successfully")
+                if success:
+                    generated_count += 1
+                    status_messages.append(f"Page {i}: {message}")
                 else:
-                    status_messages.append(f"Page {i}: {result['error']}")
-                    print(f"❌ Page {i} failed: {result['error']}")
+                    status_messages.append(f"Page {i}: {message}")
                 
                 # Clean memory after each page
                 cleanup_memory()
                 
-                # Add small delay to prevent resource exhaustion
+                # Add small delay between pages
                 if i < len(request.scenes):
-                    time.sleep(2)
+                    time.sleep(1)
                     
             except Exception as e:
                 error_msg = f"Page {i} failed: {str(e)}"
@@ -376,10 +341,35 @@ def process_storybook_background(request):
                 print(f"❌ {error_msg}")
         
         total_time = time.time() - start_time
-        print(f"✅ Background processing completed in {total_time:.2f} seconds")
+        
+        # Create response
+        response_data = {
+            "status": "success" if generated_count > 0 else "partial",
+            "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),
+            "folder_path": f"storybook-library/stories/{request.story_title.replace(' ', '_')}/",
+            "pages": [
+                {
+                    "page_number": i+1,
+                    "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))
+            ]
+        }
+        
+        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:
-        print(f"❌ Background processing failed: {str(e)}")
+        error_msg = f"Storybook generation failed: {str(e)}"
+        print(f"❌ {error_msg}")
+        raise HTTPException(status_code=500, detail=error_msg)
 
 @app.get("/api/health")
 async def health_check():
@@ -397,16 +387,27 @@ with gr.Blocks(title="Storybook Generator", theme="soft") as demo:
     gr.Markdown("# 📚 Storybook Generator")
     
     with gr.Row():
-        story_title = gr.Textbox(label="Story Title")
-        prompt_input = gr.Textbox(label="Scene Description", lines=3)
-        generate_btn = gr.Button("Generate")
+        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")
+        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
+        except Exception as e:
+            return None, f"Error: {str(e)}"
+    
     generate_btn.click(
-        fn=lambda p, t: generate_single_page(
-            {"visual": p, "text": ""}, t, 1, [], "dreamshaper-8", "childrens_book"
-        ),
+        fn=generate_test_page,
         inputs=[prompt_input, story_title],
         outputs=[output_image, status]
     )