Spaces:
Sleeping
Sleeping
| 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 | |