


'''simple docstring'''



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





class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    lowerCamelCase  :					int
    lowerCamelCase  :					jnp.dtype							     =			jnp.floataa







    def       lowercase_				(							self	)				->     int:
         __lowerCamelCase   :	Union[str, Any]								=		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_	)				->     Dict:
         __lowerCamelCase				,       __lowerCamelCase				,       __lowerCamelCase				,       __lowerCamelCase   :	Any								=		hidden_states.shape
         __lowerCamelCase   :	Optional[Any]								=		jax.image.resize(
             SCREAMING_SNAKE_CASE_			,     shape=(batch, height * 2, width * 2, channels)			,     method='nearest'			,     )
         __lowerCamelCase   :	Optional[int]								=		self.conv(SCREAMING_SNAKE_CASE_	)
         return hidden_states







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    lowerCamelCase  :					int
    lowerCamelCase  :					jnp.dtype							     =			jnp.floataa







    def       lowercase_				(							self	)				->     int:
         __lowerCamelCase   :	Any								=		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_	)				->     List[str]:
         # pad = ((0, 0), (0, 1), (0, 1), (0, 0))  # pad height and width dim
         # hidden_states = jnp.pad(hidden_states, pad_width=pad)
         __lowerCamelCase   :	Optional[Any]								=		self.conv(SCREAMING_SNAKE_CASE_	)
         return hidden_states







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    lowerCamelCase  :					int
    lowerCamelCase  :					int							     =			None
    lowerCamelCase  :					float							     =			0.0
    lowerCamelCase  :					bool							     =			None
    lowerCamelCase  :					jnp.dtype							     =			jnp.floataa







    def       lowercase_				(							self	)				->     Dict:
         __lowerCamelCase   :	Optional[int]								=		self.in_channels if self.out_channels is None else self.out_channels

         __lowerCamelCase   :	Optional[Any]								=		nn.GroupNorm(num_groups=32			,     epsilon=1E-5	)
         __lowerCamelCase   :	List[Any]								=		nn.Conv(
             SCREAMING_SNAKE_CASE_			,     kernel_size=(3, 3)			,     strides=(1, 1)			,     padding=((1, 1), (1, 1))			,     dtype=self.dtype			,     )

         __lowerCamelCase   :	Dict								=		nn.Dense(SCREAMING_SNAKE_CASE_			,     dtype=self.dtype	)

         __lowerCamelCase   :	Tuple								=		nn.GroupNorm(num_groups=32			,     epsilon=1E-5	)
         __lowerCamelCase   :	List[Any]								=		nn.Dropout(self.dropout_prob	)
         __lowerCamelCase   :	List[str]								=		nn.Conv(
             SCREAMING_SNAKE_CASE_			,     kernel_size=(3, 3)			,     strides=(1, 1)			,     padding=((1, 1), (1, 1))			,     dtype=self.dtype			,     )

         __lowerCamelCase   :	Optional[Any]								=		self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

         __lowerCamelCase   :	Any								=		None
         if use_nin_shortcut:
              __lowerCamelCase   :	Dict								=		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	)				->     Any:
         __lowerCamelCase   :	str								=		hidden_states
         __lowerCamelCase   :	Tuple								=		self.norma(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Optional[Any]								=		nn.swish(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Any								=		self.conva(SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	List[Any]								=		self.time_emb_proj(nn.swish(SCREAMING_SNAKE_CASE_	)	)
         __lowerCamelCase   :	int								=		jnp.expand_dims(jnp.expand_dims(SCREAMING_SNAKE_CASE_			,     1	)			,     1	)
         __lowerCamelCase   :	List[Any]								=		hidden_states + temb

         __lowerCamelCase   :	Any								=		self.norma(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	int								=		nn.swish(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Optional[Any]								=		self.dropout(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	List[str]								=		self.conva(SCREAMING_SNAKE_CASE_	)

         if self.conv_shortcut is not None:
              __lowerCamelCase   :	int								=		self.conv_shortcut(SCREAMING_SNAKE_CASE_	)

         return hidden_states + residual










'''simple docstring'''



import unittest

import torch

from diffusers import VQModel
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()





class    UpperCAmelCase_  (_UpperCAmelCase					,		_UpperCAmelCase					,		unittest.TestCase  ):





    """simple docstring"""






    lowerCamelCase  :					Union[str, Any]							     =			VQModel
    lowerCamelCase  :					Union[str, Any]							     =			'sample'







    @property
    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=(32, 32)	)				->     Any:
         __lowerCamelCase   :	Tuple								=		4
         __lowerCamelCase   :	Union[str, Any]								=		3

         __lowerCamelCase   :	List[Any]								=		floats_tensor((batch_size, num_channels) + sizes	).to(SCREAMING_SNAKE_CASE_	)

         return {"sample": image}







    @property
    def       lowercase_				(							self	)				->     Optional[int]:
         return (3, 32, 32)







    @property
    def       lowercase_				(							self	)				->     List[Any]:
         return (3, 32, 32)







    def       lowercase_				(							self	)				->     Optional[int]:
         __lowerCamelCase   :	Optional[Any]								=		{
             'block_out_channels': [32, 64],
             'in_channels': 3,
             'out_channels': 3,
             'down_block_types': ['DownEncoderBlock2D', 'DownEncoderBlock2D'],
             'up_block_types': ['UpDecoderBlock2D', 'UpDecoderBlock2D'],
             'latent_channels': 3,
         }
         __lowerCamelCase   :	Union[str, Any]								=		self.dummy_input
         return init_dict, inputs_dict







    def       lowercase_				(							self	)				->     str:
         pass







    def       lowercase_				(							self	)				->     Optional[int]:
         pass







    def       lowercase_				(							self	)				->     Any:
         __lowerCamelCase				,       __lowerCamelCase   :	Union[str, Any]								=		VQModel.from_pretrained('fusing/vqgan-dummy'			,     output_loading_info=SCREAMING_SNAKE_CASE_	)
         self.assertIsNotNone(SCREAMING_SNAKE_CASE_	)
         self.assertEqual(len(loading_info['missing_keys']	)			,     0	)

         model.to(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	List[Any]								=		model(**self.dummy_input	)

         assert image is not None, "Make sure output is not None"







    def       lowercase_				(							self	)				->     int:
         __lowerCamelCase   :	List[Any]								=		VQModel.from_pretrained('fusing/vqgan-dummy'	)
         model.to(SCREAMING_SNAKE_CASE_	).eval()

         torch.manual_seed(0	)
         if torch.cuda.is_available():
              torch.cuda.manual_seed_all(0	)

         __lowerCamelCase   :	List[str]								=		torch.randn(1			,     model.config.in_channels			,     model.config.sample_size			,     model.config.sample_size	)
         __lowerCamelCase   :	Optional[int]								=		image.to(SCREAMING_SNAKE_CASE_	)
         with torch.no_grad():
              __lowerCamelCase   :	int								=		model(SCREAMING_SNAKE_CASE_	).sample

         __lowerCamelCase   :	List[str]								=		output[0, -1, -3:, -3:].flatten().cpu()
         # fmt: off
         __lowerCamelCase   :	Union[str, Any]								=		torch.tensor([-0.0_1_5_3, -0.4_0_4_4, -0.1_8_8_0, -0.5_1_6_1, -0.2_4_1_8, -0.4_0_7_2, -0.1_6_1_2, -0.0_6_3_3, -0.0_1_4_3]	)
         # fmt: on
         self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     atol=1E-3	)	)






