


import flax.linen as nn
import jax
import jax.numpy as jnp

class 		__a				(    nn.Module      ):
		_a     :  int
		_a     :  jnp.dtype           =     jnp.floataa
		def      UpperCAmelCase__    (      self						)			->	Tuple:




									"""simple docstring"""





									_UpperCAmelCase					=      nn.Conv(
									    self.out_channels      ,   kernel_size=(3, 3)      ,   strides=(1, 1)      ,   padding=((1, 1), (1, 1))      ,   dtype=self.dtype      ,   )






		def __call__(      self      ,   _SCREAMING_SNAKE_CASE						)			->	Optional[int]:




									"""simple docstring"""





									_UpperCAmelCase		,	_UpperCAmelCase		,	_UpperCAmelCase		,	_UpperCAmelCase					=      hidden_states.shape
									_UpperCAmelCase					=      jax.image.resize(
									    _SCREAMING_SNAKE_CASE      ,   shape=(batch, height * 2, width * 2, channels)      ,   method='nearest'      ,   )
									_UpperCAmelCase					=      self.conv(_SCREAMING_SNAKE_CASE						)
									return hidden_states



class 		__a				(    nn.Module      ):
		_a     :  int
		_a     :  jnp.dtype           =     jnp.floataa
		def      UpperCAmelCase__    (      self						)			->	Tuple:




									"""simple docstring"""





									_UpperCAmelCase					=      nn.Conv(
									    self.out_channels      ,   kernel_size=(3, 3)      ,   strides=(2, 2)      ,   padding=((1, 1), (1, 1))      ,   dtype=self.dtype      ,   )






		def __call__(      self      ,   _SCREAMING_SNAKE_CASE						)			->	Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      self.conv(_SCREAMING_SNAKE_CASE						)
									return hidden_states



class 		__a				(    nn.Module      ):
		_a     :  int
		_a     :  int           =     None
		_a     :  float           =     0.0
		_a     :  bool           =     None
		_a     :  jnp.dtype           =     jnp.floataa
		def      UpperCAmelCase__    (      self						)			->	str:




									"""simple docstring"""





									_UpperCAmelCase					=      self.in_channels if self.out_channels is None else self.out_channels

									_UpperCAmelCase					=      nn.GroupNorm(num_groups=32      ,   epsilon=1e-5						)
									_UpperCAmelCase					=      nn.Conv(
									    _SCREAMING_SNAKE_CASE      ,   kernel_size=(3, 3)      ,   strides=(1, 1)      ,   padding=((1, 1), (1, 1))      ,   dtype=self.dtype      ,   )

									_UpperCAmelCase					=      nn.Dense(_SCREAMING_SNAKE_CASE      ,   dtype=self.dtype						)

									_UpperCAmelCase					=      nn.GroupNorm(num_groups=32      ,   epsilon=1e-5						)
									_UpperCAmelCase					=      nn.Dropout(self.dropout_prob						)
									_UpperCAmelCase					=      nn.Conv(
									    _SCREAMING_SNAKE_CASE      ,   kernel_size=(3, 3)      ,   strides=(1, 1)      ,   padding=((1, 1), (1, 1))      ,   dtype=self.dtype      ,   )

									_UpperCAmelCase					=      self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

									_UpperCAmelCase					=      None
									if use_nin_shortcut:
																_UpperCAmelCase					=      nn.Conv(
																    _SCREAMING_SNAKE_CASE      ,   kernel_size=(1, 1)      ,   strides=(1, 1)      ,   padding='VALID'      ,   dtype=self.dtype      ,   )






		def __call__(      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE=True						)			->	Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      hidden_states
									_UpperCAmelCase					=      self.norma(_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      nn.swish(_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      self.conva(_SCREAMING_SNAKE_CASE						)

									_UpperCAmelCase					=      self.time_emb_proj(nn.swish(_SCREAMING_SNAKE_CASE						)						)
									_UpperCAmelCase					=      jnp.expand_dims(jnp.expand_dims(_SCREAMING_SNAKE_CASE      ,   1						)      ,   1						)
									_UpperCAmelCase					=      hidden_states + temb

									_UpperCAmelCase					=      self.norma(_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      nn.swish(_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      self.dropout(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      self.conva(_SCREAMING_SNAKE_CASE						)

									if self.conv_shortcut is not None:
																_UpperCAmelCase					=      self.conv_shortcut(_SCREAMING_SNAKE_CASE						)

									return hidden_states + residual




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 		__a				(    UpperCAmelCase      ):
		_a     :  Union[List[np.ndarray], torch.FloatTensor]


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
