Spaces:
Running on Zero
Running on Zero
| # Copyright 2025 The Hunyuan Team and The HuggingFace Team. All rights reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import torch.utils.checkpoint | |
| from ...configuration_utils import ConfigMixin, register_to_config | |
| from ...loaders import FromOriginalModelMixin | |
| from ...utils import logging | |
| from ...utils.accelerate_utils import apply_forward_hook | |
| from ..activations import get_activation | |
| from ..modeling_outputs import AutoencoderKLOutput | |
| from ..modeling_utils import ModelMixin | |
| from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution | |
| logger = logging.get_logger(__name__) # pylint: disable=invalid-name | |
| class HunyuanImageResnetBlock(nn.Module): | |
| r""" | |
| Residual block with two convolutions and optional channel change. | |
| Args: | |
| in_channels (int): Number of input channels. | |
| out_channels (int): Number of output channels. | |
| non_linearity (str, optional): Type of non-linearity to use. Default is "silu". | |
| """ | |
| def __init__(self, in_channels: int, out_channels: int, non_linearity: str = "silu") -> None: | |
| super().__init__() | |
| self.in_channels = in_channels | |
| self.out_channels = out_channels | |
| self.nonlinearity = get_activation(non_linearity) | |
| # layers | |
| self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) | |
| self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) | |
| self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True) | |
| self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) | |
| if in_channels != out_channels: | |
| self.conv_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) | |
| else: | |
| self.conv_shortcut = None | |
| def forward(self, x): | |
| # Apply shortcut connection | |
| residual = x | |
| # First normalization and activation | |
| x = self.norm1(x) | |
| x = self.nonlinearity(x) | |
| x = self.conv1(x) | |
| x = self.norm2(x) | |
| x = self.nonlinearity(x) | |
| x = self.conv2(x) | |
| if self.conv_shortcut is not None: | |
| x = self.conv_shortcut(x) | |
| # Add residual connection | |
| return x + residual | |
| class HunyuanImageAttentionBlock(nn.Module): | |
| r""" | |
| Self-attention with a single head. | |
| Args: | |
| in_channels (int): The number of channels in the input tensor. | |
| """ | |
| def __init__(self, in_channels: int): | |
| super().__init__() | |
| # layers | |
| self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) | |
| self.to_q = nn.Conv2d(in_channels, in_channels, 1) | |
| self.to_k = nn.Conv2d(in_channels, in_channels, 1) | |
| self.to_v = nn.Conv2d(in_channels, in_channels, 1) | |
| self.proj = nn.Conv2d(in_channels, in_channels, 1) | |
| def forward(self, x): | |
| identity = x | |
| x = self.norm(x) | |
| # compute query, key, value | |
| query = self.to_q(x) | |
| key = self.to_k(x) | |
| value = self.to_v(x) | |
| batch_size, channels, height, width = query.shape | |
| query = query.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels).contiguous() | |
| key = key.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels).contiguous() | |
| value = value.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels).contiguous() | |
| # apply attention | |
| x = F.scaled_dot_product_attention(query, key, value) | |
| x = x.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2) | |
| # output projection | |
| x = self.proj(x) | |
| return x + identity | |
| class HunyuanImageDownsample(nn.Module): | |
| """ | |
| Downsampling block for spatial reduction. | |
| Args: | |
| in_channels (int): Number of input channels. | |
| out_channels (int): Number of output channels. | |
| """ | |
| def __init__(self, in_channels: int, out_channels: int): | |
| super().__init__() | |
| factor = 4 | |
| if out_channels % factor != 0: | |
| raise ValueError(f"out_channels % factor != 0: {out_channels % factor}") | |
| self.conv = nn.Conv2d(in_channels, out_channels // factor, kernel_size=3, stride=1, padding=1) | |
| self.group_size = factor * in_channels // out_channels | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| h = self.conv(x) | |
| B, C, H, W = h.shape | |
| h = h.reshape(B, C, H // 2, 2, W // 2, 2) | |
| h = h.permute(0, 3, 5, 1, 2, 4) # b, r1, r2, c, h, w | |
| h = h.reshape(B, 4 * C, H // 2, W // 2) | |
| B, C, H, W = x.shape | |
| shortcut = x.reshape(B, C, H // 2, 2, W // 2, 2) | |
| shortcut = shortcut.permute(0, 3, 5, 1, 2, 4) # b, r1, r2, c, h, w | |
| shortcut = shortcut.reshape(B, 4 * C, H // 2, W // 2) | |
| B, C, H, W = shortcut.shape | |
| shortcut = shortcut.view(B, h.shape[1], self.group_size, H, W).mean(dim=2) | |
| return h + shortcut | |
| class HunyuanImageUpsample(nn.Module): | |
| """ | |
| Upsampling block for spatial expansion. | |
| Args: | |
| in_channels (int): Number of input channels. | |
| out_channels (int): Number of output channels. | |
| """ | |
| def __init__(self, in_channels: int, out_channels: int): | |
| super().__init__() | |
| factor = 4 | |
| self.conv = nn.Conv2d(in_channels, out_channels * factor, kernel_size=3, stride=1, padding=1) | |
| self.repeats = factor * out_channels // in_channels | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| h = self.conv(x) | |
| B, C, H, W = h.shape | |
| h = h.reshape(B, 2, 2, C // 4, H, W) # b, r1, r2, c, h, w | |
| h = h.permute(0, 3, 4, 1, 5, 2) # b, c, h, r1, w, r2 | |
| h = h.reshape(B, C // 4, H * 2, W * 2) | |
| shortcut = x.repeat_interleave(repeats=self.repeats, dim=1) | |
| B, C, H, W = shortcut.shape | |
| shortcut = shortcut.reshape(B, 2, 2, C // 4, H, W) # b, r1, r2, c, h, w | |
| shortcut = shortcut.permute(0, 3, 4, 1, 5, 2) # b, c, h, r1, w, r2 | |
| shortcut = shortcut.reshape(B, C // 4, H * 2, W * 2) | |
| return h + shortcut | |
| class HunyuanImageMidBlock(nn.Module): | |
| """ | |
| Middle block for HunyuanImageVAE encoder and decoder. | |
| Args: | |
| in_channels (int): Number of input channels. | |
| num_layers (int): Number of layers. | |
| """ | |
| def __init__(self, in_channels: int, num_layers: int = 1): | |
| super().__init__() | |
| resnets = [HunyuanImageResnetBlock(in_channels=in_channels, out_channels=in_channels)] | |
| attentions = [] | |
| for _ in range(num_layers): | |
| attentions.append(HunyuanImageAttentionBlock(in_channels)) | |
| resnets.append(HunyuanImageResnetBlock(in_channels=in_channels, out_channels=in_channels)) | |
| self.resnets = nn.ModuleList(resnets) | |
| self.attentions = nn.ModuleList(attentions) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| x = self.resnets[0](x) | |
| for attn, resnet in zip(self.attentions, self.resnets[1:]): | |
| x = attn(x) | |
| x = resnet(x) | |
| return x | |
| class HunyuanImageEncoder2D(nn.Module): | |
| r""" | |
| Encoder network that compresses input to latent representation. | |
| Args: | |
| in_channels (int): Number of input channels. | |
| z_channels (int): Number of latent channels. | |
| block_out_channels (list of int): Output channels for each block. | |
| num_res_blocks (int): Number of residual blocks per block. | |
| spatial_compression_ratio (int): Spatial downsampling factor. | |
| non_linearity (str): Type of non-linearity to use. Default is "silu". | |
| downsample_match_channel (bool): Whether to match channels during downsampling. | |
| """ | |
| def __init__( | |
| self, | |
| in_channels: int, | |
| z_channels: int, | |
| block_out_channels: tuple[int, ...], | |
| num_res_blocks: int, | |
| spatial_compression_ratio: int, | |
| non_linearity: str = "silu", | |
| downsample_match_channel: bool = True, | |
| ): | |
| super().__init__() | |
| if block_out_channels[-1] % (2 * z_channels) != 0: | |
| raise ValueError( | |
| f"block_out_channels[-1 has to be divisible by 2 * out_channels, you have block_out_channels = {block_out_channels[-1]} and out_channels = {z_channels}" | |
| ) | |
| self.in_channels = in_channels | |
| self.z_channels = z_channels | |
| self.block_out_channels = block_out_channels | |
| self.num_res_blocks = num_res_blocks | |
| self.spatial_compression_ratio = spatial_compression_ratio | |
| self.group_size = block_out_channels[-1] // (2 * z_channels) | |
| self.nonlinearity = get_activation(non_linearity) | |
| # init block | |
| self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1) | |
| # downsample blocks | |
| self.down_blocks = nn.ModuleList([]) | |
| block_in_channel = block_out_channels[0] | |
| for i in range(len(block_out_channels)): | |
| block_out_channel = block_out_channels[i] | |
| # residual blocks | |
| for _ in range(num_res_blocks): | |
| self.down_blocks.append( | |
| HunyuanImageResnetBlock(in_channels=block_in_channel, out_channels=block_out_channel) | |
| ) | |
| block_in_channel = block_out_channel | |
| # downsample block | |
| if i < np.log2(spatial_compression_ratio) and i != len(block_out_channels) - 1: | |
| if downsample_match_channel: | |
| block_out_channel = block_out_channels[i + 1] | |
| self.down_blocks.append( | |
| HunyuanImageDownsample(in_channels=block_in_channel, out_channels=block_out_channel) | |
| ) | |
| block_in_channel = block_out_channel | |
| # middle blocks | |
| self.mid_block = HunyuanImageMidBlock(in_channels=block_out_channels[-1], num_layers=1) | |
| # output blocks | |
| # Output layers | |
| self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_out_channels[-1], eps=1e-6, affine=True) | |
| self.conv_out = nn.Conv2d(block_out_channels[-1], 2 * z_channels, kernel_size=3, stride=1, padding=1) | |
| self.gradient_checkpointing = False | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| x = self.conv_in(x) | |
| ## downsamples | |
| for down_block in self.down_blocks: | |
| if torch.is_grad_enabled() and self.gradient_checkpointing: | |
| x = self._gradient_checkpointing_func(down_block, x) | |
| else: | |
| x = down_block(x) | |
| ## middle | |
| if torch.is_grad_enabled() and self.gradient_checkpointing: | |
| x = self._gradient_checkpointing_func(self.mid_block, x) | |
| else: | |
| x = self.mid_block(x) | |
| ## head | |
| B, C, H, W = x.shape | |
| residual = x.view(B, C // self.group_size, self.group_size, H, W).mean(dim=2) | |
| x = self.norm_out(x) | |
| x = self.nonlinearity(x) | |
| x = self.conv_out(x) | |
| return x + residual | |
| class HunyuanImageDecoder2D(nn.Module): | |
| r""" | |
| Decoder network that reconstructs output from latent representation. | |
| Args: | |
| z_channels : int | |
| Number of latent channels. | |
| out_channels : int | |
| Number of output channels. | |
| block_out_channels : tuple[int, ...] | |
| Output channels for each block. | |
| num_res_blocks : int | |
| Number of residual blocks per block. | |
| spatial_compression_ratio : int | |
| Spatial upsampling factor. | |
| upsample_match_channel : bool | |
| Whether to match channels during upsampling. | |
| non_linearity (str): Type of non-linearity to use. Default is "silu". | |
| """ | |
| def __init__( | |
| self, | |
| z_channels: int, | |
| out_channels: int, | |
| block_out_channels: tuple[int, ...], | |
| num_res_blocks: int, | |
| spatial_compression_ratio: int, | |
| upsample_match_channel: bool = True, | |
| non_linearity: str = "silu", | |
| ): | |
| super().__init__() | |
| if block_out_channels[0] % z_channels != 0: | |
| raise ValueError( | |
| f"block_out_channels[0] should be divisible by z_channels but has block_out_channels[0] = {block_out_channels[0]} and z_channels = {z_channels}" | |
| ) | |
| self.z_channels = z_channels | |
| self.block_out_channels = block_out_channels | |
| self.num_res_blocks = num_res_blocks | |
| self.repeat = block_out_channels[0] // z_channels | |
| self.spatial_compression_ratio = spatial_compression_ratio | |
| self.nonlinearity = get_activation(non_linearity) | |
| self.conv_in = nn.Conv2d(z_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1) | |
| # Middle blocks with attention | |
| self.mid_block = HunyuanImageMidBlock(in_channels=block_out_channels[0], num_layers=1) | |
| # Upsampling blocks | |
| block_in_channel = block_out_channels[0] | |
| self.up_blocks = nn.ModuleList() | |
| for i in range(len(block_out_channels)): | |
| block_out_channel = block_out_channels[i] | |
| for _ in range(self.num_res_blocks + 1): | |
| self.up_blocks.append( | |
| HunyuanImageResnetBlock(in_channels=block_in_channel, out_channels=block_out_channel) | |
| ) | |
| block_in_channel = block_out_channel | |
| if i < np.log2(spatial_compression_ratio) and i != len(block_out_channels) - 1: | |
| if upsample_match_channel: | |
| block_out_channel = block_out_channels[i + 1] | |
| self.up_blocks.append(HunyuanImageUpsample(block_in_channel, block_out_channel)) | |
| block_in_channel = block_out_channel | |
| # Output layers | |
| self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_out_channels[-1], eps=1e-6, affine=True) | |
| self.conv_out = nn.Conv2d(block_out_channels[-1], out_channels, kernel_size=3, stride=1, padding=1) | |
| self.gradient_checkpointing = False | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| h = self.conv_in(x) + x.repeat_interleave(repeats=self.repeat, dim=1) | |
| if torch.is_grad_enabled() and self.gradient_checkpointing: | |
| h = self._gradient_checkpointing_func(self.mid_block, h) | |
| else: | |
| h = self.mid_block(h) | |
| for up_block in self.up_blocks: | |
| if torch.is_grad_enabled() and self.gradient_checkpointing: | |
| h = self._gradient_checkpointing_func(up_block, h) | |
| else: | |
| h = up_block(h) | |
| h = self.norm_out(h) | |
| h = self.nonlinearity(h) | |
| h = self.conv_out(h) | |
| return h | |
| class AutoencoderKLHunyuanImage(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): | |
| r""" | |
| A VAE model for 2D images with spatial tiling support. | |
| This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented | |
| for all models (such as downloading or saving). | |
| """ | |
| _supports_gradient_checkpointing = False | |
| # fmt: off | |
| def __init__( | |
| self, | |
| in_channels: int, | |
| out_channels: int, | |
| latent_channels: int, | |
| block_out_channels: tuple[int, ...], | |
| layers_per_block: int, | |
| spatial_compression_ratio: int, | |
| sample_size: int, | |
| scaling_factor: float = None, | |
| downsample_match_channel: bool = True, | |
| upsample_match_channel: bool = True, | |
| ) -> None: | |
| # fmt: on | |
| super().__init__() | |
| self.encoder = HunyuanImageEncoder2D( | |
| in_channels=in_channels, | |
| z_channels=latent_channels, | |
| block_out_channels=block_out_channels, | |
| num_res_blocks=layers_per_block, | |
| spatial_compression_ratio=spatial_compression_ratio, | |
| downsample_match_channel=downsample_match_channel, | |
| ) | |
| self.decoder = HunyuanImageDecoder2D( | |
| z_channels=latent_channels, | |
| out_channels=out_channels, | |
| block_out_channels=list(reversed(block_out_channels)), | |
| num_res_blocks=layers_per_block, | |
| spatial_compression_ratio=spatial_compression_ratio, | |
| upsample_match_channel=upsample_match_channel, | |
| ) | |
| # Tiling and slicing configuration | |
| self.use_slicing = False | |
| self.use_tiling = False | |
| # Tiling parameters | |
| self.tile_sample_min_size = sample_size | |
| self.tile_latent_min_size = sample_size // spatial_compression_ratio | |
| self.tile_overlap_factor = 0.25 | |
| def enable_tiling( | |
| self, | |
| tile_sample_min_size: int | None = None, | |
| tile_overlap_factor: float | None = None, | |
| ) -> None: | |
| r""" | |
| Enable spatial tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles | |
| to compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to | |
| allow processing larger images. | |
| Args: | |
| tile_sample_min_size (`int`, *optional*): | |
| The minimum size required for a sample to be separated into tiles across the spatial dimension. | |
| tile_overlap_factor (`float`, *optional*): | |
| The overlap factor required for a latent to be separated into tiles across the spatial dimension. | |
| """ | |
| self.use_tiling = True | |
| self.tile_sample_min_size = tile_sample_min_size or self.tile_sample_min_size | |
| self.tile_overlap_factor = tile_overlap_factor or self.tile_overlap_factor | |
| self.tile_latent_min_size = self.tile_sample_min_size // self.config.spatial_compression_ratio | |
| def _encode(self, x: torch.Tensor): | |
| batch_size, num_channels, height, width = x.shape | |
| if self.use_tiling and (width > self.tile_sample_min_size or height > self.tile_sample_min_size): | |
| return self.tiled_encode(x) | |
| enc = self.encoder(x) | |
| return enc | |
| def encode( | |
| self, x: torch.Tensor, return_dict: bool = True | |
| ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: | |
| r""" | |
| Encode a batch of images into latents. | |
| Args: | |
| x (`torch.Tensor`): Input batch of images. | |
| return_dict (`bool`, *optional*, defaults to `True`): | |
| Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. | |
| Returns: | |
| The latent representations of the encoded videos. If `return_dict` is True, a | |
| [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. | |
| """ | |
| if self.use_slicing and x.shape[0] > 1: | |
| encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)] | |
| h = torch.cat(encoded_slices) | |
| else: | |
| h = self._encode(x) | |
| posterior = DiagonalGaussianDistribution(h) | |
| if not return_dict: | |
| return (posterior,) | |
| return AutoencoderKLOutput(latent_dist=posterior) | |
| def _decode(self, z: torch.Tensor, return_dict: bool = True): | |
| batch_size, num_channels, height, width = z.shape | |
| if self.use_tiling and (width > self.tile_latent_min_size or height > self.tile_latent_min_size): | |
| return self.tiled_decode(z, return_dict=return_dict) | |
| dec = self.decoder(z) | |
| if not return_dict: | |
| return (dec,) | |
| return DecoderOutput(sample=dec) | |
| def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: | |
| r""" | |
| Decode a batch of images. | |
| Args: | |
| z (`torch.Tensor`): Input batch of latent vectors. | |
| return_dict (`bool`, *optional*, defaults to `True`): | |
| Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. | |
| Returns: | |
| [`~models.vae.DecoderOutput`] or `tuple`: | |
| If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is | |
| returned. | |
| """ | |
| if self.use_slicing and z.shape[0] > 1: | |
| decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] | |
| decoded = torch.cat(decoded_slices) | |
| else: | |
| decoded = self._decode(z).sample | |
| if not return_dict: | |
| return (decoded,) | |
| return DecoderOutput(sample=decoded) | |
| def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: | |
| blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) | |
| for y in range(blend_extent): | |
| b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * ( | |
| y / blend_extent | |
| ) | |
| return b | |
| def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: | |
| blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) | |
| for x in range(blend_extent): | |
| b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * ( | |
| x / blend_extent | |
| ) | |
| return b | |
| def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: | |
| """ | |
| Encode input using spatial tiling strategy. | |
| Args: | |
| x (`torch.Tensor`): Input tensor of shape (B, C, T, H, W). | |
| Returns: | |
| `torch.Tensor`: | |
| The latent representation of the encoded images. | |
| """ | |
| _, _, _, height, width = x.shape | |
| overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) | |
| blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) | |
| row_limit = self.tile_latent_min_size - blend_extent | |
| rows = [] | |
| for i in range(0, height, overlap_size): | |
| row = [] | |
| for j in range(0, width, overlap_size): | |
| tile = x[:, :, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size] | |
| tile = self.encoder(tile) | |
| row.append(tile) | |
| rows.append(row) | |
| result_rows = [] | |
| for i, row in enumerate(rows): | |
| result_row = [] | |
| for j, tile in enumerate(row): | |
| if i > 0: | |
| tile = self.blend_v(rows[i - 1][j], tile, blend_extent) | |
| if j > 0: | |
| tile = self.blend_h(row[j - 1], tile, blend_extent) | |
| result_row.append(tile[:, :, :, :row_limit, :row_limit]) | |
| result_rows.append(torch.cat(result_row, dim=-1)) | |
| moments = torch.cat(result_rows, dim=-2) | |
| return moments | |
| def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: | |
| """ | |
| Decode latent using spatial tiling strategy. | |
| Args: | |
| z (`torch.Tensor`): Latent tensor of shape (B, C, H, W). | |
| return_dict (`bool`, *optional*, defaults to `True`): | |
| Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. | |
| Returns: | |
| [`~models.vae.DecoderOutput`] or `tuple`: | |
| If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is | |
| returned. | |
| """ | |
| _, _, height, width = z.shape | |
| overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) | |
| blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) | |
| row_limit = self.tile_sample_min_size - blend_extent | |
| rows = [] | |
| for i in range(0, height, overlap_size): | |
| row = [] | |
| for j in range(0, width, overlap_size): | |
| tile = z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size] | |
| decoded = self.decoder(tile) | |
| row.append(decoded) | |
| rows.append(row) | |
| result_rows = [] | |
| for i, row in enumerate(rows): | |
| result_row = [] | |
| for j, tile in enumerate(row): | |
| if i > 0: | |
| tile = self.blend_v(rows[i - 1][j], tile, blend_extent) | |
| if j > 0: | |
| tile = self.blend_h(row[j - 1], tile, blend_extent) | |
| result_row.append(tile[:, :, :row_limit, :row_limit]) | |
| result_rows.append(torch.cat(result_row, dim=-1)) | |
| dec = torch.cat(result_rows, dim=-2) | |
| if not return_dict: | |
| return (dec,) | |
| return DecoderOutput(sample=dec) | |
| def forward( | |
| self, | |
| sample: torch.Tensor, | |
| sample_posterior: bool = False, | |
| return_dict: bool = True, | |
| generator: torch.Generator | None = None, | |
| ) -> DecoderOutput | torch.Tensor: | |
| """ | |
| Args: | |
| sample (`torch.Tensor`): Input sample. | |
| sample_posterior (`bool`, *optional*, defaults to `False`): | |
| Whether to sample from the posterior. | |
| return_dict (`bool`, *optional*, defaults to `True`): | |
| Whether or not to return a [`DecoderOutput`] instead of a plain tuple. | |
| generator (`torch.Generator`, *optional*): | |
| A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling | |
| deterministic. | |
| Returns: | |
| [`~models.vae.DecoderOutput`] or `tuple`: | |
| If `return_dict` is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is | |
| returned. | |
| """ | |
| posterior = self.encode(sample).latent_dist | |
| if sample_posterior: | |
| z = posterior.sample(generator=generator) | |
| else: | |
| z = posterior.mode() | |
| dec = self.decode(z, return_dict=return_dict) | |
| return dec | |