






import gc
import unittest

import torch
from parameterized import parameterized

from diffusers import AutoencoderKL
from diffusers.utils import floats_tensor, load_hf_numpy, require_torch_gpu, slow, torch_all_close, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()




class 			snake_case__   (							lowercase_	,    lowercase_	,    unittest.TestCase):

       '''simple docstring'''
       lowerCamelCase							:				Optional[int] 				= AutoencoderKL
       lowerCamelCase							:				Union[str, Any] 				= "sample"
       lowerCamelCase							:				Optional[Any] 				= 1E-2



       @property
       def 					__lowercase						(							self    ) ->      Optional[Any]:






            '''simple docstring'''





            __snake_case			:List[str]          =						4
            __snake_case			:Any          =						3
            __snake_case			:List[Any]          =						(32, 32)

            __snake_case			:int          =						floats_tensor((batch_size, num_channels) + sizes    ).to(a__    )

            return {"sample": image}



       @property
       def 					__lowercase						(							self    ) ->      Dict:






            '''simple docstring'''





            return (3, 32, 32)



       @property
       def 					__lowercase						(							self    ) ->      Optional[Any]:






            '''simple docstring'''





            return (3, 32, 32)



       def 					__lowercase						(							self    ) ->      Optional[Any]:






            '''simple docstring'''





            __snake_case			:Dict          =						{
                """block_out_channels""": [32, 64],
                """in_channels""": 3,
                """out_channels""": 3,
                """down_block_types""": ["""DownEncoderBlock2D""", """DownEncoderBlock2D"""],
                """up_block_types""": ["""UpDecoderBlock2D""", """UpDecoderBlock2D"""],
                """latent_channels""": 4,
            }
            __snake_case			:Dict          =						self.dummy_input
            return init_dict, inputs_dict



       def 					__lowercase						(							self    ) ->      List[Any]:






            '''simple docstring'''





            pass



       def 					__lowercase						(							self    ) ->      List[Any]:






            '''simple docstring'''





            pass



       @unittest.skipIf(torch_device == """mps"""			,						"""Gradient checkpointing skipped on MPS"""    )
       def 					__lowercase						(							self    ) ->      Tuple:






            '''simple docstring'''





            __snake_case       , __snake_case			:List[Any]          =						self.prepare_init_args_and_inputs_for_common()
            __snake_case			:List[Any]          =						self.model_class(**a__    )
            model.to(a__    )

            assert not model.is_gradient_checkpointing and model.training

            __snake_case			:Union[str, Any]          =						model(**a__    ).sample
            # run the backwards pass on the model. For backwards pass, for simplicity purpose,
            # we won't calculate the loss and rather backprop on out.sum()
            model.zero_grad()

            __snake_case			:Union[str, Any]          =						torch.randn_like(a__    )
            __snake_case			:int          =						(out - labels).mean()
            loss.backward()

            # re-instantiate the model now enabling gradient checkpointing
            __snake_case			:Optional[Any]          =						self.model_class(**a__    )
            # clone model
            model_a.load_state_dict(model.state_dict()    )
            model_a.to(a__    )
            model_a.enable_gradient_checkpointing()

            assert model_a.is_gradient_checkpointing and model_a.training

            __snake_case			:Union[str, Any]          =						model_a(**a__    ).sample
            # run the backwards pass on the model. For backwards pass, for simplicity purpose,
            # we won't calculate the loss and rather backprop on out.sum()
            model_a.zero_grad()
            __snake_case			:Optional[Any]          =						(out_a - labels).mean()
            loss_a.backward()

            # compare the output and parameters gradients
            self.assertTrue((loss - loss_a).abs() < 1e-5    )
            __snake_case			:Dict          =						dict(model.named_parameters()    )
            __snake_case			:Dict          =						dict(model_a.named_parameters()    )
            for name, param in named_params.items():
                 self.assertTrue(torch_all_close(param.grad.data			,						named_params_a[name].grad.data			,						atol=5e-5    )    )



       def 					__lowercase						(							self    ) ->      Optional[int]:






            '''simple docstring'''





            __snake_case       , __snake_case			:Optional[Any]          =						AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""			,						output_loading_info=a__    )
            self.assertIsNotNone(a__    )
            self.assertEqual(len(loading_info["""missing_keys"""]    )			,						0    )

            model.to(a__    )
            __snake_case			:Dict          =						model(**self.dummy_input    )

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



       def 					__lowercase						(							self    ) ->      Optional[Any]:






            '''simple docstring'''





            __snake_case			:Optional[Any]          =						AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""    )
            __snake_case			:List[Any]          =						model.to(a__    )
            model.eval()

            if torch_device == "mps":
                 __snake_case			:Any          =						torch.manual_seed(0    )
            else:
                 __snake_case			:Dict          =						torch.Generator(device=a__    ).manual_seed(0    )

            __snake_case			:Dict          =						torch.randn(
                1			,						model.config.in_channels			,						model.config.sample_size			,						model.config.sample_size			,						generator=torch.manual_seed(0    )			,						)
            __snake_case			:Optional[int]          =						image.to(a__    )
            with torch.no_grad():
                 __snake_case			:Tuple          =						model(a__			,						sample_posterior=a__			,						generator=a__    ).sample

            __snake_case			:int          =						output[0, -1, -3:, -3:].flatten().cpu()

            # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
            # the expected output slices are not the same for CPU and GPU.
            if torch_device == "mps":
                 __snake_case			:Optional[Any]          =						torch.tensor(
                     [
                         -4.00_78e-01,
                         -3.83_23e-04,
                         -1.26_81e-01,
                         -1.14_62e-01,
                         2.00_95e-01,
                         1.08_93e-01,
                         -8.82_47e-02,
                         -3.03_61e-01,
                         -9.86_44e-03,
                     ]    )
            elif torch_device == "cpu":
                 __snake_case			:str          =						torch.tensor(
                     [-0.13_52, 0.08_78, 0.04_19, -0.08_18, -0.10_69, 0.06_88, -0.14_58, -0.44_46, -0.00_26]    )
            else:
                 __snake_case			:Optional[int]          =						torch.tensor(
                     [-0.24_21, 0.46_42, 0.25_07, -0.04_38, 0.06_82, 0.31_60, -0.20_18, -0.07_27, 0.24_85]    )

            self.assertTrue(torch_all_close(a__			,						a__			,						rtol=1e-2    )    )




@slow
class 			snake_case__   (							unittest.TestCase):

       '''simple docstring'''



       def 					__lowercase						(							self			,						a__			,						a__    ) ->      Any:






            '''simple docstring'''





            return F'''gaussian_noise_s={seed}_shape={'_'.join([str(a__    ) for s in shape]    )}.npy'''



       def 					__lowercase						(							self    ) ->      Dict:






            '''simple docstring'''





            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()



       def 					__lowercase						(							self			,						a__=0			,						a__=(4, 3, 5_12, 5_12)			,						a__=False    ) ->      Optional[int]:






            '''simple docstring'''





            __snake_case			:Optional[Any]          =						torch.floataa if fpaa else torch.floataa
            __snake_case			:Optional[int]          =						torch.from_numpy(load_hf_numpy(self.get_file_format(a__			,						a__    )    )    ).to(a__    ).to(a__    )
            return image



       def 					__lowercase						(							self			,						a__="CompVis/stable-diffusion-v1-4"			,						a__=False    ) ->      int:






            '''simple docstring'''





            __snake_case			:str          =						"""fp16""" if fpaa else None
            __snake_case			:Optional[int]          =						torch.floataa if fpaa else torch.floataa

            __snake_case			:Optional[Any]          =						AutoencoderKL.from_pretrained(
                a__			,						subfolder="""vae"""			,						torch_dtype=a__			,						revision=a__			,						)
            model.to(a__    ).eval()

            return model



       def 					__lowercase						(							self			,						a__=0    ) ->      List[str]:






            '''simple docstring'''





            if torch_device == "mps":
                 return torch.manual_seed(a__    )
            return torch.Generator(device=a__    ).manual_seed(a__    )



       @parameterized.expand(
           [
               # fmt: off
               [33, [-0.16_03, 0.98_78, -0.04_95, -0.07_90, -0.27_09, 0.83_75, -0.20_60, -0.08_24], [-0.23_95, 0.00_98, 0.01_02, -0.07_09, -0.28_40, -0.02_74, -0.07_18, -0.18_24]],
               [47, [-0.23_76, 0.11_68, 0.13_32, -0.48_40, -0.25_08, -0.07_91, -0.04_93, -0.40_89], [0.03_50, 0.08_47, 0.04_67, 0.03_44, -0.08_42, -0.05_47, -0.06_33, -0.11_31]],
               # fmt: on
           ]    )
       def 					__lowercase						(							self			,						a__			,						a__			,						a__    ) ->      Optional[Any]:






            '''simple docstring'''





            __snake_case			:Any          =						self.get_sd_vae_model()
            __snake_case			:Union[str, Any]          =						self.get_sd_image(a__    )
            __snake_case			:List[str]          =						self.get_generator(a__    )

            with torch.no_grad():
                 __snake_case			:str          =						model(a__			,						generator=a__			,						sample_posterior=a__    ).sample

            assert sample.shape == image.shape

            __snake_case			:List[str]          =						sample[-1, -2:, -2:, :2].flatten().float().cpu()
            __snake_case			:Tuple          =						torch.tensor(expected_slice_mps if torch_device == """mps""" else expected_slice    )

            assert torch_all_close(a__			,						a__			,						atol=3e-3    )



       @parameterized.expand(
           [
               # fmt: off
               [33, [-0.05_13, 0.02_89, 1.37_99, 0.21_66, -0.25_73, -0.08_71, 0.51_03, -0.09_99]],
               [47, [-0.41_28, -0.13_20, -0.37_04, 0.19_65, -0.41_16, -0.23_32, -0.33_40, 0.22_47]],
               # fmt: on
           ]    )
       @require_torch_gpu
       def 					__lowercase						(							self			,						a__			,						a__    ) ->      List[str]:






            '''simple docstring'''





            __snake_case			:Any          =						self.get_sd_vae_model(fpaa=a__    )
            __snake_case			:Optional[int]          =						self.get_sd_image(a__			,						fpaa=a__    )
            __snake_case			:List[Any]          =						self.get_generator(a__    )

            with torch.no_grad():
                 __snake_case			:List[Any]          =						model(a__			,						generator=a__			,						sample_posterior=a__    ).sample

            assert sample.shape == image.shape

            __snake_case			:str          =						sample[-1, -2:, :2, -2:].flatten().float().cpu()
            __snake_case			:Any          =						torch.tensor(a__    )

            assert torch_all_close(a__			,						a__			,						atol=1e-2    )



       @parameterized.expand(
           [
               # fmt: off
               [33, [-0.16_09, 0.98_66, -0.04_87, -0.07_77, -0.27_16, 0.83_68, -0.20_55, -0.08_14], [-0.23_95, 0.00_98, 0.01_02, -0.07_09, -0.28_40, -0.02_74, -0.07_18, -0.18_24]],
               [47, [-0.23_77, 0.11_47, 0.13_33, -0.48_41, -0.25_06, -0.08_05, -0.04_91, -0.40_85], [0.03_50, 0.08_47, 0.04_67, 0.03_44, -0.08_42, -0.05_47, -0.06_33, -0.11_31]],
               # fmt: on
           ]    )
       def 					__lowercase						(							self			,						a__			,						a__			,						a__    ) ->      Tuple:






            '''simple docstring'''





            __snake_case			:Tuple          =						self.get_sd_vae_model()
            __snake_case			:Dict          =						self.get_sd_image(a__    )

            with torch.no_grad():
                 __snake_case			:Tuple          =						model(a__    ).sample

            assert sample.shape == image.shape

            __snake_case			:Dict          =						sample[-1, -2:, -2:, :2].flatten().float().cpu()
            __snake_case			:Union[str, Any]          =						torch.tensor(expected_slice_mps if torch_device == """mps""" else expected_slice    )

            assert torch_all_close(a__			,						a__			,						atol=3e-3    )



       @parameterized.expand(
           [
               # fmt: off
               [13, [-0.20_51, -0.18_03, -0.23_11, -0.21_14, -0.32_92, -0.35_74, -0.29_53, -0.33_23]],
               [37, [-0.26_32, -0.26_25, -0.21_99, -0.27_41, -0.45_39, -0.49_90, -0.37_20, -0.49_25]],
               # fmt: on
           ]    )
       @require_torch_gpu
       def 					__lowercase						(							self			,						a__			,						a__    ) ->      List[str]:






            '''simple docstring'''





            __snake_case			:Optional[Any]          =						self.get_sd_vae_model()
            __snake_case			:int          =						self.get_sd_image(a__			,						shape=(3, 4, 64, 64)    )

            with torch.no_grad():
                 __snake_case			:str          =						model.decode(a__    ).sample

            assert list(sample.shape    ) == [3, 3, 5_12, 5_12]

            __snake_case			:Optional[int]          =						sample[-1, -2:, :2, -2:].flatten().cpu()
            __snake_case			:str          =						torch.tensor(a__    )

            assert torch_all_close(a__			,						a__			,						atol=1e-3    )



       @parameterized.expand(
           [
               # fmt: off
               [27, [-0.03_69, 0.02_07, -0.07_76, -0.06_82, -0.17_47, -0.19_30, -0.14_65, -0.20_39]],
               [16, [-0.16_28, -0.21_34, -0.27_47, -0.26_42, -0.37_74, -0.44_04, -0.36_87, -0.42_77]],
               # fmt: on
           ]    )
       @require_torch_gpu
       def 					__lowercase						(							self			,						a__			,						a__    ) ->      List[str]:






            '''simple docstring'''





            __snake_case			:Optional[Any]          =						self.get_sd_vae_model(fpaa=a__    )
            __snake_case			:Optional[Any]          =						self.get_sd_image(a__			,						shape=(3, 4, 64, 64)			,						fpaa=a__    )

            with torch.no_grad():
                 __snake_case			:Any          =						model.decode(a__    ).sample

            assert list(sample.shape    ) == [3, 3, 5_12, 5_12]

            __snake_case			:Dict          =						sample[-1, -2:, :2, -2:].flatten().float().cpu()
            __snake_case			:Any          =						torch.tensor(a__    )

            assert torch_all_close(a__			,						a__			,						atol=5e-3    )



       @parameterized.expand([(13,), (16,), (27,)]    )
       @require_torch_gpu
       @unittest.skipIf(not is_xformers_available()			,						reason="""xformers is not required when using PyTorch 2.0."""    )
       def 					__lowercase						(							self			,						a__    ) ->      List[str]:






            '''simple docstring'''





            __snake_case			:Dict          =						self.get_sd_vae_model(fpaa=a__    )
            __snake_case			:Dict          =						self.get_sd_image(a__			,						shape=(3, 4, 64, 64)			,						fpaa=a__    )

            with torch.no_grad():
                 __snake_case			:Optional[int]          =						model.decode(a__    ).sample

            model.enable_xformers_memory_efficient_attention()
            with torch.no_grad():
                 __snake_case			:List[Any]          =						model.decode(a__    ).sample

            assert list(sample.shape    ) == [3, 3, 5_12, 5_12]

            assert torch_all_close(a__			,						a__			,						atol=1e-1    )



       @parameterized.expand([(13,), (16,), (37,)]    )
       @require_torch_gpu
       @unittest.skipIf(not is_xformers_available()			,						reason="""xformers is not required when using PyTorch 2.0."""    )
       def 					__lowercase						(							self			,						a__    ) ->      Optional[int]:






            '''simple docstring'''





            __snake_case			:List[Any]          =						self.get_sd_vae_model()
            __snake_case			:List[str]          =						self.get_sd_image(a__			,						shape=(3, 4, 64, 64)    )

            with torch.no_grad():
                 __snake_case			:Any          =						model.decode(a__    ).sample

            model.enable_xformers_memory_efficient_attention()
            with torch.no_grad():
                 __snake_case			:Any          =						model.decode(a__    ).sample

            assert list(sample.shape    ) == [3, 3, 5_12, 5_12]

            assert torch_all_close(a__			,						a__			,						atol=1e-2    )



       @parameterized.expand(
           [
               # fmt: off
               [33, [-0.30_01, 0.09_18, -2.69_84, -3.97_20, -3.20_99, -5.03_53, 1.73_38, -0.20_65, 3.42_67]],
               [47, [-1.50_30, -4.38_71, -6.03_55, -9.11_57, -1.66_61, -2.78_53, 2.16_07, -5.08_23, 2.56_33]],
               # fmt: on
           ]    )
       def 					__lowercase						(							self			,						a__			,						a__    ) ->      Tuple:






            '''simple docstring'''





            __snake_case			:str          =						self.get_sd_vae_model()
            __snake_case			:Dict          =						self.get_sd_image(a__    )
            __snake_case			:Optional[int]          =						self.get_generator(a__    )

            with torch.no_grad():
                 __snake_case			:List[str]          =						model.encode(a__    ).latent_dist
                 __snake_case			:List[Any]          =						dist.sample(generator=a__    )

            assert list(sample.shape    ) == [image.shape[0], 4] + [i // 8 for i in image.shape[2:]]

            __snake_case			:List[Any]          =						sample[0, -1, -3:, -3:].flatten().cpu()
            __snake_case			:Optional[int]          =						torch.tensor(a__    )

            __snake_case			:Any          =						3e-3 if torch_device != """mps""" else 1e-2
            assert torch_all_close(a__			,						a__			,						atol=a__    )










# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

import torch

from ..models.auto import AutoModelForVisualQuestionAnswering, AutoProcessor
from ..utils import requires_backends
from .base import PipelineTool


if TYPE_CHECKING:
     from PIL import Image




class 			snake_case__   (							lowercase_):

       '''simple docstring'''
       lowerCamelCase							:				Tuple 				= "dandelin/vilt-b32-finetuned-vqa"
       lowerCamelCase							:				List[str] 				= (
           "This is a tool that answers a question about an image. It takes an input named `image` which should be the "
           "image containing the information, as well as a `question` which should be the question in English. It "
           "returns a text that is the answer to the question."
       )
       lowerCamelCase							:				Optional[Any] 				= "image_qa"
       lowerCamelCase							:				str 				= AutoProcessor
       lowerCamelCase							:				Union[str, Any] 				= AutoModelForVisualQuestionAnswering

       lowerCamelCase							:				Any 				= ["image", "text"]
       lowerCamelCase							:				Dict 				= ["text"]



       def __init__(							self			,						*a__			,						**a__    ) ->      Any:






            '''simple docstring'''





            requires_backends(self			,						["""vision"""]    )
            super().__init__(*a__			,						**a__    )



       def 					__lowercase						(							self			,						a__			,						a__    ) ->      int:






            '''simple docstring'''





            return self.pre_processor(a__			,						a__			,						return_tensors="""pt"""    )



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






            '''simple docstring'''





            with torch.no_grad():
                 return self.model(**a__    ).logits



       def 					__lowercase						(							self			,						a__    ) ->      Tuple:






            '''simple docstring'''





            __snake_case			:str          =						outputs.argmax(-1    ).item()
            return self.model.config.idalabel[idx]


