




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   A(						__lowerCAmelCase							,				__lowerCAmelCase							,				unittest.TestCase       ):



 '''simple docstring'''


 UpperCamelCase						       =						VQModel
 UpperCamelCase						       =						"sample"

 @property
 def     a__			(    self :  Optional[int]						,							A_ :  str=(32, 32)  )       ->     Optional[Any]:

       """simple docstring"""






       lowerCamelCase_							=   4
       lowerCamelCase_							=   3

       lowerCamelCase_							=   floats_tensor((batch_size, num_channels) + sizes  ).to(lowerCamelCase__  )

       return {"sample": image}

 @property
 def     a__			(    self :  Dict  )       ->     List[str]:

       """simple docstring"""






       return (3, 32, 32)

 @property
 def     a__			(    self :  Tuple  )       ->     Union[str, Any]:

       """simple docstring"""






       return (3, 32, 32)

 def     a__			(    self :  Optional[Any]  )       ->     int:

       """simple docstring"""






       lowerCamelCase_							=   {
           '''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_							=   self.dummy_input
       return init_dict, inputs_dict

 def     a__			(    self :  int  )       ->     Union[str, Any]:

       """simple docstring"""






       pass

 def     a__			(    self :  Any  )       ->     int:

       """simple docstring"""






       pass

 def     a__			(    self :  Dict  )       ->     Tuple:

       """simple docstring"""






       lowerCamelCase_							=   VQModel.from_pretrained('fusing/vqgan-dummy'						,							output_loading_info=lowerCamelCase__  )
       self.assertIsNotNone(lowerCamelCase__  )
       self.assertEqual(len(loading_info['missing_keys']  )						,							0  )

       model.to(lowerCamelCase__  )
       lowerCamelCase_							=   model(**self.dummy_input  )

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





 def     a__			(    self :  List[str]  )       ->     Dict:

       """simple docstring"""






       lowerCamelCase_							=   VQModel.from_pretrained('fusing/vqgan-dummy'  )
       model.to(lowerCamelCase__  ).eval()

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

       lowerCamelCase_							=   torch.randn(1						,							model.config.in_channels						,							model.config.sample_size						,							model.config.sample_size  )
       lowerCamelCase_							=   image.to(lowerCamelCase__  )
       with torch.no_grad():
             lowerCamelCase_							=   model(lowerCamelCase__  ).sample

       lowerCamelCase_							=   output[0, -1, -3:, -3:].flatten().cpu()
       # fmt: off
       lowerCamelCase_							=   torch.tensor([-0.0153, -0.4044, -0.1880, -0.5161, -0.2418, -0.4072, -0.1612, -0.0633, -0.0143]  )
       # fmt: on
       self.assertTrue(torch.allclose(lowerCamelCase__						,							lowerCamelCase__						,							atol=1E-3  )  )









from ..utils import is_flax_available, is_torch_available


if is_torch_available():
 from .autoencoder_kl import AutoencoderKL
 from .controlnet import ControlNetModel
 from .dual_transformer_ad import DualTransformeraDModel
 from .modeling_utils import ModelMixin
 from .prior_transformer import PriorTransformer
 from .ta_film_transformer import TaFilmDecoder
 from .transformer_ad import TransformeraDModel
 from .unet_ad import UNetaDModel
 from .unet_ad import UNetaDModel
 from .unet_ad_condition import UNetaDConditionModel
 from .unet_ad_condition import UNetaDConditionModel
 from .vq_model import VQModel

if is_flax_available():
 from .controlnet_flax import FlaxControlNetModel
 from .unet_ad_condition_flax import FlaxUNetaDConditionModel
 from .vae_flax import FlaxAutoencoderKL



