Ref2VA/video_vae/klvae.py
47.5 KB · 1259 lines · python Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # MiniMax H3 visual VAE: 3D causal CNN encoder + ViT3D decoder (inference-only bundle).
3 import os
4 import math
5 import numpy as np
6 import torch
7 import torch.nn as nn
8 import torch.distributed as dist
9 from typing import List, Union
10 from PIL import Image
11 from contextlib import nullcontext
12 from diffusers.models import ModelMixin
13 from diffusers.configuration_utils import ConfigMixin, register_to_config
14 from diffusers.loaders.single_file_model import FromOriginalModelMixin
15 from diffusers.utils import logging
16
17 from .parallel import get_parallel_state, all_gather_var_shape
18 from .utils import apply_spatial_parallel
19 from .normalize import get_normalize_transform, get_denormalize_transform
20 from .vae_vit import ViT3DDecoder
21 from .vae_cnn import EncoderFCN3D
22 from .vae_module import DiagonalGaussianDistribution, ClsTokenAggregator
23 from .vae_processor import VAEProcessor
24
25
26 logger = logging.get_logger(__name__) # pylint: disable=invalid-name
27
28
29 def _resolve_temporal_cat_dtype():
30 raw = os.environ.get("MINIMAX_H3_VAE_DECODER_TEMPORAL_CAT_DTYPE", "").strip().lower()
31 if raw in ("", "0", "false", "no", "off", "none", "keep", "default"):
32 return None
33 mapping = {
34 "fp16": torch.float16,
35 "float16": torch.float16,
36 "half": torch.float16,
37 "bf16": torch.bfloat16,
38 "bfloat16": torch.bfloat16,
39 "fp32": torch.float32,
40 "float32": torch.float32,
41 }
42 if raw not in mapping:
43 raise ValueError(
44 "MINIMAX_H3_VAE_DECODER_TEMPORAL_CAT_DTYPE must be one of "
45 "fp16|bf16|fp32|keep, got %r" % raw
46 )
47 return mapping[raw]
48
49
50 def _resolve_temporal_stream_cat():
51 raw = os.environ.get("MINIMAX_H3_VAE_DECODER_STREAM_TEMPORAL_CAT", "1").strip().lower()
52 return raw not in ("0", "false", "no", "off", "disable", "disabled")
53
54
55 class AutoencoderKL(ModelMixin, ConfigMixin, FromOriginalModelMixin):
56 r"""
57 Abstract shared base for the MiniMax H3 visual VAE.
58
59 This class only carries the shared inference machinery (temporal
60 chunking, tiling, encode/decode entry points). Instantiate the concrete
61 subclass ``AutoencoderKLLegacy`` via ``from_pretrained`` instead.
62 """
63
64 _supports_gradient_checkpointing = True
65 _compilable_modules = ["encoder", "decoder"]
66 _deprecated_kwargs = [
67 "clip_length",
68 "token_drop",
69 "isolated_first_frame",
70 "isolated_last_frame",
71 "isolated_key_frame",
72 "encoder_tiling",
73 "decoder_tiling",
74 "parallel_tiling",
75 "stack_tiling",
76 "tile_size",
77 "tile_overlap_min",
78 "decoder_tile_size",
79 "decoder_tile_overlap_min",
80 "latent_patch_size",
81 "crop_mode",
82 "encoder_parallel",
83 "decoder_parallel",
84 "chunk_dim",
85 ] # legacy config keys accepted by from_pretrained for checkpoint compatibility
86
87
88 def _set_gradient_checkpointing(self, module, value=False):
89 if hasattr(module, "gradient_checkpointing"):
90 module.gradient_checkpointing = value
91
92 def _freeze_nested_module(self, module_path):
93 parts = module_path.split(".")
94 module = self
95 for part in parts:
96 module = getattr(module, part)
97 module.requires_grad_(False)
98
99 def setup_forward(self, **kwargs):
100 self.clip_length = kwargs.get("clip_length", 17)
101 self.token_drop = kwargs.get("token_drop", 0)
102 self.frame_drop = self.token_drop * self.vae_ratio_t
103 self.frame_pre_padding = (-self.clip_length) % self.vae_ratio_t
104 self.tokens_chunk_size = math.ceil(self.clip_length / self.vae_ratio_t)
105 self.token_overlap = (-self.token_drop) % self.tokens_chunk_size
106 self.frame_overlap = max(self.token_overlap * self.vae_ratio_t - self.frame_pre_padding, 0)
107 self.isolated_first_frame = kwargs.get("isolated_first_frame", False)
108 self.isolated_last_frame = kwargs.get("isolated_last_frame", False)
109 self.isolated_key_frame = kwargs.get("isolated_key_frame", False)
110
111 self.encoder_tiling = kwargs.get("encoder_tiling", False)
112 self.decoder_tiling = kwargs.get("decoder_tiling", False)
113 self.stack_tiling = kwargs.get("stack_tiling", False)
114 self.tile_size = kwargs.get("tile_size", 256)
115 self.tile_overlap_min = kwargs.get("tile_overlap_min", 64)
116 self.decoder_tile_size = kwargs.get("decoder_tile_size", self.tile_size)
117 self.decoder_tile_overlap_min = kwargs.get("decoder_tile_overlap_min", self.tile_overlap_min)
118 self.latent_patch_size = kwargs.get("latent_patch_size", 1)
119 self.crop_mode = kwargs.get("crop_mode", "top_left")
120 self.pixel_norm_type = kwargs.get("pixel_norm_type", "imagenet")
121
122 # spatial parallel mode
123 if hasattr(self, "_sp_initialized"):
124 if (
125 kwargs.get("chunk_dim", -1) != self.chunk_dim
126 or kwargs.get("encoder_parallel", False) != self.encoder_parallel
127 or kwargs.get("decoder_parallel", False) != self.decoder_parallel
128 or kwargs.get("parallel_tiling", False) != self.parallel_tiling
129 ):
130 logger.warning(
131 "Do not support changing parallel schema after initialization"
132 )
133 else:
134 self.chunk_dim = kwargs.get("chunk_dim", -1)
135 self.encoder_parallel = kwargs.get("encoder_parallel", False)
136 self.decoder_parallel = kwargs.get("decoder_parallel", False)
137 self.parallel_tiling = kwargs.get("parallel_tiling", False)
138 self._sp_initialized = True
139
140 processor_kwargs = {
141 "vae_ratio": self.vae_ratio,
142 "vae_ratio_t": self.vae_ratio_t,
143 "clip_length": self.clip_length,
144 "frame_overlap": self.frame_overlap,
145 "token_overlap": self.token_overlap,
146 "tokens_chunk_size": self.tokens_chunk_size,
147 "isolated_last_frame": self.isolated_last_frame,
148 "latent_patch_size": self.latent_patch_size,
149 "crop_mode": self.crop_mode,
150 "pixel_norm_type": self.pixel_norm_type,
151 "transform": self.transform,
152 "transform_rev": self.transform_rev,
153 "use_3d_conv": self.use_3d_conv,
154 }
155 if hasattr(self, "processor"):
156 for key, value in processor_kwargs.items():
157 setattr(self.processor, key, value)
158 else:
159 self.processor = VAEProcessor(**processor_kwargs)
160
161 def perform_input_slice(self, x, chunk_size_stride=1):
162 state = get_parallel_state()
163 sp_rank = state["sp_rank"]
164 sp_size = state["sp_size"]
165
166 total_size = x.shape[self.chunk_dim]
167 units = total_size // chunk_size_stride
168 base_units = units // sp_size
169 remainder_units = units % sp_size
170 if sp_rank < remainder_units:
171 start_units = sp_rank * (base_units + 1)
172 end_units = start_units + base_units + 1
173 else:
174 start_units = sp_rank * base_units + remainder_units
175 end_units = start_units + base_units
176 start = start_units * chunk_size_stride
177 end = end_units * chunk_size_stride
178
179 slice_indices = [slice(None)] * x.ndim
180 slice_indices[self.chunk_dim] = slice(start, end)
181 x = x[tuple(slice_indices)].contiguous()
182 return x
183
184 def perform_output_concat(self, x):
185 sp_process_group = get_parallel_state()["sp_process_group"]
186 gathered = all_gather_var_shape(x, group=sp_process_group)
187 x = torch.cat(gathered, dim=self.chunk_dim)
188 return x
189
190
191
192 def split_tiles(self, input_len, is_decoder=False):
193 tile_size = self.decoder_tile_size if is_decoder else self.tile_size
194 tile_overlap_min = self.decoder_tile_overlap_min if is_decoder else self.tile_overlap_min
195
196 if tile_size >= input_len:
197 return [0], [input_len], []
198
199 N = math.ceil(input_len / tile_size)
200 while True:
201 overlaps = [tile_overlap_min] * (N - 1)
202 remaining = tile_size * N - sum(overlaps) - input_len
203
204 if remaining < 0:
205 N += 1
206 else:
207 break
208
209 remaining_units = remaining // self.vae_ratio
210 for i in range(remaining_units):
211 overlaps[i % (N - 1)] += self.vae_ratio
212
213 tile_start_idx = [0]
214 for i in range(N - 1):
215 tile_start_idx.append(tile_start_idx[-1] + tile_size - overlaps[i])
216
217 tile_len = [tile_size] * N
218 return tile_start_idx, tile_len, overlaps
219
220 def blend(
221 self, a: torch.Tensor, b: torch.Tensor, blend_extent: int, dim: int
222 ) -> torch.Tensor:
223 blend_extent = min(a.shape[dim], b.shape[dim], blend_extent)
224
225 positions = torch.arange(blend_extent, device=b.device, dtype=b.dtype)
226 weight_a = 1 - positions / blend_extent
227 weight_b = positions / blend_extent
228
229 shape = [1] * a.ndim
230 shape[dim] = blend_extent
231 weight_a = weight_a.view(shape)
232 weight_b = weight_b.view(shape)
233
234 slice_a = [slice(None)] * a.ndim
235 slice_a[dim] = slice(-blend_extent, None)
236 a_overlap = a[tuple(slice_a)]
237
238 slice_b = [slice(None)] * b.ndim
239 slice_b[dim] = slice(0, blend_extent)
240 b_overlap = b[tuple(slice_b)]
241
242 blended = a_overlap * weight_a + b_overlap * weight_b
243
244 if blend_extent < b.shape[dim]:
245 slice_b_rest = [slice(None)] * b.ndim
246 slice_b_rest[dim] = slice(blend_extent, None)
247 b_rest = b[tuple(slice_b_rest)]
248 return torch.cat([blended, b_rest], dim=dim)
249 else:
250 return blended
251
252 def _all_gather_tiled_results(self, tasks, num_tiles):
253 state = get_parallel_state()
254 group = state["sp_process_group"]
255 sp_size = state["sp_size"]
256 sp_rank = state["sp_rank"]
257
258 if not tasks:
259 raise ValueError(f"Found empty tasks on sp rank {sp_rank}")
260
261 stacked = torch.stack(tasks, dim=0)
262 gathered = all_gather_var_shape(stacked, group=group)
263
264 results = [None] * num_tiles
265 for rank, rank_tensors in enumerate(gathered):
266 num_rank_tasks = rank_tensors.shape[0]
267 for k in range(num_rank_tasks):
268 global_idx = k * sp_size + rank
269 if global_idx >= num_tiles:
270 break
271 results[global_idx] = rank_tensors[k]
272
273 return results
274
275 def _local_tile_indices(self, num_tiles, sp_rank, sp_size):
276 return list(range(sp_rank, num_tiles, sp_size))
277
278 def _run_tile_tasks(self, tiles, tile_indices, forward_fn, stack_tiling, cls_agg=None):
279 if stack_tiling and tile_indices:
280 sample_batch_size = tiles[0].shape[0]
281 tile_batch = torch.cat([tiles[idx] for idx in tile_indices], dim=0)
282 output_batch = forward_fn(tile_batch)
283 output_tiles = output_batch.unflatten(
284 0, (len(tile_indices), sample_batch_size)
285 ).unbind(dim=0)
286 if cls_agg is not None:
287 cls_agg.collect_stacked(len(tile_indices), sample_batch_size)
288 return list(output_tiles)
289
290 tasks = []
291 for idx in tile_indices:
292 tasks.append(forward_fn(tiles[idx]))
293 if cls_agg is not None:
294 cls_agg.collect()
295 return tasks
296
297 def tiled_encode(self, x):
298 if self.parallel_tiling: # Fast online encoding for large videos
299 state = get_parallel_state()
300 sp_rank = state["sp_rank"]
301 sp_size = state["sp_size"]
302 else:
303 sp_rank, sp_size = 0, 1
304
305 height, width = x.shape[-2], x.shape[-1]
306 y_idx, y_len, y_overlap = self.split_tiles(height, False)
307 x_idx, x_len, x_overlap = self.split_tiles(width, False)
308
309 i_max, j_max = len(y_idx), len(x_idx)
310 num_tiles = i_max * j_max
311
312 x_tiles = []
313 for i, (i_pos, i_len) in enumerate(zip(y_idx, y_len)):
314 for j, (j_pos, j_len) in enumerate(zip(x_idx, x_len)):
315 tile = x[..., i_pos : i_pos + i_len, j_pos : j_pos + j_len]
316 x_tiles.append(tile)
317
318 with ClsTokenAggregator(self) as agg:
319 local_tile_indices = self._local_tile_indices(num_tiles, sp_rank, sp_size)
320 stack_tiling = self.stack_tiling and not (
321 self.training and getattr(self.encoder, "mask_enabled", False)
322 )
323 encoded_tasks = self._run_tile_tasks(
324 x_tiles, local_tile_indices, self.encode, stack_tiling, agg
325 )
326
327 if sp_size > 1:
328 dist.barrier(group=get_parallel_state()["sp_process_group"])
329 all_encoded = self._all_gather_tiled_results(encoded_tasks, num_tiles)
330 if agg.cls_tokens:
331 agg.cls_tokens = self._all_gather_tiled_results(agg.cls_tokens, num_tiles)
332 else:
333 all_encoded = encoded_tasks
334
335 rows = [[None for _ in range(j_max)] for _ in range(i_max)]
336 for idx, encoded in enumerate(all_encoded):
337 i, j = idx // j_max, idx % j_max
338 rows[i][j] = encoded.to(x.device)
339
340 latent_y_overlap = [
341 tile_overlap // self.vae_ratio for tile_overlap in y_overlap
342 ]
343 latent_x_overlap = [
344 tile_overlap // self.vae_ratio for tile_overlap in x_overlap
345 ]
346
347 result_rows = []
348 for i, row in enumerate(rows):
349 result_row = []
350 for j, tile in enumerate(row):
351 if i > 0:
352 tile = self.blend(rows[i - 1][j], tile, latent_y_overlap[i - 1], dim=-2)
353 if j > 0:
354 tile = self.blend(row[j - 1], tile, latent_x_overlap[j - 1], dim=-1)
355 if i < len(rows) - 1:
356 tile = tile[..., : -latent_y_overlap[i], :]
357 if j < len(row) - 1:
358 tile = tile[..., :, : -latent_x_overlap[j]]
359 result_row.append(tile)
360 result_rows.append(torch.cat(result_row, dim=-1))
361 z = torch.cat(result_rows, dim=-2)
362
363 return z
364
365 def tiled_decode(self, z):
366 if self.parallel_tiling: # Fast online decoding for large videos
367 state = get_parallel_state()
368 sp_rank = state["sp_rank"]
369 sp_size = state["sp_size"]
370 else:
371 sp_rank, sp_size = 0, 1
372
373 height, width = (
374 z.shape[-2] * self.vae_ratio,
375 z.shape[-1] * self.vae_ratio,
376 )
377 y_idx, y_len, y_overlap = self.split_tiles(height, True)
378 x_idx, x_len, x_overlap = self.split_tiles(width, True)
379
380 i_max, j_max = len(y_idx), len(x_idx)
381 num_tiles = i_max * j_max
382
383 z_tiles = []
384 for i, (i_pos, i_len) in enumerate(zip(y_idx, y_len)):
385 i_pos, i_len = (
386 i_pos // self.vae_ratio,
387 i_len // self.vae_ratio,
388 )
389 for j, (j_pos, j_len) in enumerate(zip(x_idx, x_len)):
390 j_pos, j_len = (j_pos // self.vae_ratio, j_len // self.vae_ratio)
391 tile = z[..., i_pos : i_pos + i_len, j_pos : j_pos + j_len]
392 z_tiles.append(tile)
393
394 local_tile_indices = self._local_tile_indices(num_tiles, sp_rank, sp_size)
395 stack_tiling = self.stack_tiling and not (
396 self.training and getattr(self.decoder, "mask_enabled", False)
397 )
398 decoded_tasks = self._run_tile_tasks(
399 z_tiles, local_tile_indices, self.decode, stack_tiling
400 )
401
402 if sp_size > 1:
403 dist.barrier(group=get_parallel_state()["sp_process_group"])
404 all_decoded = self._all_gather_tiled_results(decoded_tasks, num_tiles)
405 else:
406 all_decoded = decoded_tasks
407
408
409 rows = [[None for _ in range(j_max)] for _ in range(i_max)]
410 for idx, decoded in enumerate(all_decoded):
411 i, j = idx // j_max, idx % j_max
412 rows[i][j] = decoded.to(z.device)
413
414 result_rows = []
415 for i, row in enumerate(rows):
416 result_row = []
417 for j, tile in enumerate(row):
418 if i > 0:
419 tile = self.blend(rows[i - 1][j], tile, y_overlap[i - 1], dim=-2)
420 if j > 0:
421 tile = self.blend(row[j - 1], tile, x_overlap[j - 1], dim=-1)
422 if i < len(rows) - 1:
423 tile = tile[..., : -y_overlap[i], :]
424 if j < len(row) - 1:
425 tile = tile[..., :, : -x_overlap[j]]
426 result_row.append(tile)
427 result_rows.append(torch.cat(result_row, dim=-1))
428 dec = torch.cat(result_rows, dim=-2)
429 return dec
430
431 def _adaptive_encode(self, x):
432 if self.encoder_tiling:
433 return self.tiled_encode(x)
434 else:
435 return self.encode(x)
436
437 def _adaptive_decode(self, z):
438 if self.decoder_tiling:
439 return self.tiled_decode(z)
440 else:
441 return self.decode(z)
442
443 def trim_code(self, z, target_codes):
444 if target_codes < z.shape[2]:
445 if self.causal_encoder:
446 z = z[:, :, -target_codes:, :, :]
447 else:
448 start_frame = (z.shape[2] - target_codes) // 2
449 z = z[:, :, start_frame : start_frame + target_codes, :, :]
450 return z
451
452 def trim_output(self, dec, target_frames):
453 if target_frames < dec.shape[2]:
454 if self.causal_encoder: # This is defined by encoder, not decoder
455 dec = dec[:, :, -target_frames:, :, :]
456 else:
457 start_frame = (dec.shape[2] - target_frames) // 2
458 dec = dec[:, :, start_frame : start_frame + target_frames, :, :]
459 return dec
460
461 def encode_temporal(self, x):
462 offset_frame = 1 if self.isolated_first_frame and self.frame_pre_padding == 0 else 0
463
464 if x.shape[2] % self.clip_length != offset_frame:
465 pad_size = (offset_frame - x.shape[2]) % self.clip_length
466 pad_frames = x[:, :, -1:].repeat(1, 1, pad_size, 1, 1)
467 x = torch.cat([x, pad_frames], dim=2)
468
469 num_chunks = (x.shape[2] - offset_frame) // self.clip_length
470
471 z_list = []
472 for i in range(num_chunks):
473 start_idx = i * self.clip_length + offset_frame
474 end_idx = (i + 1) * self.clip_length + offset_frame
475 clip_x = x[:, :, start_idx:end_idx, :, :]
476
477 if self.isolated_key_frame:
478 key_frame = clip_x[:, :, :1, :, :]
479 z_key = self._adaptive_encode(key_frame)
480
481 if clip_x.shape[2] > 1:
482 video_frames = clip_x[:, :, 1:, :, :]
483 z_video = self._adaptive_encode(video_frames)
484 z = torch.cat([z_key, z_video], dim=2)
485 else:
486 z = z_key
487 else:
488 z = self._adaptive_encode(clip_x)
489
490 z_list.append(z)
491
492 z = torch.cat(z_list, dim=2)
493 if self.token_drop > 0:
494 z = z[:, :, : -self.token_drop]
495
496 if self.isolated_first_frame:
497 input_first_frame = x[:, :, :1, :, :]
498 z_first_frame = self._adaptive_encode(input_first_frame)
499
500 if self.frame_pre_padding == 0:
501 z = torch.cat([z_first_frame, z], dim=2)
502 else:
503 z = torch.cat([z_first_frame, z[:, :, 1:, :, :]], dim=2)
504
505 if self.isolated_last_frame:
506 frame_num = x.shape[2]
507 last_frame_idx = frame_num - self.frame_drop + offset_frame
508 input_last_frame = x[:, :, last_frame_idx : last_frame_idx + 1, :, :]
509 z_last_frame = self._adaptive_encode(input_last_frame)
510 z = torch.cat([z, z_last_frame], dim=2)
511
512 return z
513
514 def _decode_temporal_pad_frames(self, z, pad_tokens):
515 if pad_tokens <= 0:
516 return 0
517 intra_tail = self.clip_length % self.vae_ratio_t
518 if intra_tail == 0:
519 return int(pad_tokens) * int(self.vae_ratio_t)
520
521 z_len_before_pad = z.shape[2] - pad_tokens
522 return sum(
523 (
524 intra_tail
525 if (z_len_before_pad + k) % self.tokens_chunk_size == 0
526 else self.vae_ratio_t
527 )
528 for k in range(pad_tokens)
529 )
530
531 def _decode_temporal_output_frame_plan(self, z, z_head, z_tail, num_chunks, pad_tokens):
532 chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
533 split_count = int(self.token_drop > 0) + 1
534 total_frames = 0
535 final_overlap_frames = 0
536
537 if z_head is not None:
538 total_frames += 1
539
540 for i in range(num_chunks):
541 t_start_idx = i * self.tokens_chunk_size
542 t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap
543 clip_token_len = max(0, min(t_end_idx, z.shape[2]) - min(t_start_idx, z.shape[2]))
544 if i == 0 and z_head is not None:
545 clip_token_len += z_head.shape[2]
546 if i == num_chunks - 1 and z_tail is not None:
547 clip_token_len += z_tail.shape[2]
548
549 clip_frame_len = clip_token_len * self.vae_ratio_t
550 if i == 0 and z_head is not None:
551 clip_frame_len = max(0, clip_frame_len - self.vae_ratio_t)
552 if i == num_chunks - 1 and z_tail is not None:
553 clip_frame_len = max(0, clip_frame_len - self.vae_ratio_t)
554
555 for j in range(split_count):
556 f_start_idx = j * chunk_dec
557 f_end_idx = min(f_start_idx + chunk_dec, clip_frame_len)
558 chunk_frames = max(0, f_end_idx - f_start_idx - self.frame_pre_padding)
559 if j == 0:
560 total_frames += chunk_frames
561 else:
562 final_overlap_frames = chunk_frames
563
564 total_frames += final_overlap_frames
565 if z_tail is not None:
566 total_frames += 1
567
568 pad_frames = self._decode_temporal_pad_frames(z, pad_tokens)
569 return int(total_frames), int(pad_frames), int(total_frames - pad_frames)
570
571 def _decode_temporal_streaming(self, z, z_head, z_tail, num_chunks, pad_tokens, temporal_cat_dtype):
572 total_frames, pad_frames, output_frames = self._decode_temporal_output_frame_plan(
573 z, z_head, z_tail, num_chunks, pad_tokens
574 )
575 if output_frames <= 0:
576 raise ValueError(
577 f"decode_temporal streaming planned non-positive output_frames={output_frames} "
578 f"total_frames={total_frames} pad_frames={pad_frames}"
579 )
580
581
582 chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
583 split_count = int(self.token_drop > 0) + 1
584 dec = None
585 dec_overlap = None
586 write_pos = 0
587 logical_frames = 0
588 dropped_frames = 0
589 decoded_count = 0
590
591 def write_part(part):
592 nonlocal dec, write_pos, logical_frames, dropped_frames
593 part_frames = int(part.shape[2])
594 if part_frames <= 0:
595 return
596 logical_frames += part_frames
597 if dec is None:
598 out_shape = list(part.shape)
599 out_shape[2] = output_frames
600 dec = torch.empty(out_shape, dtype=part.dtype, device=part.device)
601
602 remaining = int(dec.shape[2]) - write_pos
603 copy_frames = min(part_frames, max(0, remaining))
604 if copy_frames > 0:
605 dec[:, :, write_pos : write_pos + copy_frames, :, :].copy_(
606 part[:, :, :copy_frames, :, :]
607 )
608 write_pos += copy_frames
609 dropped_frames += part_frames - copy_frames
610
611 for i in range(num_chunks):
612 t_start_idx = i * self.tokens_chunk_size
613 t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap
614 clip_z = z[:, :, t_start_idx:t_end_idx, :, :]
615
616 if i == 0 and z_head is not None:
617 clip_z = torch.cat([z_head, clip_z], dim=2)
618
619 if i == num_chunks - 1 and z_tail is not None:
620 clip_z = torch.cat([clip_z, z_tail], dim=2)
621
622 clip_dec = self._adaptive_decode(clip_z)
623 decoded_count += 1
624 if temporal_cat_dtype is not None and clip_dec.dtype != temporal_cat_dtype:
625 clip_dec = clip_dec.to(temporal_cat_dtype)
626 if clip_dec.device != z.device:
627 clip_dec = clip_dec.to(z.device)
628
629
630 dec_tail = None
631 if i == 0 and z_head is not None:
632 write_part(clip_dec[:, :, self.vae_ratio_t - 1 : self.vae_ratio_t, :, :])
633 clip_dec = clip_dec[:, :, self.vae_ratio_t :, :, :]
634
635 if i == num_chunks - 1 and z_tail is not None:
636 dec_tail = clip_dec[:, :, -1:, :, :]
637 clip_dec = clip_dec[:, :, : -self.vae_ratio_t, :, :]
638
639 for j in range(split_count):
640 f_start_idx = j * chunk_dec
641 f_end_idx = min(f_start_idx + chunk_dec, clip_dec.shape[2])
642 clip_dec_chunk = clip_dec[:, :, f_start_idx:f_end_idx, :, :]
643 clip_dec_chunk = clip_dec_chunk[:, :, self.frame_pre_padding :, :, :]
644
645 if j == 0:
646 if dec_overlap is not None:
647 clip_dec_chunk = self.blend(
648 dec_overlap, clip_dec_chunk, self.frame_overlap, dim=-3
649 )
650 dec_overlap = None
651 write_part(clip_dec_chunk)
652 else:
653 # Break the view's reference to the full decoded clip so earlier
654 # temporal chunks can be released before the final output exists.
655 dec_overlap = clip_dec_chunk.contiguous()
656
657 if i == num_chunks - 1:
658 if dec_overlap is not None:
659 write_part(dec_overlap)
660 dec_overlap = None
661 if dec_tail is not None:
662 write_part(dec_tail)
663
664 del clip_dec, clip_z
665
666 if dec is None:
667 raise RuntimeError("decode_temporal streaming produced no output tensor")
668 if logical_frames != total_frames or dropped_frames != pad_frames or write_pos != output_frames:
669 raise RuntimeError(
670 "decode_temporal streaming frame plan mismatch: "
671 f"logical_frames={logical_frames} total_frames={total_frames} "
672 f"dropped_frames={dropped_frames} pad_frames={pad_frames} "
673 f"write_pos={write_pos} output_frames={output_frames}"
674 )
675
676 return dec
677
678 def decode_temporal(self, z):
679 chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
680
681 isolated_token_num = 0
682 if self.isolated_first_frame and self.frame_pre_padding == 0:
683 isolated_token_num = isolated_token_num + 1
684 if self.isolated_last_frame:
685 isolated_token_num = isolated_token_num + 1
686
687 pseudo_total_tokens = z.shape[2] - isolated_token_num + self.token_drop
688
689 pad_tokens = 0
690 remainder = pseudo_total_tokens % self.tokens_chunk_size
691 if remainder != 0:
692 if self.training:
693 raise ValueError(f"Temporal token length {z.shape[2]} is wrong!")
694 else:
695 pad_tokens = self.tokens_chunk_size - remainder
696 pseudo_total_tokens = pseudo_total_tokens + pad_tokens
697
698 pseudo_num_chunks = pseudo_total_tokens // self.tokens_chunk_size
699 num_chunks = pseudo_num_chunks - int(self.token_drop > 0)
700
701 z_head = None
702 if self.isolated_first_frame and self.frame_pre_padding == 0:
703 z_head = z[:, :, :1, :, :]
704 z = z[:, :, 1:, :, :]
705
706 z_tail = None
707 if self.isolated_last_frame:
708 z_tail = z[:, :, -1:, :, :]
709 z = z[:, :, :-1, :, :]
710
711 if pad_tokens > 0:
712 pad_z = z[:, :, -1:, :, :].repeat(1, 1, pad_tokens, 1, 1)
713 z = torch.cat([z, pad_z], dim=2)
714
715 temporal_cat_dtype = _resolve_temporal_cat_dtype()
716 if not self.training and _resolve_temporal_stream_cat():
717 return self._decode_temporal_streaming(
718 z, z_head, z_tail, num_chunks, pad_tokens, temporal_cat_dtype
719 )
720
721 decoded_tasks = []
722 for i in range(num_chunks):
723 t_start_idx = i * self.tokens_chunk_size
724 t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap
725 clip_z = z[:, :, t_start_idx:t_end_idx, :, :]
726
727 if i == 0 and z_head is not None:
728 clip_z = torch.cat([z_head, clip_z], dim=2)
729
730 if i == num_chunks - 1 and z_tail is not None:
731 clip_z = torch.cat([clip_z, z_tail], dim=2)
732
733 clip_dec = self._adaptive_decode(clip_z)
734 if temporal_cat_dtype is not None and clip_dec.dtype != temporal_cat_dtype:
735 clip_dec = clip_dec.to(temporal_cat_dtype)
736
737 decoded_tasks.append((i, clip_dec))
738
739 clip_dec_list = [clip_dec.to(z.device) for _, clip_dec in decoded_tasks]
740
741 dec_list = []
742 dec_overlap = None
743
744 dec_head = None
745 if z_head is not None:
746 dec_head = clip_dec_list[0][:, :, self.vae_ratio_t - 1 : self.vae_ratio_t, :, :]
747 clip_dec_list[0] = clip_dec_list[0][:, :, self.vae_ratio_t :, :, :]
748
749 dec_tail = None
750 if z_tail is not None:
751 dec_tail = clip_dec_list[-1][:, :, -1:, :, :]
752 clip_dec_list[-1] = clip_dec_list[-1][:, :, : -self.vae_ratio_t, :, :]
753
754 if dec_head is not None:
755 dec_list.append(dec_head)
756
757 for i in range(num_chunks):
758 for j in range(int(self.token_drop > 0) + 1):
759 clip_dec = clip_dec_list[i]
760
761 f_start_idx = j * chunk_dec
762 f_end_idx = min(f_start_idx + chunk_dec, clip_dec.shape[2])
763 clip_dec_chunk = clip_dec[:, :, f_start_idx:f_end_idx, :, :]
764 clip_dec_chunk = clip_dec_chunk[:, :, self.frame_pre_padding :, :, :]
765
766 if j == 0:
767 if dec_overlap is not None:
768 clip_dec_chunk = self.blend(
769 dec_overlap, clip_dec_chunk, self.frame_overlap, dim=-3
770 )
771 dec_list.append(clip_dec_chunk)
772 else:
773 dec_overlap = clip_dec_chunk
774
775 if dec_overlap is not None:
776 dec_list.append(dec_overlap)
777
778 if dec_tail is not None:
779 dec_list.append(dec_tail)
780
781
782 dec = torch.cat(dec_list, dim=2)
783
784 pad_frames = self._decode_temporal_pad_frames(z, pad_tokens)
785 if pad_frames > 0:
786 dec = dec[:, :, :-pad_frames, :, :]
787
788 return dec
789
790 def decode_base(self, z, frame_num=None, process_image=False):
791 if process_image or not self.use_3d_conv:
792 if not self.use_3d_conv and z.ndim == 5:
793 z = z.squeeze(2)
794
795 recon = self._adaptive_decode(z)
796 else:
797 recon = self.decode_temporal(z)
798
799 if self.use_3d_conv:
800 if frame_num is not None:
801 target_frames = frame_num
802 else:
803 target_frames = recon.shape[2]
804
805 recon = self.trim_output(recon, target_frames)
806 if process_image:
807 recon = recon.squeeze(2)
808
809 return recon
810
811 #########################################################
812 # freeze_scope is retained from the training codebase: in this
813 # inference-only bundle (self.training is always False) it simply
814 # provides the no_grad() context used by encode()/decode().
815 #########################################################
816
817
818 def freeze_scope(self, module_name):
819 if not self.training:
820 return torch.no_grad()
821
822 if_freeze = module_name in self.fix_modules
823 if if_freeze:
824 return torch.no_grad()
825 else:
826 return nullcontext()
827
828
829
830
831
832
833
834
835 #########################################################
836 # following methods are for inference
837 #########################################################
838
839 @torch.no_grad()
840 def encode_images(
841 self,
842 images: Union[List[np.ndarray], List[torch.Tensor]],
843 transform_input: bool = False,
844 use_fp16_latent: bool = False,
845 verbose: bool = False,
846 ) -> List[torch.Tensor]:
847 """encode images into latents
848
849 Args:
850 images (Union[List[np.ndarray], List[torch.Tensor]]):
851 List of images, single input will be wrapped in a list.
852 If input is a list of np.ndarray, it should be in shape B * (H, W, 3), dtype uint8.
853 If input is a list of torch.Tensor, it should be in shape B * (3, H, W), dtype float32.
854 transform_input (bool, optional):
855 Whether to transform input using ImageNet std/mean. Defaults to False.
856 If input is a list of np.ndarray, it will always be set to True.
857 use_fp16_latent (bool, optional):
858 Whether to use fp16 latent. Defaults to False.
859 verbose (bool, optional):
860 Whether to print debug information. Defaults to False.
861
862 Returns:
863 List[torch.Tensor]:
864 List of image latents.
865 If self.use_3d_conv is True, it should be in shape B * (D, 1, H', W').
866 Otherwise, it should be in shape B * (D, H', W').
867 """
868
869 images = self.processor._ensure_list(images)
870
871 if isinstance(images[0], Image.Image):
872 images = [np.array(image) for image in images]
873
874 if isinstance(images[0], np.ndarray):
875 device = next(self.parameters()).device
876 images = self.processor.convert_numpy_to_tensor(images, device)
877 images = torch.split(images, 1, dim=0)
878 transform_input = True
879
880 if transform_input:
881 images = [
882 image.unsqueeze(0) if image.ndim == 3 else image for image in images
883 ]
884 images = [self.processor.transform_tensor(image) for image in images]
885
886 prepared = []
887 for image_tensor in images:
888 if image_tensor.ndim == 3:
889 image_tensor = image_tensor.unsqueeze(0)
890 _, _, h, w = image_tensor.shape
891 new_h, new_w = self.processor._align_to_total_patch_size(h, w)
892 image_tensor = self.processor._crop_to_align(image_tensor, new_h, new_w)
893 prepared.append(image_tensor)
894
895 if len(prepared) > 1 and len(set(t.shape for t in prepared)) == 1:
896 stacked = torch.cat(prepared, dim=0)
897 if verbose:
898 logger.info(f"batch encode input shape {tuple(stacked.shape)}")
899 all_latents = self.encode_base(stacked, True)
900 image_latents = [all_latents[i].contiguous() for i in range(all_latents.shape[0])]
901 else:
902 image_latents = []
903 for image_tensor in prepared:
904 if verbose:
905 logger.info(f"input shape {tuple(image_tensor.shape)}")
906 image_latent = self.encode_base(image_tensor, True)
907 image_latents.append(image_latent.squeeze(0).contiguous())
908
909 if use_fp16_latent:
910 image_latents = [lat.to(torch.float16) for lat in image_latents]
911
912 if verbose:
913 for lat in image_latents:
914 logger.info(f"image latent shape {tuple(lat.shape)}")
915
916 return image_latents
917
918 @torch.no_grad()
919 def encode_videos(
920 self,
921 videos: Union[List[np.ndarray], List[torch.Tensor]],
922 transform_input: bool = False,
923 use_fp16_latent: bool = False,
924 verbose: bool = False,
925 encode_prefix: bool = False,
926 ) -> List[torch.Tensor]:
927 """encode videos into latents
928
929 Args:
930 videos (Union[List[np.ndarray], List[torch.Tensor]]):
931 List of videos, single input will be wrapped in a list.
932 If input is a list of np.ndarray, it should be in shape B * (T, H, W, 3), dtype uint8.
933 If input is a list of torch.Tensor, it should be in shape B * (3, T, H, W), dtype float32.
934 transform_input (bool, optional):
935 Whether to transform input using ImageNet std/mean. Defaults to False.
936 If input is a list of np.ndarray, it will always be set to True.
937 use_fp16_latent (bool, optional):
938 Whether to use fp16 latent. Defaults to False.
939 verbose (bool, optional):
940 Whether to print debug information. Defaults to False.
941 encode_prefix (bool, optional):
942 Continuation (prefix) mode: prepend normalized
943 black frames to token alignment, append black frames to chunk
944 alignment, encode with token_drop disabled, then discard only
945 the trailing padding tokens. Returns both latents and leading
946 pad-frame counts. Defaults to False.
947
948 Returns:
949 List[torch.Tensor]:
950 List of video latents, shape B * (D, T', H', W').
951 With encode_prefix=True, returns
952 (List[torch.Tensor], List[int]).
953 """
954
955 videos = self.processor._ensure_list(videos)
956
957 if isinstance(videos[0], np.ndarray):
958 device = next(self.parameters()).device
959 videos = [self.processor.convert_numpy_to_tensor(video, device) for video in videos]
960 transform_input = True
961
962 if transform_input:
963 videos = [self.processor.transform_tensor(video) for video in videos]
964 videos = [video.transpose(0, 1) for video in videos]
965
966 if encode_prefix:
967 if self.isolated_last_frame:
968 raise ValueError(
969 "encode_prefix does not support isolated_last_frame"
970 )
971
972 video_latents = []
973 prefix_pad_frames = []
974 for video in videos:
975 if video.ndim == 4:
976 video = video.unsqueeze(0)
977 _, _, _, h, w = video.shape
978 new_h, new_w = self.processor._align_to_total_patch_size(h, w)
979 video = self.processor._crop_to_align(
980 video, new_h, new_w, is_video=True
981 )
982
983 model_alignment = (
984 self.token_drop,
985 self.frame_drop,
986 self.token_overlap,
987 self.frame_overlap,
988 )
989 processor_alignment = (
990 self.processor.token_overlap,
991 self.processor.frame_overlap,
992 )
993 self.token_drop = 0
994 self.frame_drop = 0
995 self.token_overlap = 0
996 self.frame_overlap = 0
997 self.processor.token_overlap = 0
998 self.processor.frame_overlap = 0
999 try:
1000 orig_frames = video.shape[2]
1001 leading, trailing, drop_tokens = (
1002 self.processor.align_video_length_2pass(orig_frames)
1003 )
1004 _, _, _, cropped_h, cropped_w = video.shape
1005 if leading > 0:
1006 black = self.processor.transform(
1007 video.new_zeros(leading, 3, cropped_h, cropped_w)
1008 )
1009 black = black.unsqueeze(0).permute(0, 2, 1, 3, 4)
1010 video = torch.cat([black, video], dim=2)
1011 if trailing > 0:
1012 black = self.processor.transform(
1013 video.new_zeros(trailing, 3, cropped_h, cropped_w)
1014 )
1015 black = black.unsqueeze(0).permute(0, 2, 1, 3, 4)
1016 video = torch.cat([video, black], dim=2)
1017
1018 if verbose:
1019 logger.info(
1020 f"[encode_prefix] {orig_frames} frames -> "
1021 f"pad leading={leading}, trailing={trailing} -> "
1022 f"{video.shape[2]} frames"
1023 )
1024
1025 video_latent = self.encode_base(video, False)
1026 if drop_tokens > 0:
1027 video_latent = video_latent[:, :, :-drop_tokens, :, :]
1028 prefix_pad_frames.append(leading)
1029 finally:
1030 (
1031 self.token_drop,
1032 self.frame_drop,
1033 self.token_overlap,
1034 self.frame_overlap,
1035 ) = model_alignment
1036 (
1037 self.processor.token_overlap,
1038 self.processor.frame_overlap,
1039 ) = processor_alignment
1040
1041 video_latents.append(video_latent.squeeze(0).contiguous())
1042
1043 if use_fp16_latent:
1044 video_latents = [lat.to(torch.float16) for lat in video_latents]
1045 if verbose:
1046 for latent in video_latents:
1047 logger.info(f"video latent shape {tuple(latent.shape)}")
1048 return video_latents, prefix_pad_frames
1049
1050 prepared = []
1051 for video in videos:
1052 if video.ndim == 4:
1053 video = video.unsqueeze(0)
1054 used_frame_length = self.processor.get_suitable_video_length(video.shape[2], verbose)
1055 _, _, _, h, w = video.shape
1056 new_h, new_w = self.processor._align_to_total_patch_size(h, w)
1057 video = video[:, :, :used_frame_length, :, :]
1058 video = self.processor._crop_to_align(video, new_h, new_w, is_video=True)
1059 prepared.append(video)
1060
1061 if len(prepared) > 1 and len(set(t.shape for t in prepared)) == 1:
1062 stacked = torch.cat(prepared, dim=0)
1063 if verbose:
1064 logger.info(f"batch encode input shape {tuple(stacked.shape)}")
1065 all_latents = self.encode_base(stacked, False)
1066 video_latents = [all_latents[i].contiguous() for i in range(all_latents.shape[0])]
1067 else:
1068 video_latents = []
1069 for video in prepared:
1070 if verbose:
1071 logger.info(f"input shape {tuple(video.shape)}")
1072 video_latent = self.encode_base(video, False)
1073 video_latents.append(video_latent.squeeze(0).contiguous())
1074
1075 if use_fp16_latent:
1076 video_latents = [lat.to(torch.float16) for lat in video_latents]
1077
1078 if verbose:
1079 for lat in video_latents:
1080 logger.info(f"video latent shape {tuple(lat.shape)}")
1081
1082 return video_latents
1083
1084
1085
1086
1087 # ============================================================================
1088 # Legacy CNN VAE
1089 # ============================================================================
1090
1091
1092 class AutoencoderKLLegacy(AutoencoderKL):
1093 r"""
1094 A VAE model (legacy CNN-based) for encoding pixels into latents and decoding latent representations into pixels.
1095 """
1096
1097 @register_to_config
1098 def __init__(
1099 self,
1100 in_channels=3,
1101 out_ch=3,
1102 ch=128,
1103 embed_dim=16,
1104 z_channels=16,
1105 use_3d_conv=False,
1106 # cnn vae
1107 zq_ch_encoder=None,
1108 zq_ch_decoder=None,
1109 num_res_blocks=2,
1110 num_res_blocks_decoder=None,
1111 ch_mult=[1, 2, 2, 4, 4, 8],
1112 space_down=[2, 2, 2, 2, 1, 1],
1113 space_up=[1, 2, 2, 2, 2, 1],
1114 time_down=None,
1115 time_up=None,
1116 padding_mode="zeros",
1117 padding_mode_t=None,
1118 use_t_isolated_gn=False,
1119 causal_encoder=True,
1120 causal_decoder=True,
1121 use_vit_decoder=False,
1122 vit_decoder_kwargs=None,
1123 # stats
1124 shift_factor=0.0,
1125 scaling_factor=1.0,
1126 # pixel normalization
1127 pixel_norm_type="imagenet",
1128 # others
1129 **kwargs,
1130 ):
1131 ModelMixin.__init__(self) # NOTE: avoid wrong @register_to_config
1132
1133 if not use_3d_conv or not use_vit_decoder:
1134 raise NotImplementedError(
1135 "this release only supports use_3d_conv=True with use_vit_decoder=True"
1136 )
1137
1138 self.transform = get_normalize_transform(pixel_norm_type)
1139 self.transform_rev = get_denormalize_transform(pixel_norm_type)
1140
1141 self.use_3d_conv = use_3d_conv
1142 self.causal_encoder = causal_encoder
1143 self.causal_decoder = causal_decoder
1144 self.slidedec = self.causal_encoder and not self.causal_decoder
1145
1146 # some registered parameters for simplicity
1147 self.vae_ratio = int(np.cumprod(space_down)[-1])
1148 self.vae_ratio_t = int(np.cumprod(time_down)[-1]) if time_down else 1
1149 self.config["vae_ratio"] = self.vae_ratio
1150 self.config["vae_ratio_t"] = self.vae_ratio_t
1151
1152 # some registered parameters for inference and training
1153 self.setup_forward(**kwargs)
1154 self.setup_training(**kwargs)
1155
1156 # init encoder
1157 encoder_config = {
1158 "double_z": True,
1159 "z_channels": z_channels,
1160 "zq_ch": zq_ch_encoder,
1161 "in_channels": in_channels,
1162 "ch": ch,
1163 "num_res_blocks": num_res_blocks,
1164 "ch_mult": ch_mult,
1165 "space_down": space_down,
1166 "time_down": time_down,
1167 "padding_mode": padding_mode,
1168 "padding_mode_t": padding_mode_t,
1169 "causal": causal_encoder,
1170 "use_t_isolated_gn": use_t_isolated_gn,
1171 }
1172 self.encoder = EncoderFCN3D(**encoder_config)
1173
1174 # init pointwise quant/post_quant conv
1175 self.quant_conv = nn.Conv3d(z_channels * 2, 2 * embed_dim, 1)
1176 self.post_quant_conv = nn.Conv3d(embed_dim, z_channels, 1)
1177
1178 self.use_vit_decoder = use_vit_decoder
1179
1180 # init decoder
1181 vit_kwargs = {
1182 "patch_size": self.vae_ratio,
1183 "in_channels": z_channels,
1184 "out_channels": out_ch,
1185 **(vit_decoder_kwargs or {}),
1186 }
1187 vit_kwargs.setdefault("patch_size_t", self.vae_ratio_t)
1188 vit_kwargs.setdefault("t_causal", causal_decoder)
1189 self.decoder = ViT3DDecoder(**vit_kwargs)
1190
1191 apply_spatial_parallel(self.encoder, self.encoder_parallel, self.chunk_dim)
1192 apply_spatial_parallel(self.decoder, self.decoder_parallel, self.chunk_dim)
1193
1194 for module in set(self.fix_modules + self.frozen_modules):
1195 self._freeze_nested_module(module)
1196
1197 self.gradient_checkpointing = False
1198
1199 def encode(self, x):
1200 if self.encoder_parallel:
1201 x = self.perform_input_slice(x, self.vae_ratio)
1202
1203 with self.freeze_scope("encoder"):
1204 h = self.encoder(x)
1205
1206 with self.freeze_scope("quant_conv"):
1207 moments = self.quant_conv(h)
1208
1209 if self.encoder_parallel:
1210 moments = self.perform_output_concat(moments)
1211
1212 return moments
1213
1214 def decode(self, z):
1215 if self.decoder_parallel and not self.use_vit_decoder:
1216 z = self.perform_input_slice(z)
1217
1218 with self.freeze_scope("post_quant_conv"):
1219 z2 = self.post_quant_conv(z)
1220
1221 with self.freeze_scope("decoder"):
1222 if self.use_vit_decoder:
1223 dec = self.decoder(z2)
1224 else:
1225 dec = self.decoder(z2, z)
1226
1227 if self.decoder_parallel and not self.use_vit_decoder:
1228 dec = self.perform_output_concat(dec)
1229 return dec
1230
1231 def encode_base(self, input, process_image=False):
1232 if self.use_3d_conv and input.ndim == 4:
1233 input = input.unsqueeze(2)
1234
1235 if process_image or not self.use_3d_conv:
1236 moments = self._adaptive_encode(input)
1237 else:
1238 moments = self.encode_temporal(input)
1239
1240 z = DiagonalGaussianDistribution(moments).sample()
1241
1242 if process_image and self.use_3d_conv:
1243 z = self.trim_code(z, 1)
1244
1245 return z
1246
1247 #########################################################
1248 # training-related knobs kept only for checkpoint/config compatibility
1249 #########################################################
1250
1251 def setup_training(self, **kwargs):
1252 self.fix_modules = kwargs.get("fix_modules", [])
1253 self.frozen_modules = kwargs.get("frozen_modules", [])
1254
1255
1256
1257
1258
1259