File size: 6,050 Bytes
cac3654 |
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 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://github.com/Tencent-Hunyuan/HunyuanImage-3.0/blob/main/LICENSE
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from typing import Tuple
from PIL import Image
from torchvision import transforms
from transformers import Siglip2ImageProcessorFast
from .tokenizer_wrapper import ImageInfo, JointImageInfo, ResolutionGroup
def resize_and_crop(image: Image.Image, target_size: Tuple[int, int]) -> Image.Image:
tw, th = target_size
w, h = image.size
tr = th / tw
r = h / w
# resize
if r < tr:
resize_height = th
resize_width = int(round(th / h * w))
else:
resize_width = tw
resize_height = int(round(tw / w * h))
image = image.resize((resize_width, resize_height), resample=Image.Resampling.LANCZOS)
# center crop
crop_top = int(round((resize_height - th) / 2.0))
crop_left = int(round((resize_width - tw) / 2.0))
image = image.crop((crop_left, crop_top, crop_left + tw, crop_top + th))
return image
class HunyuanImage3ImageProcessor(object):
def __init__(self, config):
self.config = config
self.reso_group = ResolutionGroup(base_size=config.image_base_size)
self.vae_processor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]), # transform to [-1, 1]
])
self.vision_encoder_processor = Siglip2ImageProcessorFast.from_dict(config.vit_processor)
def build_image_info(self, image_size):
# parse image size (HxW, H:W, or <img_ratio_i>)
if isinstance(image_size, str):
if image_size.startswith("<img_ratio_"):
ratio_index = int(image_size.split("_")[-1].rstrip(">"))
reso = self.reso_group[ratio_index]
image_size = reso.height, reso.width
elif 'x' in image_size:
image_size = [int(s) for s in image_size.split('x')]
elif ':' in image_size:
image_size = [int(s) for s in image_size.split(':')]
else:
raise ValueError(
f"`image_size` should be in the format of 'HxW', 'H:W' or <img_ratio_i>, got {image_size}.")
assert len(image_size) == 2, f"`image_size` should be in the format of 'HxW', got {image_size}."
elif isinstance(image_size, (list, tuple)):
assert len(image_size) == 2 and all(isinstance(s, int) for s in image_size), \
f"`image_size` should be a tuple of two integers or a string in the format of 'HxW', got {image_size}."
else:
raise ValueError(f"`image_size` should be a tuple of two integers or a string in the format of 'WxH', "
f"got {image_size}.")
image_width, image_height = self.reso_group.get_target_size(image_size[1], image_size[0])
token_height = image_height // (self.config.vae_downsample_factor[0] * self.config.patch_size)
token_width = image_width // (self.config.vae_downsample_factor[1] * self.config.patch_size)
base_size, ratio_idx = self.reso_group.get_base_size_and_ratio_index(image_size[1], image_size[0])
image_info = ImageInfo(
image_type="gen_image", image_width=image_width, image_height=image_height,
token_width=token_width, token_height=token_height, base_size=base_size, ratio_index=ratio_idx,
)
return image_info
def preprocess(self, image: Image.Image):
# ==== VAE processor ====
image_width, image_height = self.reso_group.get_target_size(image.width, image.height)
resized_image = resize_and_crop(image, (image_width, image_height))
image_tensor = self.vae_processor(resized_image)
token_height = image_height // (self.config.vae_downsample_factor[0] * self.config.patch_size)
token_width = image_width // (self.config.vae_downsample_factor[1] * self.config.patch_size)
base_size, ratio_index = self.reso_group.get_base_size_and_ratio_index(width=image_width, height=image_height)
vae_image_info = ImageInfo(
image_type="vae",
image_tensor=image_tensor.unsqueeze(0), # include batch dim
image_width=image_width, image_height=image_height,
token_width=token_width, token_height=token_height,
base_size=base_size, ratio_index=ratio_index,
)
# ==== ViT processor ====
inputs = self.vision_encoder_processor(image)
image = inputs["pixel_values"].squeeze(0) # seq_len x dim
pixel_attention_mask = inputs["pixel_attention_mask"].squeeze(0) # seq_len
spatial_shapes = inputs["spatial_shapes"].squeeze(0) # 2 (h, w)
vision_encoder_kwargs = dict(
pixel_attention_mask=pixel_attention_mask,
spatial_shapes=spatial_shapes,
)
vision_image_info = ImageInfo(
image_type="vit",
image_tensor=image.unsqueeze(0), # 1 x seq_len x dim
image_width=spatial_shapes[1].item() * self.config.vit_processor["patch_size"],
image_height=spatial_shapes[0].item() * self.config.vit_processor["patch_size"],
token_width=spatial_shapes[1].item(),
token_height=spatial_shapes[0].item(),
image_token_length=self.config.vit_processor["max_num_patches"],
# may not equal to token_width * token_height
)
return JointImageInfo(vae_image_info, vision_image_info, vision_encoder_kwargs)
|