|
import torch |
|
import torch.nn as nn |
|
from diffusers.models import AutoencoderKL, AutoencoderKLTemporalDecoder |
|
from einops import rearrange |
|
|
|
from opensora.registry import MODELS |
|
|
|
|
|
@MODELS.register_module() |
|
class VideoAutoencoderKL(nn.Module): |
|
def __init__(self, from_pretrained=None, micro_batch_size=None): |
|
super().__init__() |
|
self.module = AutoencoderKL.from_pretrained(from_pretrained) |
|
self.out_channels = self.module.config.latent_channels |
|
self.patch_size = (1, 8, 8) |
|
self.micro_batch_size = micro_batch_size |
|
|
|
def encode(self, x): |
|
|
|
B = x.shape[0] |
|
x = rearrange(x, "B C T H W -> (B T) C H W") |
|
|
|
if self.micro_batch_size is None: |
|
x = self.module.encode(x).latent_dist.sample().mul_(0.18215) |
|
else: |
|
bs = self.micro_batch_size |
|
x_out = [] |
|
for i in range(0, x.shape[0], bs): |
|
x_bs = x[i : i + bs] |
|
x_bs = self.module.encode(x_bs).latent_dist.sample().mul_(0.18215) |
|
x_out.append(x_bs) |
|
x = torch.cat(x_out, dim=0) |
|
x = rearrange(x, "(B T) C H W -> B C T H W", B=B) |
|
return x |
|
|
|
def decode(self, x): |
|
|
|
B = x.shape[0] |
|
x = rearrange(x, "B C T H W -> (B T) C H W") |
|
if self.micro_batch_size is None: |
|
x = self.module.decode(x / 0.18215).sample |
|
else: |
|
bs = self.micro_batch_size |
|
x_out = [] |
|
for i in range(0, x.shape[0], bs): |
|
x_bs = x[i : i + bs] |
|
x_bs = self.module.decode(x_bs / 0.18215).sample |
|
x_out.append(x_bs) |
|
x = torch.cat(x_out, dim=0) |
|
x = rearrange(x, "(B T) C H W -> B C T H W", B=B) |
|
return x |
|
|
|
def get_latent_size(self, input_size): |
|
for i in range(3): |
|
assert input_size[i] % self.patch_size[i] == 0, "Input size must be divisible by patch size" |
|
input_size = [input_size[i] // self.patch_size[i] for i in range(3)] |
|
return input_size |
|
|
|
|
|
@MODELS.register_module() |
|
class VideoAutoencoderKLTemporalDecoder(nn.Module): |
|
def __init__(self, from_pretrained=None): |
|
super().__init__() |
|
self.module = AutoencoderKLTemporalDecoder.from_pretrained(from_pretrained) |
|
self.out_channels = self.module.config.latent_channels |
|
self.patch_size = (1, 8, 8) |
|
|
|
def encode(self, x): |
|
raise NotImplementedError |
|
|
|
def decode(self, x): |
|
B, _, T = x.shape[:3] |
|
x = rearrange(x, "B C T H W -> (B T) C H W") |
|
x = self.module.decode(x / 0.18215, num_frames=T).sample |
|
x = rearrange(x, "(B T) C H W -> B C T H W", B=B) |
|
return x |
|
|
|
def get_latent_size(self, input_size): |
|
for i in range(3): |
|
assert input_size[i] % self.patch_size[i] == 0, "Input size must be divisible by patch size" |
|
input_size = [input_size[i] // self.patch_size[i] for i in range(3)] |
|
return input_size |
|
|
