File size: 3,408 Bytes
101858b ad2ce18 101858b ad2ce18 101858b ad2ce18 101858b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | from __future__ import annotations
import hashlib
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from PIL import Image
from .storyboard import VideoStoryboard
@dataclass
class KeyframeResult:
frames: list[Image.Image]
paths: list[Path]
cache_hits: int
metadata: list[dict[str, Any]]
def cache_key(payload: dict[str, Any]) -> str:
raw = json.dumps(payload, sort_keys=True, ensure_ascii=True)
return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:24]
def _json_safe(value: Any) -> Any:
if isinstance(value, (str, int, float, bool)) or value is None:
return value
if isinstance(value, Path):
return str(value)
if isinstance(value, (list, tuple)):
return [_json_safe(item) for item in value]
if isinstance(value, dict):
return {str(key): _json_safe(item) for key, item in sorted(value.items(), key=lambda pair: str(pair[0]))}
return repr(value)
def generate_keyframes(
router: Any,
storyboard: VideoStoryboard,
*,
cache_dir: str | Path,
width: int,
height: int,
image_steps: int,
guidance_scale: float,
seed: int,
motion: str,
version: str,
use_cache: bool = True,
**kwargs: Any,
) -> KeyframeResult:
root = Path(cache_dir)
root.mkdir(parents=True, exist_ok=True)
generation_strategy = kwargs.pop("generation_strategy", "diffusion")
frames: list[Image.Image] = []
paths: list[Path] = []
metadata: list[dict[str, Any]] = []
cache_hits = 0
for idx, prompt in enumerate(storyboard.keyframe_prompts):
image_kwargs = {
"use_memory": False,
"reference_pass_steps": 0,
"unload_after_call": False,
**kwargs,
}
payload = {
"prompt": prompt,
"negative": storyboard.negative_prompt,
"width": width,
"height": height,
"steps": image_steps,
"guidance_scale": guidance_scale,
"seed": seed + idx,
"motion": motion,
"videogen_version": version,
"image_options": {
"generation_strategy": generation_strategy,
**{str(key): _json_safe(value) for key, value in sorted(image_kwargs.items(), key=lambda pair: str(pair[0]))},
},
}
path = root / f"{cache_key(payload)}.png"
if use_cache and path.exists():
frame = Image.open(path).convert("RGB")
cache_hits += 1
route_meta = {"cache_hit": True, "saved_path": str(path)}
else:
result = router.generate_image(
prompt=prompt,
height=height,
width=width,
steps=image_steps,
guidance_scale=guidance_scale,
seed=seed + idx,
generation_strategy=generation_strategy,
**image_kwargs,
)
frame = result.payload.images[0].convert("RGB")
frame.save(path)
route_meta = dict(getattr(result, "metadata", {}))
route_meta["cache_hit"] = False
route_meta["saved_path"] = str(path)
frames.append(frame)
paths.append(path)
metadata.append(route_meta)
return KeyframeResult(frames=frames, paths=paths, cache_hits=cache_hits, metadata=metadata)
|