Upload keyframe_utils.py

utils/keyframe_utils.py

@@ -0,0 +1,78 @@
+import json
+import random
+import os
+from diffusers import StableDiffusionPipeline
+import torch
+
+# Load and cache the diffusion pipeline (only once)
+pipe = StableDiffusionPipeline.from_pretrained(
+    "CompVis/stable-diffusion-v1-4",
+    torch_dtype=torch.float16
+)
+pipe = pipe.to("cuda")
+
+
+def generate_keyframe_prompt(segment):
+    """
+    Generates a detailed prompt optimized for Stable Diffusion (low-resolution, preview style)
+    based on the segment description.
+    """
+    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
+    }
+
+
+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 = []
+
+    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")
+
+        frame_images = []
+        for i in range(3):
+            image = pipe(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}")
+
+    return keyframe_outputs