File size: 2,113 Bytes
a64dea9 | 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 | """Stage 3 — keyframes. Two backends:
hunyuan-frame : text-to-video for N frames, steal frame 0. No extra weights.
flux / sdxl : a real image model, resolved from DWCT_IMAGE_MODEL."""
from __future__ import annotations
import logging
from pathlib import Path
from . import config, models
LOG = logging.getLogger("dwct.keyframes")
def _apply_palette(prompt: str, palette: list[str]) -> str:
if not palette:
return prompt
return f"{prompt}, colour palette {' '.join(palette)}"
def _via_hunyuan(prompt: str, negative: str, out: Path, seed: int) -> Path:
import torch
from PIL import Image
pipe = models.hunyuan()
gen = torch.Generator(device="cpu").manual_seed(seed)
common = dict(prompt=prompt, negative_prompt=negative,
num_inference_steps=config.STEPS, generator=gen,
height=config.HEIGHT, width=config.WIDTH)
try:
frames = pipe(num_frames=5, **common).frames[0]
except TypeError:
frames = pipe(**common).frames[0]
Image.fromarray(frames[0]).save(out)
return out
def _via_image_model(prompt: str, negative: str, out: Path, seed: int) -> Path:
import torch
pipe = models.image_model()
if pipe is None:
raise RuntimeError("DWCT_IMAGE_MODEL is unset")
gen = torch.Generator(device="cpu").manual_seed(seed)
img = pipe(prompt=prompt, negative_prompt=negative,
num_inference_steps=config.STEPS, guidance_scale=3.5,
height=config.HEIGHT, width=config.WIDTH, generator=gen).images[0]
img.save(out)
return out
def make(prompt: str, negative: str, out: Path, seed: int,
palette: list[str]) -> Path:
out.parent.mkdir(parents=True, exist_ok=True)
backend = config.KEYFRAME_BACKEND
if backend == "auto":
backend = "flux" if config.IMAGE_MODEL else "hunyuan-frame"
p = _apply_palette(prompt, palette)
LOG.info("keyframe -> %s via %s", out.name, backend)
if backend == "hunyuan-frame":
return _via_hunyuan(prompt, negative, out, seed)
return _via_image_model(p, negative, out, seed)
|