Update utils/keyframe_utils.py

utils/keyframe_utils.py

@@ -3,6 +3,7 @@ import random
 import os
 from diffusers import StableDiffusionPipeline
 import torch
+import openai
 
 # Load and cache the diffusion pipeline (only once)
 pipe = StableDiffusionPipeline.from_pretrained(
@@ -12,37 +13,57 @@ 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)
+story_context_cn = "《博物馆的全能ACE》是一部拟人化博物馆文物与AI讲解助手互动的短片，讲述太阳人石刻在闭馆后的博物馆中，遇到了新来的AI助手博小翼，两者展开对话，AI展示了自己的多模态讲解能力与文化知识，最终被文物们认可，并一起展开智慧导览服务的故事。该片融合了文物拟人化、夜间博物馆奇妙氛围、科技感界面与中国地方文化元素，风格活泼、具未来感。"
+
 def generate_keyframe_prompt(segment):
     """
-    Generates a detailed prompt optimized for Stable Diffusion (low-resolution, preview style)
-    based on the segment description.
+    Calls GPT-4o to generate an image prompt optimized for Stable Diffusion,
+    based on segment content and full story context.
     """
     description = segment.get("description", "")
     speaker = segment.get("speaker", "")
     narration = segment.get("narration", "")
     segment_id = segment.get("segment_id")
 
-    prompt_parts = []
-
-    if description:
-        prompt_parts.append(f"Scene: {description}.")
-
-    if speaker and narration:
-        prompt_parts.append(f"Character '{speaker}' speaking: \"{narration}\".")
-    elif narration:
-        prompt_parts.append(f"Narration: \"{narration}\".")
-
-    prompt_parts.append("Style: Simple, cartoonish, line art, sketch, low detail, illustrative, minimal background, focus on main subject.")
-    prompt_parts.append("Resolution: lowres, 256x256.")
-    prompt_parts.append("Lighting: Nighttime museum, dim lighting.")
-    prompt_parts.append("Setting: Museum interior, exhibits.")
-
-    negative_prompt = "blurry, distorted, ugly, tiling, poorly drawn, out of frame, disfigured, deformed, bad anatomy, watermark, text, signature, high detail, realistic, photorealistic, complex"
-
-    return {
-        "prompt": " ".join(prompt_parts).strip(),
-        "negative_prompt": negative_prompt
-    }
+    input_prompt = f"你是一个擅长视觉脚本设计的AI，请基于以下故事整体背景与分镜内容，帮我生成一个适合用于Stable Diffusion图像生成的英文提示词（image prompt），用于生成低分辨率草图风格的关键帧。请注意突出主要角色、镜头氛围、光影、构图、动作，避免复杂背景和细节。
+
+【整体故事背景】：\n{story_context_cn}
+
+【当前分镜描述】：\n{description}
+【角色】：{speaker}\n【台词或画外音】：{narration}
+
+请用英文输出一个简洁但具体的prompt，风格偏草图、线稿、卡通、简洁构图，并指出一个negative prompt。"
+
+    try:
+        response = openai.ChatCompletion.create(
+            model="gpt-4o",
+            messages=[
+                {"role": "system", "content": "You are an expert visual prompt designer for image generation."},
+                {"role": "user", "content": input_prompt}
+            ],
+            temperature=0.7
+        )
+        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 {
+            "prompt": prompt.strip(),
+            "negative_prompt": negative.strip()
+        }
+    except Exception as e:
+        print(f"[Error] GPT-4o prompt generation failed for segment {segment_id}: {e}")
+        return {
+            "prompt": description,
+            "negative_prompt": ""
+        }
 
 
 def generate_all_keyframe_images(script_data, output_dir="keyframes"):