Update app.py

app.py

@@ -60,7 +60,7 @@ MODEL_CHOICES = {
 # GLOBAL MODEL CACHE - Load once, reuse forever
 model_cache = {}
 current_model_name = None
-pipe = None
+current_pipe = None
 
 # Character consistency tracking
 character_descriptions = {}
@@ -68,12 +68,14 @@ 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
+    global model_cache, current_model_name, current_pipe
     
     # Return cached model if already loaded
     if model_name in model_cache:
         print(f"✅ Using cached model: {model_name}")
-        return model_cache[model_name]
+        current_pipe = model_cache[model_name]
+        current_model_name = model_name
+        return current_pipe
     
     print(f"🔄 Loading model for the first time: {model_name}")
     try:
@@ -92,6 +94,7 @@ def load_model(model_name="dreamshaper-8"):
         
         # Cache the model for future use
         model_cache[model_name] = pipe
+        current_pipe = pipe
         current_model_name = model_name
         
         print(f"✅ Model loaded and cached: {model_name}")
@@ -100,16 +103,19 @@ def load_model(model_name="dreamshaper-8"):
     except Exception as e:
         print(f"❌ Model loading failed: {e}")
         # Fallback to SD 1.5
-        return StableDiffusionPipeline.from_pretrained(
+        pipe = StableDiffusionPipeline.from_pretrained(
             "runwayml/stable-diffusion-v1-5", 
             torch_dtype=torch.float32,
             safety_checker=None,
             requires_safety_checker=False
         ).to("cpu")
+        model_cache[model_name] = pipe
+        current_pipe = pipe
+        return pipe
 
 # Load the default model once at startup
 print("🚀 Initializing Storybook Generator...")
-pipe = load_model("dreamshaper-8")
+current_pipe = load_model("dreamshaper-8")
 print("✅ Default model loaded and ready!")
 
 # PROFESSIONAL PROMPT ENGINEERING
@@ -220,13 +226,13 @@ def get_character_seed(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
+    global current_pipe, current_model_name
     
     try:
-        # Switch model if different from current
+        # Switch model if different from current - BUT DON'T RELOAD UNLESS NECESSARY
         if model_choice != current_model_name:
             print(f"🔄 Switching to model: {model_choice}")
-            pipe = load_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)
@@ -247,11 +253,14 @@ def generate_storybook_page(scene_visual, story_title, sequence_number, scene_te
             # Use seed from main character for consistency
             main_char_seed = get_character_seed(story_title, characters[0].name)
             generator.manual_seed(main_char_seed)
+            print(f"🌱 Using seed {main_char_seed} for character consistency")
         else:
-            generator.manual_seed(int(time.time()))
+            seed = int(time.time())
+            generator.manual_seed(seed)
+            print(f"🌱 Using timestamp seed {seed}")
         
-        # Generate high-quality image
-        image = pipe(
+        # Generate high-quality image - USE THE GLOBAL current_pipe
+        image = current_pipe(
             prompt=enhanced_prompt,
             negative_prompt=negative_prompt,
             num_inference_steps=30,
@@ -271,7 +280,7 @@ def generate_storybook_page(scene_visual, story_title, sequence_number, scene_te
 
 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
+    global character_descriptions, current_pipe
     
     results = []
     status_messages = []
@@ -286,9 +295,9 @@ def batch_generate_complete_storybook(story_title, scenes_data, characters, mode
         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)
+    # Load model once at the beginning - THIS IS THE KEY FIX
+    print(f"🔧 Loading model for this storybook...")
+    current_pipe = load_model(model_choice)
     
     start_time = time.time()
     
@@ -324,6 +333,8 @@ async def api_generate_storybook(request: StorybookRequest):
             for char in request.characters:
                 print(f"   - {char.name}: {char.description[:50]}...")
         
+        start_time = time.time()
+        
         # Convert to scene data format
         scenes_data = [{"visual": scene.visual, "text": scene.text} for scene in request.scenes]
         
@@ -336,12 +347,15 @@ async def api_generate_storybook(request: StorybookRequest):
             request.style
         )
         
+        generation_time = time.time() - start_time
+        
         return {
             "status": "success",
             "story_title": request.story_title,
             "total_pages": len(request.scenes),
             "characters_used": len(request.characters),
             "generated_pages": len(results),
+            "generation_time": round(generation_time, 2),
             "message": status,
             "folder_path": f"storybook-library/stories/{request.story_title.replace(' ', '_')}/",
             "pages": [
@@ -379,6 +393,11 @@ def generate_single_page(prompt, story_title, scene_text, model_choice, style):
     if not prompt or not story_title:
         return None, "❌ Please enter both scene description and story title"
     
+    # Ensure model is loaded
+    global current_pipe
+    if current_model_name != model_choice:
+        current_pipe = load_model(model_choice)
+    
     image, status = generate_storybook_page(
         prompt, story_title, 1, scene_text or "", [], model_choice, style
     )