| 2198 | DESCRIPTION = "Encodes a batch of images individually to create a latent video batch where each video is a single frame, useful for I2V init purposes, for example as multiple context window inits" |
| 2199 | |
| 2200 | def encode(self, vae, images, enable_vae_tiling=False, tile_x=272, tile_y=272, tile_stride_x=144, tile_stride_y=128, latent_strength=1.0): |
| 2201 | vae.to(device) |
| 2202 | |
| 2203 | images = images.clone() |
| 2204 | |
| 2205 | B, H, W, C = images.shape |
| 2206 | if W % 16 != 0 or H % 16 != 0: |
| 2207 | new_height = (H // 16) * 16 |
| 2208 | new_width = (W // 16) * 16 |
| 2209 | log.warning(f"Image size {W}x{H} is not divisible by 16, resizing to {new_width}x{new_height}") |
| 2210 | images = common_upscale(images.movedim(-1, 1), new_width, new_height, "lanczos", "disabled").movedim(1, -1) |
| 2211 | |
| 2212 | if images.shape[-1] == 4: |
| 2213 | images = images[..., :3] |
| 2214 | images = images.to(vae.dtype).to(device) * 2.0 - 1.0 |
| 2215 | |
| 2216 | latent_list = [] |
| 2217 | for img in images: |
| 2218 | if enable_vae_tiling and tile_x is not None: |
| 2219 | latent = vae.encode(img.unsqueeze(0).unsqueeze(0).permute(0, 4, 1, 2, 3), device=device, tiled=enable_vae_tiling, tile_size=(tile_x//vae.upsampling_factor, tile_y//vae.upsampling_factor), tile_stride=(tile_stride_x//vae.upsampling_factor, tile_stride_y//vae.upsampling_factor)) |
| 2220 | else: |
| 2221 | latent = vae.encode(img.unsqueeze(0).unsqueeze(0).permute(0, 4, 1, 2, 3), device=device, tiled=enable_vae_tiling) |
| 2222 | |
| 2223 | if latent_strength != 1.0: |
| 2224 | latent *= latent_strength |
| 2225 | latent_list.append(latent.squeeze(0).cpu()) |
| 2226 | latents_out = torch.stack(latent_list, dim=0) |
| 2227 | |
| 2228 | log.info(f"WanVideoEncode: Encoded latents shape {latents_out.shape}") |
| 2229 | vae.to(offload_device) |
| 2230 | mm.soft_empty_cache() |
| 2231 | |
| 2232 | return ({"samples": latents_out},) |
| 2233 | |
| 2234 | class WanVideoEncode: |
| 2235 | @classmethod |