| """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) |
|
|
