Update utils/keyframe_utils.py

utils/keyframe_utils.py

@@ -1,38 +1,26 @@
-import json
-import random
+import openai
 import os
+import json
+from pathlib import Path
 from diffusers import StableDiffusionPipeline
 import torch
-import openai
-from pathlib import Path
-
-# Load and cache the diffusion pipeline (only once)
-pipe = StableDiffusionPipeline.from_pretrained(
-    "CompVis/stable-diffusion-v1-4",
-    torch_dtype=torch.float16
-)
-pipe = pipe.to("cpu")
-
-
 
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
 # Global story context
 story_context_cn = "《博物馆的全能ACE》是一部拟人化博物馆文物与AI讲解助手互动的短片，讲述太阳人石刻在闭馆后的博物馆中，遇到了新来的AI助手博小翼，两者展开对话，AI展示了自己的多模态讲解能力与文化知识，最终被文物们认可，并一起展开智慧导览服务的故事。该片融合了文物拟人化、夜间博物馆奇妙氛围、科技感界面与中国地方文化元素，风格活泼、具未来感。"
 
-# Cache directory for prompts
+# Cache and log directories
 CACHE_DIR = Path("prompt_cache")
 CACHE_DIR.mkdir(exist_ok=True)
 LOG_PATH = Path("prompt_log.jsonl")
 
+pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
+pipe = pipe.to("cuda")
+
 def generate_keyframe_prompt(segment):
-    """
-    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)
@@ -41,14 +29,7 @@ def generate_keyframe_prompt(segment):
     speaker = segment.get("speaker", "")
     narration = segment.get("narration", "")
 
-    input_prompt = f"你是一个擅长视觉脚本设计的AI，请基于以下故事整体背景与分镜内容，帮我生成一个适合用于Stable Diffusion图像生成的英文提示词（image prompt），用于生成低分辨率草图风格的关键帧。请注意突出主要角色、镜头氛围、光影、构图、动作，避免复杂背景和细节。
-
-【整体故事背景】：\n{story_context_cn}
-
-【当前分镜描述】：\n{description}
-【角色】：{speaker}\n【台词或画外音】：{narration}
-
-请用英文输出一个简洁但具体的prompt，风格偏草图、线稿、卡通、简洁构图，并指出一个negative prompt。"
+    input_prompt = f"你是一个擅长视觉脚本设计的AI，请基于以下故事整体背景与分镜内容，帮我生成一个适合用于Stable Diffusion图像生成的英文提示词（image prompt），用于生成低分辨率草图风格的关键帧。请注意突出主要角色、镜头氛围、光影、构图、动作，避免复杂背景和细节。\n\n【整体故事背景】：\n{story_context_cn}\n\n【当前分镜描述】：\n{description}\n【角色】：{speaker}\n【台词或画外音】：{narration}\n\n请用英文输出一个简洁但具体的prompt，风格偏草图、线稿、卡通、简洁构图，并指出一个negative prompt。"
 
     try:
         response = openai.ChatCompletion.create(
@@ -60,7 +41,6 @@ def generate_keyframe_prompt(segment):
             temperature=0.7
         )
         output_text = response["choices"][0]["message"]["content"]
-
         if "Negative prompt:" in output_text:
             prompt, negative = output_text.split("Negative prompt:", 1)
         else:
@@ -70,17 +50,11 @@ def generate_keyframe_prompt(segment):
             "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 {
@@ -89,10 +63,6 @@ def generate_keyframe_prompt(segment):
         }
 
 def generate_all_keyframe_images(script_data, output_dir="keyframes"):
-    """
-    Generates 3 keyframe images per segment using Stable Diffusion,
-    stores them in the given output directory.
-    """
     os.makedirs(output_dir, exist_ok=True)
     keyframe_outputs = []
 
@@ -118,4 +88,7 @@ def generate_all_keyframe_images(script_data, output_dir="keyframes"):
 
         print(f"✓ Generated 3 images for Segment {segment_id}")
 
+    with open("all_prompts_output.json", "w", encoding="utf-8") as f:
+        json.dump(keyframe_outputs, f, ensure_ascii=False, indent=2)
+
     return keyframe_outputs