DFAGWE/infinitetalk2
0
1# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.2import numpy as np3import torch4import torch.nn.functional as F5import torchvision.transforms.functional as TF6from PIL import Image7 8 9class VaceImageProcessor(object):10 11 def __init__(self, downsample=None, seq_len=None):12 self.downsample = downsample13 self.seq_len = seq_len14 15 def _pillow_convert(self, image, cvt_type='RGB'):16 if image.mode != cvt_type:17 if image.mode == 'P':18 image = image.convert(f'{cvt_type}A')19 if image.mode == f'{cvt_type}A':20 bg = Image.new(21 cvt_type,22 size=(image.width, image.height),23 color=(255, 255, 255))24 bg.paste(image, (0, 0), mask=image)25 image = bg26 else:27 image = image.convert(cvt_type)28 return image29 30 def _load_image(self, img_path):31 if img_path is None or img_path == '':32 return None33 img = Image.open(img_path)34 img = self._pillow_convert(img)35 return img36 37 def _resize_crop(self, img, oh, ow, normalize=True):38 """39 Resize, center crop, convert to tensor, and normalize.40 """41 # resize and crop42 iw, ih = img.size43 if iw != ow or ih != oh:44 # resize45 scale = max(ow / iw, oh / ih)46 img = img.resize((round(scale * iw), round(scale * ih)),47 resample=Image.Resampling.LANCZOS)48 assert img.width >= ow and img.height >= oh49 50 # center crop51 x1 = (img.width - ow) // 252 y1 = (img.height - oh) // 253 img = img.crop((x1, y1, x1 + ow, y1 + oh))54 55 # normalize56 if normalize:57 img = TF.to_tensor(img).sub_(0.5).div_(0.5).unsqueeze(1)58 return img59 60 def _image_preprocess(self, img, oh, ow, normalize=True, **kwargs):61 return self._resize_crop(img, oh, ow, normalize)62 63 def load_image(self, data_key, **kwargs):64 return self.load_image_batch(data_key, **kwargs)65 66 def load_image_pair(self, data_key, data_key2, **kwargs):67 return self.load_image_batch(data_key, data_key2, **kwargs)68 69 def load_image_batch(self,70 *data_key_batch,71 normalize=True,72 seq_len=None,73 **kwargs):74 seq_len = self.seq_len if seq_len is None else seq_len75 imgs = []76 for data_key in data_key_batch:77 img = self._load_image(data_key)78 imgs.append(img)79 w, h = imgs[0].size80 dh, dw = self.downsample[1:]81 82 # compute output size83 scale = min(1., np.sqrt(seq_len / ((h / dh) * (w / dw))))84 oh = int(h * scale) // dh * dh85 ow = int(w * scale) // dw * dw86 assert (oh // dh) * (ow // dw) <= seq_len87 imgs = [self._image_preprocess(img, oh, ow, normalize) for img in imgs]88 return *imgs, (oh, ow)89 90 91class VaceVideoProcessor(object):92 93 def __init__(self, downsample, min_area, max_area, min_fps, max_fps,94 zero_start, seq_len, keep_last, **kwargs):95 self.downsample = downsample96 self.min_area = min_area97 self.max_area = max_area98 self.min_fps = min_fps99 self.max_fps = max_fps100 self.zero_start = zero_start101 self.keep_last = keep_last102 self.seq_len = seq_len103 assert seq_len >= min_area / (self.downsample[1] * self.downsample[2])104 105 def set_area(self, area):106 self.min_area = area107 self.max_area = area108 109 def set_seq_len(self, seq_len):110 self.seq_len = seq_len111 112 @staticmethod113 def resize_crop(video: torch.Tensor, oh: int, ow: int):114 """115 Resize, center crop and normalize for decord loaded video (torch.Tensor type)116 117 Parameters:118 video - video to process (torch.Tensor): Tensor from `reader.get_batch(frame_ids)`, in shape of (T, H, W, C)119 oh - target height (int)120 ow - target width (int)121 122 Returns:123 The processed video (torch.Tensor): Normalized tensor range [-1, 1], in shape of (C, T, H, W)124 125 Raises:126 """127 # permute ([t, h, w, c] -> [t, c, h, w])128 video = video.permute(0, 3, 1, 2)129 130 # resize and crop131 ih, iw = video.shape[2:]132 if ih != oh or iw != ow:133 # resize134 scale = max(ow / iw, oh / ih)135 video = F.interpolate(136 video,137 size=(round(scale * ih), round(scale * iw)),138 mode='bicubic',139 antialias=True)140 assert video.size(3) >= ow and video.size(2) >= oh141 142 # center crop143 x1 = (video.size(3) - ow) // 2144 y1 = (video.size(2) - oh) // 2145 video = video[:, :, y1:y1 + oh, x1:x1 + ow]146 147 # permute ([t, c, h, w] -> [c, t, h, w]) and normalize148 video = video.transpose(0, 1).float().div_(127.5).sub_(1.)149 return video150 151 def _video_preprocess(self, video, oh, ow):152 return self.resize_crop(video, oh, ow)153 154 def _get_frameid_bbox_default(self, fps, frame_timestamps, h, w, crop_box,155 rng):156 target_fps = min(fps, self.max_fps)157 duration = frame_timestamps[-1].mean()158 x1, x2, y1, y2 = [0, w, 0, h] if crop_box is None else crop_box159 h, w = y2 - y1, x2 - x1160 ratio = h / w161 df, dh, dw = self.downsample162 163 area_z = min(self.seq_len, self.max_area / (dh * dw),164 (h // dh) * (w // dw))165 of = min((int(duration * target_fps) - 1) // df + 1,166 int(self.seq_len / area_z))167 168 # deduce target shape of the [latent video]169 target_area_z = min(area_z, int(self.seq_len / of))170 oh = round(np.sqrt(target_area_z * ratio))171 ow = int(target_area_z / oh)172 of = (of - 1) * df + 1173 oh *= dh174 ow *= dw175 176 # sample frame ids177 target_duration = of / target_fps178 begin = 0. if self.zero_start else rng.uniform(179 0, duration - target_duration)180 timestamps = np.linspace(begin, begin + target_duration, of)181 frame_ids = np.argmax(182 np.logical_and(timestamps[:, None] >= frame_timestamps[None, :, 0],183 timestamps[:, None] < frame_timestamps[None, :, 1]),184 axis=1).tolist()185 return frame_ids, (x1, x2, y1, y2), (oh, ow), target_fps186 187 def _get_frameid_bbox_adjust_last(self, fps, frame_timestamps, h, w,188 crop_box, rng):189 duration = frame_timestamps[-1].mean()190 x1, x2, y1, y2 = [0, w, 0, h] if crop_box is None else crop_box191 h, w = y2 - y1, x2 - x1192 ratio = h / w193 df, dh, dw = self.downsample194 195 area_z = min(self.seq_len, self.max_area / (dh * dw),196 (h // dh) * (w // dw))197 of = min((len(frame_timestamps) - 1) // df + 1,198 int(self.seq_len / area_z))199 200 # deduce target shape of the [latent video]201 target_area_z = min(area_z, int(self.seq_len / of))202 oh = round(np.sqrt(target_area_z * ratio))203 ow = int(target_area_z / oh)204 of = (of - 1) * df + 1205 oh *= dh206 ow *= dw207 208 # sample frame ids209 target_duration = duration210 target_fps = of / target_duration211 timestamps = np.linspace(0., target_duration, of)212 frame_ids = np.argmax(213 np.logical_and(timestamps[:, None] >= frame_timestamps[None, :, 0],214 timestamps[:, None] <= frame_timestamps[None, :, 1]),215 axis=1).tolist()216 # print(oh, ow, of, target_duration, target_fps, len(frame_timestamps), len(frame_ids))217 return frame_ids, (x1, x2, y1, y2), (oh, ow), target_fps218 219 def _get_frameid_bbox(self, fps, frame_timestamps, h, w, crop_box, rng):220 if self.keep_last:221 return self._get_frameid_bbox_adjust_last(fps, frame_timestamps, h,222 w, crop_box, rng)223 else:224 return self._get_frameid_bbox_default(fps, frame_timestamps, h, w,225 crop_box, rng)226 227 def load_video(self, data_key, crop_box=None, seed=2024, **kwargs):228 return self.load_video_batch(229 data_key, crop_box=crop_box, seed=seed, **kwargs)230 231 def load_video_pair(self,232 data_key,233 data_key2,234 crop_box=None,235 seed=2024,236 **kwargs):237 return self.load_video_batch(238 data_key, data_key2, crop_box=crop_box, seed=seed, **kwargs)239 240 def load_video_batch(self,241 *data_key_batch,242 crop_box=None,243 seed=2024,244 **kwargs):245 rng = np.random.default_rng(seed + hash(data_key_batch[0]) % 10000)246 # read video247 import decord248 decord.bridge.set_bridge('torch')249 readers = []250 for data_k in data_key_batch:251 reader = decord.VideoReader(data_k)252 readers.append(reader)253 254 fps = readers[0].get_avg_fps()255 length = min([len(r) for r in readers])256 frame_timestamps = [257 readers[0].get_frame_timestamp(i) for i in range(length)258 ]259 frame_timestamps = np.array(frame_timestamps, dtype=np.float32)260 h, w = readers[0].next().shape[:2]261 frame_ids, (x1, x2, y1, y2), (oh, ow), fps = self._get_frameid_bbox(262 fps, frame_timestamps, h, w, crop_box, rng)263 264 # preprocess video265 videos = [266 reader.get_batch(frame_ids)[:, y1:y2, x1:x2, :]267 for reader in readers268 ]269 videos = [self._video_preprocess(video, oh, ow) for video in videos]270 return *videos, frame_ids, (oh, ow), fps271 # return videos if len(videos) > 1 else videos[0]272 273 274def prepare_source(src_video, src_mask, src_ref_images, num_frames, image_size,275 device):276 for i, (sub_src_video, sub_src_mask) in enumerate(zip(src_video, src_mask)):277 if sub_src_video is None and sub_src_mask is None:278 src_video[i] = torch.zeros(279 (3, num_frames, image_size[0], image_size[1]), device=device)280 src_mask[i] = torch.ones(281 (1, num_frames, image_size[0], image_size[1]), device=device)282 for i, ref_images in enumerate(src_ref_images):283 if ref_images is not None:284 for j, ref_img in enumerate(ref_images):285 if ref_img is not None and ref_img.shape[-2:] != image_size:286 canvas_height, canvas_width = image_size287 ref_height, ref_width = ref_img.shape[-2:]288 white_canvas = torch.ones(289 (3, 1, canvas_height, canvas_width),290 device=device) # [-1, 1]291 scale = min(canvas_height / ref_height,292 canvas_width / ref_width)293 new_height = int(ref_height * scale)294 new_width = int(ref_width * scale)295 resized_image = F.interpolate(296 ref_img.squeeze(1).unsqueeze(0),297 size=(new_height, new_width),298 mode='bilinear',299 align_corners=False).squeeze(0).unsqueeze(1)300 top = (canvas_height - new_height) // 2301 left = (canvas_width - new_width) // 2302 white_canvas[:, :, top:top + new_height,303 left:left + new_width] = resized_image304 src_ref_images[i][j] = white_canvas305 return src_video, src_mask, src_ref_images306