Update utils/keyframe_utils.py

utils/keyframe_utils.py

@@ -4,6 +4,7 @@ import os
 from diffusers import StableDiffusionPipeline
 import torch
 import openai
+from pathlib import Path
 
 # Load and cache the diffusion pipeline (only once)
 pipe = StableDiffusionPipeline.from_pretrained(
@@ -13,20 +14,32 @@ pipe = StableDiffusionPipeline.from_pretrained(
 pipe = pipe.to("cpu")
 
 
-openai.api_key = os.getenv("OPENAI_API_KEY")  # Make sure this is set in your environment
 
-# Global story context (in Chinese)
+openai.api_key = os.getenv("OPENAI_API_KEY")
+
+# Global story context
 story_context_cn = "《博物馆的全能ACE》是一部拟人化博物馆文物与AI讲解助手互动的短片，讲述太阳人石刻在闭馆后的博物馆中，遇到了新来的AI助手博小翼，两者展开对话，AI展示了自己的多模态讲解能力与文化知识，最终被文物们认可，并一起展开智慧导览服务的故事。该片融合了文物拟人化、夜间博物馆奇妙氛围、科技感界面与中国地方文化元素，风格活泼、具未来感。"
 
+# Cache directory for prompts
+CACHE_DIR = Path("prompt_cache")
+CACHE_DIR.mkdir(exist_ok=True)
+LOG_PATH = Path("prompt_log.jsonl")
+
 def generate_keyframe_prompt(segment):
     """
-    Calls GPT-4o to generate an image prompt optimized for Stable Diffusion,
-    based on segment content and full story context.
+    Generates and caches image prompts using GPT-4o for a given segment.
     """
+    segment_id = segment.get("segment_id")
+    cache_file = CACHE_DIR / f"segment_{segment_id}.json"
+
+    # Return cached result if exists
+    if cache_file.exists():
+        with open(cache_file, "r", encoding="utf-8") as f:
+            return json.load(f)
+
     description = segment.get("description", "")
     speaker = segment.get("speaker", "")
     narration = segment.get("narration", "")
-    segment_id = segment.get("segment_id")
 
     input_prompt = f"你是一个擅长视觉脚本设计的AI，请基于以下故事整体背景与分镜内容，帮我生成一个适合用于Stable Diffusion图像生成的英文提示词（image prompt），用于生成低分辨率草图风格的关键帧。请注意突出主要角色、镜头氛围、光影、构图、动作，避免复杂背景和细节。
 
@@ -48,16 +61,26 @@ def generate_keyframe_prompt(segment):
         )
         output_text = response["choices"][0]["message"]["content"]
 
-        # Split response into prompt + negative if possible
         if "Negative prompt:" in output_text:
             prompt, negative = output_text.split("Negative prompt:", 1)
         else:
             prompt, negative = output_text, "blurry, distorted, low quality, text, watermark"
 
-        return {
+        result = {
             "prompt": prompt.strip(),
             "negative_prompt": negative.strip()
         }
+
+        # Save to cache
+        with open(cache_file, "w", encoding="utf-8") as f:
+            json.dump(result, f, ensure_ascii=False, indent=2)
+
+        # Log to JSONL for review
+        with open(LOG_PATH, "a", encoding="utf-8") as logf:
+            logf.write(json.dumps({"segment_id": segment_id, **result}, ensure_ascii=False) + "\n")
+
+        return result
+
     except Exception as e:
         print(f"[Error] GPT-4o prompt generation failed for segment {segment_id}: {e}")
         return {
@@ -65,7 +88,6 @@ def generate_keyframe_prompt(segment):
             "negative_prompt": ""
         }
 
-
 def generate_all_keyframe_images(script_data, output_dir="keyframes"):
     """
     Generates 3 keyframe images per segment using Stable Diffusion,