Update utils/keyframe_utils.py

utils/keyframe_utils.py

@@ -72,3 +72,45 @@ def generate_keyframe_prompt(segment):
             "prompt": description,
             "negative_prompt": ""
         }
+
+def generate_all_keyframe_images(script_data, output_dir="keyframes"):
+    os.makedirs(output_dir, exist_ok=True)
+    keyframe_outputs = []
+
+    for segment in script_data:
+        sd_prompts = generate_keyframe_prompt(segment)
+        prompt = sd_prompts["prompt"]
+        negative_prompt = sd_prompts["negative_prompt"]
+        segment_id = segment.get("segment_id")
+
+        description = segment.get("description", "")
+        use_reference = any(name in description for name in ASSET_IMAGES)
+
+        if use_reference:
+            ref_key = next(k for k in ASSET_IMAGES if k in description)
+            init_image = Image.open(ASSET_IMAGES[ref_key]).convert("RGB").resize((512, 512))
+
+        frame_images = []
+        for i in range(3):
+            if use_reference:
+                image = pipe_img2img(prompt=prompt, image=init_image, negative_prompt=negative_prompt, strength=0.6, guidance_scale=7.5).images[0]
+            else:
+                image = pipe_txt2img(prompt, negative_prompt=negative_prompt, num_inference_steps=20, guidance_scale=7.5, height=256, width=256).images[0]
+
+            image_path = os.path.join(output_dir, f"segment_{segment_id}_v{i+1}.png")
+            image.save(image_path)
+            frame_images.append(image_path)
+
+        keyframe_outputs.append({
+            "segment_id": segment_id,
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "frame_images": frame_images
+        })
+
+        print(f"✓ Generated 3 images for Segment {segment_id} ({'img2img' if use_reference else 'txt2img'})")
+
+    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