File size: 2,665 Bytes
4fe31ed | 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 | """Inference-time utilities for HY-Unified.
Self-contained slice of the pieces the inference scripts need — the image
transform used to derive train/inference resolution parity, and the
task-instruction prompts that cue each generation mode.
The SAME task strings are used at train and inference (parity is load-bearing:
``interleave_inference`` / ``multiframe_inference`` build the prompt the same
way the model was trained). Single source of truth, imported by the scripts.
"""
import torch
from PIL import Image
import torchvision.transforms as T
import torchvision.transforms.functional as TF
from torchvision.transforms import InterpolationMode
# Task-instruction prompts appended to the user turn so each generation MODE is
# explicitly cued.
TASK_INSTRUCTION_SINGLE_FRAME = "Generate future image of the goal"
TASK_INSTRUCTION_MULTI_FRAME = "Generate future video of the goal"
TASK_INSTRUCTION_INTERLEAVE = "Generate interleave goal planning"
class VAEImageTransform:
"""Aspect-preserving resize + normalize to [-1, 1] for VAE encoding.
Sides are clamped to be ``stride``-divisible and within ``[min_size,
max_size]``; extreme aspect ratios are re-scaled to stay inside ``max_size``
(otherwise latent position IDs could go out of bounds).
"""
def __init__(self, max_size=512, min_size=256, stride=16):
self.max_size = max_size
self.min_size = min_size
self.stride = stride
def _make_divisible(self, val):
return max(self.stride, round(val / self.stride) * self.stride)
def __call__(self, img: Image.Image) -> torch.Tensor:
w, h = img.size
scale = min(self.max_size / max(w, h), 1.0)
scale = max(scale, self.min_size / min(w, h))
new_w = self._make_divisible(round(w * scale))
new_h = self._make_divisible(round(h * scale))
# Clamp to max_size to handle extreme aspect ratios where the min_size
# constraint overrides the max_size constraint (e.g. w=1000,h=50 would
# produce new_w=5120 which causes position ID out-of-bounds).
max_div = self._make_divisible(self.max_size)
if new_w > max_div or new_h > max_div:
# Re-scale to fit within max_size while preserving aspect ratio.
shrink = min(max_div / new_w, max_div / new_h)
new_w = self._make_divisible(round(new_w * shrink))
new_h = self._make_divisible(round(new_h * shrink))
img = TF.resize(img, (new_h, new_w), InterpolationMode.BICUBIC, antialias=True)
tensor = TF.to_tensor(img) # [0, 1]
tensor = T.Normalize([0.5] * 3, [0.5] * 3)(tensor) # [-1, 1]
return tensor
|