Update core/image_generator.py

core/image_generator.py

@@ -181,19 +181,9 @@ def pil_to_base64(img: Image.Image) -> str:
 # UNIFIED IMAGE GENERATION FUNCTION
 # --------------------------------------------------------------
 async def generate_images(prompt_or_json, seed: int | None = None, num_images: int = 3):
-    """
-    Universal entrypoint.
-    - If input is a string → generate list of images
-    - If input is a JSON (dict) → generate character & keyframe images and return updated JSON
-    """
     global pipe, img2img_pipe
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    if pipe is None:
-        pipe = load_pipeline()
-    if img2img_pipe is None:
-        img2img_pipe = load_img2img_pipeline()
-
     # ----------------------------------------------------------
     # CASE 1: STRUCTURED JSON (story mode)
     # ----------------------------------------------------------
@@ -201,7 +191,8 @@ async def generate_images(prompt_or_json, seed: int | None = None, num_images: i
         story_json = prompt_or_json
         print("[ImageGen] Detected structured JSON input. Generating cinematic visuals...")
 
-        # ---------------- Step 1: Character Images ----------------
+        # Step 1: Load only txt2img for character generation
+        pipe = load_pipeline()
         seed_to_char_image = {}
         for char in story_json.get("characters", []):
             char_name = char["name"]
@@ -218,10 +209,13 @@ async def generate_images(prompt_or_json, seed: int | None = None, num_images: i
                 image = pipe(f"{char_name}, {char_desc}", num_inference_steps=30, generator=generator).images[0]
                 image.save(seed_image_path)
 
-            # Map seed → character image
             seed_to_char_image[char_seed] = image
 
-        # ---------------- Step 2: Keyframe Images ----------------
+        # Free txt2img pipeline
+        unload_pipelines(target="pipe")
+
+        # Step 2: Load only img2img for keyframes
+        img2img_pipe = load_img2img_pipeline()
         for key, scene_data in story_json.items():
             if not key.startswith("scene"):
                 continue
@@ -246,13 +240,11 @@ async def generate_images(prompt_or_json, seed: int | None = None, num_images: i
                             generator=generator
                         ).images[0]
 
-                        # Save temporarily
                         out_path = TMP_DIR / f"{key}_{kf_key}_seed{frame_seed}.png"
                         img.save(out_path)
-
-                        # Replace prompt with actual base64 image
                         frame[kf_key] = pil_to_base64(img)
 
+        unload_pipelines(target="all")  # unload both just in case
         print("[ImageGen] ✅ Story JSON image generation complete.")
         return story_json
 
@@ -260,6 +252,7 @@ async def generate_images(prompt_or_json, seed: int | None = None, num_images: i
     # CASE 2: NORMAL PROMPT
     # ----------------------------------------------------------
     print(f"[ImageGen] Generating {num_images} image(s) for prompt='{prompt_or_json}' seed={seed}")
+    pipe = load_pipeline()
     images = []
     for i in range(num_images):
         gen = torch.Generator(device).manual_seed(seed + i) if seed is not None else None
@@ -271,5 +264,6 @@ async def generate_images(prompt_or_json, seed: int | None = None, num_images: i
         except Exception as e:
             print(f"[ImageGen] ⚠️ Failed on image {i}: {e}")
 
+    unload_pipelines(target="pipe")
     print(f"[ImageGen] ✅ Generated {len(images)} image(s) successfully.")
     return images