





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     SCREAMING_SNAKE_CASE__ (						a__				,					a__				,					unittest.TestCase      ):
 __SCREAMING_SNAKE_CASE							       =	VQModel
 __SCREAMING_SNAKE_CASE							       =	'''sample'''



 @property
 def 				UpperCamelCase			(							self,__lowerCamelCase=(32, 32) ):
        A__											=						4
        A__											=						3

        A__											=						floats_tensor((batch_size, num_channels) + sizes ).to(_UpperCAmelCase )

        return {"sample": image}



 @property
 def 				UpperCamelCase			(							self ):
        return (3, 32, 32)



 @property
 def 				UpperCamelCase			(							self ):
        return (3, 32, 32)



 def 				UpperCamelCase			(							self ):
        A__											=						{
            'block_out_channels': [32, 64],
            'in_channels': 3,
            'out_channels': 3,
            'down_block_types': ['DownEncoderBlock2D', 'DownEncoderBlock2D'],
            'up_block_types': ['UpDecoderBlock2D', 'UpDecoderBlock2D'],
            'latent_channels': 3,
        }
        A__											=						self.dummy_input
        return init_dict, inputs_dict



 def 				UpperCamelCase			(							self ):
        pass



 def 				UpperCamelCase			(							self ):
        pass



 def 				UpperCamelCase			(							self ):
        A__											=						VQModel.from_pretrained('''fusing/vqgan-dummy''',output_loading_info=_UpperCAmelCase )
        self.assertIsNotNone(_UpperCAmelCase )
        self.assertEqual(len(loading_info['''missing_keys'''] ),0 )

        model.to(_UpperCAmelCase )
        A__											=						model(**self.dummy_input )

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




 def 				UpperCamelCase			(							self ):
        A__											=						VQModel.from_pretrained('''fusing/vqgan-dummy''' )
        model.to(_UpperCAmelCase ).eval()

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

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

        A__											=						output[0, -1, -3:, -3:].flatten().cpu()
        # fmt: off
        A__											=						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(_UpperCAmelCase,_UpperCAmelCase,atol=1E-3 ) )












'''simple docstring'''




lowercase__    :      Dict						=    {
    '''Pillow''': '''Pillow''',
    '''accelerate''': '''accelerate>=0.11.0''',
    '''compel''': '''compel==0.1.8''',
    '''black''': '''black~=23.1''',
    '''datasets''': '''datasets''',
    '''filelock''': '''filelock''',
    '''flax''': '''flax>=0.4.1''',
    '''hf-doc-builder''': '''hf-doc-builder>=0.3.0''',
    '''huggingface-hub''': '''huggingface-hub>=0.13.2''',
    '''requests-mock''': '''requests-mock==1.10.0''',
    '''importlib_metadata''': '''importlib_metadata''',
    '''invisible-watermark''': '''invisible-watermark''',
    '''isort''': '''isort>=5.5.4''',
    '''jax''': '''jax>=0.2.8,!=0.3.2''',
    '''jaxlib''': '''jaxlib>=0.1.65''',
    '''Jinja2''': '''Jinja2''',
    '''k-diffusion''': '''k-diffusion>=0.0.12''',
    '''torchsde''': '''torchsde''',
    '''note_seq''': '''note_seq''',
    '''librosa''': '''librosa''',
    '''numpy''': '''numpy''',
    '''omegaconf''': '''omegaconf''',
    '''parameterized''': '''parameterized''',
    '''protobuf''': '''protobuf>=3.20.3,<4''',
    '''pytest''': '''pytest''',
    '''pytest-timeout''': '''pytest-timeout''',
    '''pytest-xdist''': '''pytest-xdist''',
    '''ruff''': '''ruff>=0.0.241''',
    '''safetensors''': '''safetensors''',
    '''sentencepiece''': '''sentencepiece>=0.1.91,!=0.1.92''',
    '''scipy''': '''scipy''',
    '''onnx''': '''onnx''',
    '''regex''': '''regex!=2019.12.17''',
    '''requests''': '''requests''',
    '''tensorboard''': '''tensorboard''',
    '''torch''': '''torch>=1.4''',
    '''torchvision''': '''torchvision''',
    '''transformers''': '''transformers>=4.25.1''',
    '''urllib3''': '''urllib3<=2.0.0''',
}