| 172 | |
| 173 | |
| 174 | def encode_image_condition(self, image, inter_images, num_frames, height, width,cam_id): |
| 175 | image = self.preprocess_image(image.resize((width, height))).to(self.device) |
| 176 | clip_context = self.image_encoder.encode_image([image]) |
| 177 | |
| 178 | msk = torch.ones(1, num_frames, height//8, width//8, device=self.device) |
| 179 | msk = torch.concat([torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]], dim=1) |
| 180 | msk = msk.view(1, msk.shape[1] // 4, 4, height//8, width//8) |
| 181 | msk = msk.transpose(1, 2)[0] |
| 182 | |
| 183 | vae_input = torch.concat([image.transpose(0, 1), torch.zeros(3, num_frames-1, height, width).to(image.device)], dim=1) |
| 184 | |
| 185 | for idx in range(1,len(inter_images)): |
| 186 | inter_image = self.preprocess_image(inter_images[idx].resize((width, height))).to(self.device) |
| 187 | vae_input[:,idx,...] = inter_image |
| 188 | y = self.vae.encode([vae_input.to(dtype=self.torch_dtype, device=self.device)], device=self.device)[0] |
| 189 | y = torch.concat([msk, y]) |
| 190 | y = y.unsqueeze(0) |
| 191 | clip_context = clip_context.to(dtype=self.torch_dtype, device=self.device) |
| 192 | y = y.to(dtype=self.torch_dtype, device=self.device) |
| 193 | return {"clip_feature": clip_context, "y": y} |
| 194 | |
| 195 | def tensor2video(self, frames): |
| 196 | frames = rearrange(frames, "C T H W -> T H W C") |