56 lines
2.0 KiB
Python
56 lines
2.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from diffusers import AutoencoderKL as _DiffusersAutoencoderKL
|
|
|
|
from fastvideo.configs.models import VAEConfig
|
|
|
|
|
|
class AutoencoderKL(_DiffusersAutoencoderKL):
|
|
|
|
def __init__(self, config: VAEConfig, **kwargs: Any) -> None:
|
|
self.fastvideo_config = config
|
|
arch = config.arch_config
|
|
|
|
down_block_types = arch.down_block_types
|
|
if isinstance(down_block_types, list):
|
|
down_block_types = tuple(down_block_types)
|
|
up_block_types = arch.up_block_types
|
|
if isinstance(up_block_types, list):
|
|
up_block_types = tuple(up_block_types)
|
|
block_out_channels = arch.block_out_channels
|
|
if isinstance(block_out_channels, list):
|
|
block_out_channels = tuple(block_out_channels)
|
|
|
|
latents_mean = arch.latents_mean
|
|
if isinstance(latents_mean, list):
|
|
latents_mean = tuple(latents_mean)
|
|
latents_std = arch.latents_std
|
|
if isinstance(latents_std, list):
|
|
latents_std = tuple(latents_std)
|
|
|
|
super().__init__(
|
|
in_channels=arch.in_channels,
|
|
out_channels=arch.out_channels,
|
|
down_block_types=down_block_types,
|
|
up_block_types=up_block_types,
|
|
block_out_channels=block_out_channels,
|
|
layers_per_block=arch.layers_per_block,
|
|
act_fn=arch.act_fn,
|
|
latent_channels=arch.latent_channels,
|
|
norm_num_groups=arch.norm_num_groups,
|
|
sample_size=arch.sample_size,
|
|
scaling_factor=arch.scaling_factor,
|
|
shift_factor=arch.shift_factor,
|
|
latents_mean=latents_mean,
|
|
latents_std=latents_std,
|
|
force_upcast=arch.force_upcast,
|
|
use_quant_conv=arch.use_quant_conv,
|
|
use_post_quant_conv=arch.use_post_quant_conv,
|
|
mid_block_add_attention=arch.mid_block_add_attention,
|
|
**kwargs,
|
|
)
|