




from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import torch

from ...utils import BaseOutput, OptionalDependencyNotAvailable, is_torch_available, is_transformers_available





@dataclass
class 			__lowerCamelCase	(   UpperCamelCase__   ):




    """simple docstring"""

    snake_case__ =				42


try:
     if not (is_transformers_available() and is_torch_available()):
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
     from .pipeline_text_to_video_synth import TextToVideoSDPipeline
     from .pipeline_text_to_video_synth_imgaimg import VideoToVideoSDPipeline  # noqa: F401
     from .pipeline_text_to_video_zero import TextToVideoZeroPipeline






from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block





@dataclass
class 			__lowerCamelCase	(   UpperCamelCase__   ):




    """simple docstring"""

    snake_case__ =				42






class 			__lowerCamelCase	(   nn.Module   ):




    """simple docstring"""

    def __init__(	self					:					List[Any]       ,     SCREAMING_SNAKE_CASE__					:					Any=3       ,     SCREAMING_SNAKE_CASE__					:					int=3       ,     SCREAMING_SNAKE_CASE__					:					List[str]=("DownEncoderBlock2D",)       ,     SCREAMING_SNAKE_CASE__					:					Tuple=(64,)       ,     SCREAMING_SNAKE_CASE__					:					str=2       ,     SCREAMING_SNAKE_CASE__					:					List[Any]=32       ,     SCREAMING_SNAKE_CASE__					:					List[str]="silu"       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]=True       ,     )				->	str:
         super().__init__()
         lowerCAmelCase__          =			layers_per_block

         lowerCAmelCase__          =			torch.nn.Convad(
             SCREAMING_SNAKE_CASE__       ,     block_out_channels[0]       ,     kernel_size=3       ,     stride=1       ,     padding=1       ,     )

         lowerCAmelCase__          =			None
         lowerCAmelCase__          =			nn.ModuleList([]       )

         # down
         lowerCAmelCase__          =			block_out_channels[0]
         for i, down_block_type in enumerate(SCREAMING_SNAKE_CASE__       ):
              lowerCAmelCase__          =			output_channel
              lowerCAmelCase__          =			block_out_channels[i]
              lowerCAmelCase__          =			i == len(SCREAMING_SNAKE_CASE__       ) - 1

              lowerCAmelCase__          =			get_down_block(
                  SCREAMING_SNAKE_CASE__       ,     num_layers=self.layers_per_block       ,     in_channels=SCREAMING_SNAKE_CASE__       ,     out_channels=SCREAMING_SNAKE_CASE__       ,     add_downsample=not is_final_block       ,     resnet_eps=1e-6       ,     downsample_padding=0       ,     resnet_act_fn=SCREAMING_SNAKE_CASE__       ,     resnet_groups=SCREAMING_SNAKE_CASE__       ,     attention_head_dim=SCREAMING_SNAKE_CASE__       ,     temb_channels=SCREAMING_SNAKE_CASE__       ,     )
              self.down_blocks.append(SCREAMING_SNAKE_CASE__       )

         # mid
         lowerCAmelCase__          =			UNetMidBlockaD(
             in_channels=block_out_channels[-1]       ,     resnet_eps=1e-6       ,     resnet_act_fn=SCREAMING_SNAKE_CASE__       ,     output_scale_factor=1       ,     resnet_time_scale_shift="default"       ,     attention_head_dim=block_out_channels[-1]       ,     resnet_groups=SCREAMING_SNAKE_CASE__       ,     temb_channels=SCREAMING_SNAKE_CASE__       ,     )

         # out
         lowerCAmelCase__          =			nn.GroupNorm(num_channels=block_out_channels[-1]       ,     num_groups=SCREAMING_SNAKE_CASE__       ,     eps=1e-6       )
         lowerCAmelCase__          =			nn.SiLU()

         lowerCAmelCase__          =			2 * out_channels if double_z else out_channels
         lowerCAmelCase__          =			nn.Convad(block_out_channels[-1]       ,     SCREAMING_SNAKE_CASE__       ,     3       ,     padding=1       )

         lowerCAmelCase__          =			False


    def        a				(	self					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					Dict       )				->	int:
         lowerCAmelCase__          =			x
         lowerCAmelCase__          =			self.conv_in(SCREAMING_SNAKE_CASE__       )

         if self.training and self.gradient_checkpointing:

              def create_custom_forward(SCREAMING_SNAKE_CASE__					:					Optional[int]       ):
                   def custom_forward(*SCREAMING_SNAKE_CASE__					:					Dict       ):
                        return module(*SCREAMING_SNAKE_CASE__       )

                   return custom_forward

              # down
              if is_torch_version(">="       ,     "1.11.0"       ):
                   for down_block in self.down_blocks:
                        lowerCAmelCase__          =			torch.utils.checkpoint.checkpoint(
                            create_custom_forward(SCREAMING_SNAKE_CASE__       )       ,     SCREAMING_SNAKE_CASE__       ,     use_reentrant=SCREAMING_SNAKE_CASE__       )
                   # middle
                   lowerCAmelCase__          =			torch.utils.checkpoint.checkpoint(
                       create_custom_forward(self.mid_block       )       ,     SCREAMING_SNAKE_CASE__       ,     use_reentrant=SCREAMING_SNAKE_CASE__       )
              else:
                   for down_block in self.down_blocks:
                        lowerCAmelCase__          =			torch.utils.checkpoint.checkpoint(create_custom_forward(SCREAMING_SNAKE_CASE__       )       ,     SCREAMING_SNAKE_CASE__       )
                   # middle
                   lowerCAmelCase__          =			torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block       )       ,     SCREAMING_SNAKE_CASE__       )

         else:
              # down
              for down_block in self.down_blocks:
                   lowerCAmelCase__          =			down_block(SCREAMING_SNAKE_CASE__       )

              # middle
              lowerCAmelCase__          =			self.mid_block(SCREAMING_SNAKE_CASE__       )

         # post-process
         lowerCAmelCase__          =			self.conv_norm_out(SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			self.conv_act(SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			self.conv_out(SCREAMING_SNAKE_CASE__       )

         return sample






class 			__lowerCamelCase	(   nn.Module   ):




    """simple docstring"""

    def __init__(	self					:					Any       ,     SCREAMING_SNAKE_CASE__					:					str=3       ,     SCREAMING_SNAKE_CASE__					:					Dict=3       ,     SCREAMING_SNAKE_CASE__					:					int=("UpDecoderBlock2D",)       ,     SCREAMING_SNAKE_CASE__					:					Tuple=(64,)       ,     SCREAMING_SNAKE_CASE__					:					Tuple=2       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]=32       ,     SCREAMING_SNAKE_CASE__					:					List[str]="silu"       ,     SCREAMING_SNAKE_CASE__					:					List[str]="group"       ,     )				->	Union[str, Any]:
         super().__init__()
         lowerCAmelCase__          =			layers_per_block

         lowerCAmelCase__          =			nn.Convad(
             SCREAMING_SNAKE_CASE__       ,     block_out_channels[-1]       ,     kernel_size=3       ,     stride=1       ,     padding=1       ,     )

         lowerCAmelCase__          =			None
         lowerCAmelCase__          =			nn.ModuleList([]       )

         lowerCAmelCase__          =			in_channels if norm_type == "spatial" else None

         # mid
         lowerCAmelCase__          =			UNetMidBlockaD(
             in_channels=block_out_channels[-1]       ,     resnet_eps=1e-6       ,     resnet_act_fn=SCREAMING_SNAKE_CASE__       ,     output_scale_factor=1       ,     resnet_time_scale_shift="default" if norm_type == "group" else norm_type       ,     attention_head_dim=block_out_channels[-1]       ,     resnet_groups=SCREAMING_SNAKE_CASE__       ,     temb_channels=SCREAMING_SNAKE_CASE__       ,     )

         # up
         lowerCAmelCase__          =			list(reversed(SCREAMING_SNAKE_CASE__       )       )
         lowerCAmelCase__          =			reversed_block_out_channels[0]
         for i, up_block_type in enumerate(SCREAMING_SNAKE_CASE__       ):
              lowerCAmelCase__          =			output_channel
              lowerCAmelCase__          =			reversed_block_out_channels[i]

              lowerCAmelCase__          =			i == len(SCREAMING_SNAKE_CASE__       ) - 1

              lowerCAmelCase__          =			get_up_block(
                  SCREAMING_SNAKE_CASE__       ,     num_layers=self.layers_per_block + 1       ,     in_channels=SCREAMING_SNAKE_CASE__       ,     out_channels=SCREAMING_SNAKE_CASE__       ,     prev_output_channel=SCREAMING_SNAKE_CASE__       ,     add_upsample=not is_final_block       ,     resnet_eps=1e-6       ,     resnet_act_fn=SCREAMING_SNAKE_CASE__       ,     resnet_groups=SCREAMING_SNAKE_CASE__       ,     attention_head_dim=SCREAMING_SNAKE_CASE__       ,     temb_channels=SCREAMING_SNAKE_CASE__       ,     resnet_time_scale_shift=SCREAMING_SNAKE_CASE__       ,     )
              self.up_blocks.append(SCREAMING_SNAKE_CASE__       )
              lowerCAmelCase__          =			output_channel

         # out
         if norm_type == "spatial":
              lowerCAmelCase__          =			SpatialNorm(block_out_channels[0]       ,     SCREAMING_SNAKE_CASE__       )
         else:
              lowerCAmelCase__          =			nn.GroupNorm(num_channels=block_out_channels[0]       ,     num_groups=SCREAMING_SNAKE_CASE__       ,     eps=1e-6       )
         lowerCAmelCase__          =			nn.SiLU()
         lowerCAmelCase__          =			nn.Convad(block_out_channels[0]       ,     SCREAMING_SNAKE_CASE__       ,     3       ,     padding=1       )

         lowerCAmelCase__          =			False


    def        a				(	self					:					Any       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]       ,     SCREAMING_SNAKE_CASE__					:					Dict=None       )				->	Dict:
         lowerCAmelCase__          =			z
         lowerCAmelCase__          =			self.conv_in(SCREAMING_SNAKE_CASE__       )

         lowerCAmelCase__          =			next(iter(self.up_blocks.parameters()       )       ).dtype
         if self.training and self.gradient_checkpointing:

              def create_custom_forward(SCREAMING_SNAKE_CASE__					:					Union[str, Any]       ):
                   def custom_forward(*SCREAMING_SNAKE_CASE__					:					List[str]       ):
                        return module(*SCREAMING_SNAKE_CASE__       )

                   return custom_forward

              if is_torch_version(">="       ,     "1.11.0"       ):
                   # middle
                   lowerCAmelCase__          =			torch.utils.checkpoint.checkpoint(
                       create_custom_forward(self.mid_block       )       ,     SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     use_reentrant=SCREAMING_SNAKE_CASE__       )
                   lowerCAmelCase__          =			sample.to(SCREAMING_SNAKE_CASE__       )

                   # up
                   for up_block in self.up_blocks:
                        lowerCAmelCase__          =			torch.utils.checkpoint.checkpoint(
                            create_custom_forward(SCREAMING_SNAKE_CASE__       )       ,     SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     use_reentrant=SCREAMING_SNAKE_CASE__       )
              else:
                   # middle
                   lowerCAmelCase__          =			torch.utils.checkpoint.checkpoint(
                       create_custom_forward(self.mid_block       )       ,     SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
                   lowerCAmelCase__          =			sample.to(SCREAMING_SNAKE_CASE__       )

                   # up
                   for up_block in self.up_blocks:
                        lowerCAmelCase__          =			torch.utils.checkpoint.checkpoint(create_custom_forward(SCREAMING_SNAKE_CASE__       )       ,     SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
         else:
              # middle
              lowerCAmelCase__          =			self.mid_block(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
              lowerCAmelCase__          =			sample.to(SCREAMING_SNAKE_CASE__       )

              # up
              for up_block in self.up_blocks:
                   lowerCAmelCase__          =			up_block(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )

        # post-process
         if latent_embeds is None:
              lowerCAmelCase__          =			self.conv_norm_out(SCREAMING_SNAKE_CASE__       )
         else:
              lowerCAmelCase__          =			self.conv_norm_out(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			self.conv_act(SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			self.conv_out(SCREAMING_SNAKE_CASE__       )

         return sample






class 			__lowerCamelCase	(   nn.Module   ):




    """simple docstring"""

    def __init__(	self					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					Any       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]=None       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]="random"       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]=False       ,     SCREAMING_SNAKE_CASE__					:					Any=True       )				->	List[Any]:
         super().__init__()
         lowerCAmelCase__          =			n_e
         lowerCAmelCase__          =			vq_embed_dim
         lowerCAmelCase__          =			beta
         lowerCAmelCase__          =			legacy

         lowerCAmelCase__          =			nn.Embedding(self.n_e       ,     self.vq_embed_dim       )
         self.embedding.weight.data.uniform_(-1.0 / self.n_e       ,     1.0 / self.n_e       )

         lowerCAmelCase__          =			remap
         if self.remap is not None:
              self.register_buffer("used"       ,     torch.tensor(np.load(self.remap       )       )       )
              lowerCAmelCase__          =			self.used.shape[0]
              lowerCAmelCase__          =			unknown_index  # "random" or "extra" or integer
              if self.unknown_index == "extra":
                   lowerCAmelCase__          =			self.re_embed
                   lowerCAmelCase__          =			self.re_embed + 1
              print(
                  f'Remapping {self.n_e} indices to {self.re_embed} indices. '
                  f'Using {self.unknown_index} for unknown indices.'       )
         else:
              lowerCAmelCase__          =			n_e

         lowerCAmelCase__          =			sane_index_shape


    def        a				(	self					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]       )				->	Optional[int]:
         lowerCAmelCase__          =			inds.shape
         assert len(SCREAMING_SNAKE_CASE__       ) > 1
         lowerCAmelCase__          =			inds.reshape(ishape[0]       ,     -1       )
         lowerCAmelCase__          =			self.used.to(SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			(inds[:, :, None] == used[None, None, ...]).long()
         lowerCAmelCase__          =			match.argmax(-1       )
         lowerCAmelCase__          =			match.sum(2       ) < 1
         if self.unknown_index == "random":
              lowerCAmelCase__          =			torch.randint(0       ,     self.re_embed       ,     size=new[unknown].shape       ).to(device=new.device       )
         else:
              lowerCAmelCase__          =			self.unknown_index
         return new.reshape(SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					Optional[Any]       ,     SCREAMING_SNAKE_CASE__					:					Dict       )				->	Optional[int]:
         lowerCAmelCase__          =			inds.shape
         assert len(SCREAMING_SNAKE_CASE__       ) > 1
         lowerCAmelCase__          =			inds.reshape(ishape[0]       ,     -1       )
         lowerCAmelCase__          =			self.used.to(SCREAMING_SNAKE_CASE__       )
         if self.re_embed > self.used.shape[0]:  # extra token
              lowerCAmelCase__          =			0  # simply set to zero
         lowerCAmelCase__          =			torch.gather(used[None, :][inds.shape[0] * [0], :]       ,     1       ,     SCREAMING_SNAKE_CASE__       )
         return back.reshape(SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					List[Any]       ,     SCREAMING_SNAKE_CASE__					:					Tuple       )				->	Tuple:
         # reshape z -> (batch, height, width, channel) and flatten
         lowerCAmelCase__          =			z.permute(0       ,     2       ,     3       ,     1       ).contiguous()
         lowerCAmelCase__          =			z.view(-1       ,     self.vq_embed_dim       )

         # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
         lowerCAmelCase__          =			torch.argmin(torch.cdist(SCREAMING_SNAKE_CASE__       ,     self.embedding.weight       )       ,     dim=1       )

         lowerCAmelCase__          =			self.embedding(SCREAMING_SNAKE_CASE__       ).view(z.shape       )
         lowerCAmelCase__          =			None
         lowerCAmelCase__          =			None

         # compute loss for embedding
         if not self.legacy:
              lowerCAmelCase__          =			self.beta * torch.mean((z_q.detach() - z) ** 2       ) + torch.mean((z_q - z.detach()) ** 2       )
         else:
              lowerCAmelCase__          =			torch.mean((z_q.detach() - z) ** 2       ) + self.beta * torch.mean((z_q - z.detach()) ** 2       )

         # preserve gradients
         lowerCAmelCase__          =			z + (z_q - z).detach()

         # reshape back to match original input shape
         lowerCAmelCase__          =			z_q.permute(0       ,     3       ,     1       ,     2       ).contiguous()

         if self.remap is not None:
              lowerCAmelCase__          =			min_encoding_indices.reshape(z.shape[0]       ,     -1       )  # add batch axis
              lowerCAmelCase__          =			self.remap_to_used(SCREAMING_SNAKE_CASE__       )
              lowerCAmelCase__          =			min_encoding_indices.reshape(-1       ,     1       )  # flatten

         if self.sane_index_shape:
              lowerCAmelCase__          =			min_encoding_indices.reshape(z_q.shape[0]       ,     z_q.shape[2]       ,     z_q.shape[3]       )

         return z_q, loss, (perplexity, min_encodings, min_encoding_indices)


    def        a				(	self					:					str       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					str       )				->	Any:
         # shape specifying (batch, height, width, channel)
         if self.remap is not None:
              lowerCAmelCase__          =			indices.reshape(shape[0]       ,     -1       )  # add batch axis
              lowerCAmelCase__          =			self.unmap_to_all(SCREAMING_SNAKE_CASE__       )
              lowerCAmelCase__          =			indices.reshape(-1       )  # flatten again

         # get quantized latent vectors
         lowerCAmelCase__          =			self.embedding(SCREAMING_SNAKE_CASE__       )

         if shape is not None:
              lowerCAmelCase__          =			z_q.view(SCREAMING_SNAKE_CASE__       )
              # reshape back to match original input shape
              lowerCAmelCase__          =			z_q.permute(0       ,     3       ,     1       ,     2       ).contiguous()

         return z_q






class 			__lowerCamelCase	(   UpperCamelCase__   ):




    """simple docstring"""

    def __init__(	self					:					Optional[Any]       ,     SCREAMING_SNAKE_CASE__					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					Any=False       )				->	Optional[int]:
         lowerCAmelCase__          =			parameters
         lowerCAmelCase__     ,		lowerCAmelCase__          =			torch.chunk(SCREAMING_SNAKE_CASE__       ,     2       ,     dim=1       )
         lowerCAmelCase__          =			torch.clamp(self.logvar       ,     -30.0       ,     20.0       )
         lowerCAmelCase__          =			deterministic
         lowerCAmelCase__          =			torch.exp(0.5 * self.logvar       )
         lowerCAmelCase__          =			torch.exp(self.logvar       )
         if self.deterministic:
              lowerCAmelCase__          =			lowerCAmelCase__          =			torch.zeros_like(
                  self.mean       ,     device=self.parameters.device       ,     dtype=self.parameters.dtype       )


    def        a				(	self					:					str       ,     SCREAMING_SNAKE_CASE__					:					Optional[torch.Generator] = None       )				->	torch.FloatTensor:
         # make sure sample is on the same device as the parameters and has same dtype
         lowerCAmelCase__          =			randn_tensor(
             self.mean.shape       ,     generator=SCREAMING_SNAKE_CASE__       ,     device=self.parameters.device       ,     dtype=self.parameters.dtype       )
         lowerCAmelCase__          =			self.mean + self.std * sample
         return x


    def        a				(	self					:					List[str]       ,     SCREAMING_SNAKE_CASE__					:					List[Any]=None       )				->	Union[str, Any]:
         if self.deterministic:
              return torch.Tensor([0.0]       )
         else:
              if other is None:
                   return 0.5 * torch.sum(torch.pow(self.mean       ,     2       ) + self.var - 1.0 - self.logvar       ,     dim=[1, 2, 3]       )
              else:
                   return 0.5 * torch.sum(
                       torch.pow(self.mean - other.mean       ,     2       ) / other.var
                       + self.var / other.var
                       - 1.0
                       - self.logvar
                       + other.logvar       ,     dim=[1, 2, 3]       ,     )


    def        a				(	self					:					Any       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]       ,     SCREAMING_SNAKE_CASE__					:					List[str]=[1, 2, 3]       )				->	Tuple:
         if self.deterministic:
              return torch.Tensor([0.0]       )
         lowerCAmelCase__          =			np.log(2.0 * np.pi       )
         return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean       ,     2       ) / self.var       ,     dim=SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					Union[str, Any]       )				->	Dict:
         return self.mean
