



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     _UpperCAmelCase			(       A__			,      A__			,      unittest.TestCase     ):
     UpperCamelCase__		    =				AutoencoderKL
     UpperCamelCase__		    =				'''sample'''
     UpperCamelCase__		    =				1e-2




     @property
     def 		snake_case_		(				self):
         A__		 =					4
         A__		 =					3
         A__		 =					(3_2, 3_2)

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

         return {"sample": image}




     @property
     def 		snake_case_		(				self):
         return (3, 3_2, 3_2)




     @property
     def 		snake_case_		(				self):
         return (3, 3_2, 3_2)




     def 		snake_case_		(				self):
         A__		 =					{
             '''block_out_channels''': [3_2, 6_4],
             '''in_channels''': 3,
             '''out_channels''': 3,
             '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
             '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
             '''latent_channels''': 4,
         }
         A__		 =					self.dummy_input
         return init_dict, inputs_dict




     def 		snake_case_		(				self):
         pass




     def 		snake_case_		(				self):
         pass




     @unittest.skipIf(torch_device == '''mps''' ,						'''Gradient checkpointing skipped on MPS''')
     def 		snake_case_		(				self):
         # enable deterministic behavior for gradient checkpointing
         A__ ,	A__		 =					self.prepare_init_args_and_inputs_for_common()
         A__		 =					self.model_class(**a__)
         model.to(a__)

         assert not model.is_gradient_checkpointing and model.training

         A__		 =					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()

         A__		 =					torch.randn_like(a__)
         A__		 =					(out - labels).mean()
         loss.backward()

         # re-instantiate the model now enabling gradient checkpointing
         A__		 =					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

         A__		 =					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()
         A__		 =					(out_a - labels).mean()
         loss_a.backward()

         # compare the output and parameters gradients
         self.assertTrue((loss - loss_a).abs() < 1e-5)
         A__		 =					dict(model.named_parameters())
         A__		 =					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 		snake_case_		(				self):
         A__ ,	A__		 =					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__)
         A__		 =					model(**self.dummy_input)

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




     def 		snake_case_		(				self):
         A__		 =					AutoencoderKL.from_pretrained('''fusing/autoencoder-kl-dummy''')
         A__		 =					model.to(a__)
         model.eval()

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

         A__		 =					torch.randn(
             1 ,						model.config.in_channels ,						model.config.sample_size ,						model.config.sample_size ,						generator=torch.manual_seed(0) ,						)
         A__		 =					image.to(a__)
         with torch.no_grad():
             A__		 =					model(a__ ,						sample_posterior=a__ ,						generator=a__).sample

         A__		 =					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":
             A__		 =					torch.tensor(
                 [
                     -4.0_078e-01,
                     -3.8_323e-04,
                     -1.2_681e-01,
                     -1.1_462e-01,
                     2.0_095e-01,
                     1.0_893e-01,
                     -8.8_247e-02,
                     -3.0_361e-01,
                     -9.8_644e-03,
                 ])
         elif torch_device == "cpu":
             A__		 =					torch.tensor(
                 [-0.1_3_5_2, 0.0_8_7_8, 0.0_4_1_9, -0.0_8_1_8, -0.1_0_6_9, 0.0_6_8_8, -0.1_4_5_8, -0.4_4_4_6, -0.0_0_2_6])
         else:
             A__		 =					torch.tensor(
                 [-0.2_4_2_1, 0.4_6_4_2, 0.2_5_0_7, -0.0_4_3_8, 0.0_6_8_2, 0.3_1_6_0, -0.2_0_1_8, -0.0_7_2_7, 0.2_4_8_5])

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







@slow
class     _UpperCAmelCase			(       unittest.TestCase     ):




     def 		snake_case_		(				self ,						a__ ,						a__):
         return F"gaussian_noise_s={seed}_shape={'_'.join([str(a__) for s in shape])}.npy"




     def 		snake_case_		(				self):
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()




     def 		snake_case_		(				self ,						a__=0 ,						a__=(4, 3, 5_1_2, 5_1_2) ,						a__=False):
         A__		 =					torch.floataa if fpaa else torch.floataa
         A__		 =					torch.from_numpy(load_hf_numpy(self.get_file_format(a__ ,						a__))).to(a__).to(a__)
         return image




     def 		snake_case_		(				self ,						a__="CompVis/stable-diffusion-v1-4" ,						a__=False):
         A__		 =					'''fp16''' if fpaa else None
         A__		 =					torch.floataa if fpaa else torch.floataa

         A__		 =					AutoencoderKL.from_pretrained(
             a__ ,						subfolder='''vae''' ,						torch_dtype=a__ ,						revision=a__ ,						)
         model.to(a__).eval()

         return model




     def 		snake_case_		(				self ,						a__=0):
         if torch_device == "mps":
             return torch.manual_seed(a__)
         return torch.Generator(device=a__).manual_seed(a__)




     @parameterized.expand(
         [
             # fmt: off
             [3_3, [-0.1_6_0_3, 0.9_8_7_8, -0.0_4_9_5, -0.0_7_9_0, -0.2_7_0_9, 0.8_3_7_5, -0.2_0_6_0, -0.0_8_2_4], [-0.2_3_9_5, 0.0_0_9_8, 0.0_1_0_2, -0.0_7_0_9, -0.2_8_4_0, -0.0_2_7_4, -0.0_7_1_8, -0.1_8_2_4]],
             [4_7, [-0.2_3_7_6, 0.1_1_6_8, 0.1_3_3_2, -0.4_8_4_0, -0.2_5_0_8, -0.0_7_9_1, -0.0_4_9_3, -0.4_0_8_9], [0.0_3_5_0, 0.0_8_4_7, 0.0_4_6_7, 0.0_3_4_4, -0.0_8_4_2, -0.0_5_4_7, -0.0_6_3_3, -0.1_1_3_1]],
             # fmt: on
         ])
     def 		snake_case_		(				self ,						a__ ,						a__ ,						a__):
         A__		 =					self.get_sd_vae_model()
         A__		 =					self.get_sd_image(a__)
         A__		 =					self.get_generator(a__)

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

         assert sample.shape == image.shape

         A__		 =					sample[-1, -2:, -2:, :2].flatten().float().cpu()
         A__		 =					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
             [3_3, [-0.0_5_1_3, 0.0_2_8_9, 1.3_7_9_9, 0.2_1_6_6, -0.2_5_7_3, -0.0_8_7_1, 0.5_1_0_3, -0.0_9_9_9]],
             [4_7, [-0.4_1_2_8, -0.1_3_2_0, -0.3_7_0_4, 0.1_9_6_5, -0.4_1_1_6, -0.2_3_3_2, -0.3_3_4_0, 0.2_2_4_7]],
             # fmt: on
         ])
     @require_torch_gpu
     def 		snake_case_		(				self ,						a__ ,						a__):
         A__		 =					self.get_sd_vae_model(fpaa=a__)
         A__		 =					self.get_sd_image(a__ ,						fpaa=a__)
         A__		 =					self.get_generator(a__)

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

         assert sample.shape == image.shape

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

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




     @parameterized.expand(
         [
             # fmt: off
             [3_3, [-0.1_6_0_9, 0.9_8_6_6, -0.0_4_8_7, -0.0_7_7_7, -0.2_7_1_6, 0.8_3_6_8, -0.2_0_5_5, -0.0_8_1_4], [-0.2_3_9_5, 0.0_0_9_8, 0.0_1_0_2, -0.0_7_0_9, -0.2_8_4_0, -0.0_2_7_4, -0.0_7_1_8, -0.1_8_2_4]],
             [4_7, [-0.2_3_7_7, 0.1_1_4_7, 0.1_3_3_3, -0.4_8_4_1, -0.2_5_0_6, -0.0_8_0_5, -0.0_4_9_1, -0.4_0_8_5], [0.0_3_5_0, 0.0_8_4_7, 0.0_4_6_7, 0.0_3_4_4, -0.0_8_4_2, -0.0_5_4_7, -0.0_6_3_3, -0.1_1_3_1]],
             # fmt: on
         ])
     def 		snake_case_		(				self ,						a__ ,						a__ ,						a__):
         A__		 =					self.get_sd_vae_model()
         A__		 =					self.get_sd_image(a__)

         with torch.no_grad():
             A__		 =					model(a__).sample

         assert sample.shape == image.shape

         A__		 =					sample[-1, -2:, -2:, :2].flatten().float().cpu()
         A__		 =					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
             [1_3, [-0.2_0_5_1, -0.1_8_0_3, -0.2_3_1_1, -0.2_1_1_4, -0.3_2_9_2, -0.3_5_7_4, -0.2_9_5_3, -0.3_3_2_3]],
             [3_7, [-0.2_6_3_2, -0.2_6_2_5, -0.2_1_9_9, -0.2_7_4_1, -0.4_5_3_9, -0.4_9_9_0, -0.3_7_2_0, -0.4_9_2_5]],
             # fmt: on
         ])
     @require_torch_gpu
     def 		snake_case_		(				self ,						a__ ,						a__):
         A__		 =					self.get_sd_vae_model()
         A__		 =					self.get_sd_image(a__ ,						shape=(3, 4, 6_4, 6_4))

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

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

         A__		 =					sample[-1, -2:, :2, -2:].flatten().cpu()
         A__		 =					torch.tensor(a__)

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




     @parameterized.expand(
         [
             # fmt: off
             [2_7, [-0.0_3_6_9, 0.0_2_0_7, -0.0_7_7_6, -0.0_6_8_2, -0.1_7_4_7, -0.1_9_3_0, -0.1_4_6_5, -0.2_0_3_9]],
             [1_6, [-0.1_6_2_8, -0.2_1_3_4, -0.2_7_4_7, -0.2_6_4_2, -0.3_7_7_4, -0.4_4_0_4, -0.3_6_8_7, -0.4_2_7_7]],
             # fmt: on
         ])
     @require_torch_gpu
     def 		snake_case_		(				self ,						a__ ,						a__):
         A__		 =					self.get_sd_vae_model(fpaa=a__)
         A__		 =					self.get_sd_image(a__ ,						shape=(3, 4, 6_4, 6_4) ,						fpaa=a__)

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

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

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

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




     @parameterized.expand([(1_3,), (1_6,), (2_7,)])
     @require_torch_gpu
     @unittest.skipIf(not is_xformers_available() ,						reason='''xformers is not required when using PyTorch 2.0.''')
     def 		snake_case_		(				self ,						a__):
         A__		 =					self.get_sd_vae_model(fpaa=a__)
         A__		 =					self.get_sd_image(a__ ,						shape=(3, 4, 6_4, 6_4) ,						fpaa=a__)

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

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

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

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




     @parameterized.expand([(1_3,), (1_6,), (3_7,)])
     @require_torch_gpu
     @unittest.skipIf(not is_xformers_available() ,						reason='''xformers is not required when using PyTorch 2.0.''')
     def 		snake_case_		(				self ,						a__):
         A__		 =					self.get_sd_vae_model()
         A__		 =					self.get_sd_image(a__ ,						shape=(3, 4, 6_4, 6_4))

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

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

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

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




     @parameterized.expand(
         [
             # fmt: off
             [3_3, [-0.3_0_0_1, 0.0_9_1_8, -2.6_9_8_4, -3.9_7_2_0, -3.2_0_9_9, -5.0_3_5_3, 1.7_3_3_8, -0.2_0_6_5, 3.4_2_6_7]],
             [4_7, [-1.5_0_3_0, -4.3_8_7_1, -6.0_3_5_5, -9.1_1_5_7, -1.6_6_6_1, -2.7_8_5_3, 2.1_6_0_7, -5.0_8_2_3, 2.5_6_3_3]],
             # fmt: on
         ])
     def 		snake_case_		(				self ,						a__ ,						a__):
         A__		 =					self.get_sd_vae_model()
         A__		 =					self.get_sd_image(a__)
         A__		 =					self.get_generator(a__)

         with torch.no_grad():
             A__		 =					model.encode(a__).latent_dist
             A__		 =					dist.sample(generator=a__)

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

         A__		 =					sample[0, -1, -3:, -3:].flatten().cpu()
         A__		 =					torch.tensor(a__)

         A__		 =					3e-3 if torch_device != '''mps''' else 1e-2
         assert torch_all_close(a__ ,						a__ ,						atol=a__)





def 		lowerCAmelCase__    (					UpperCamelCase_				:       str ,       UpperCamelCase_				:       str	)->  float:

    def get_matched_characters(UpperCamelCase_				:       str ,       UpperCamelCase_				:       str	) -> str:
        A__		 =					[]
        A__		 =					min(len(_stra	) ,       len(_stra	)	) // 2
        for i, l in enumerate(_stra	):
            A__		 =					int(max(0 ,       i - limit	)	)
            A__		 =					int(min(i + limit + 1 ,       len(_stra	)	)	)
            if l in _stra[left:right]:
                matched.append(UpperCamelCase_	)
                A__		 =					f"{_stra[0:_stra.index(UpperCamelCase_	)]} {_stra[_stra.index(UpperCamelCase_	) + 1:]}"

        return "".join(UpperCamelCase_	)

    # matching characters
    A__		 =					get_matched_characters(UpperCamelCase_ ,       UpperCamelCase_	)
    A__		 =					get_matched_characters(UpperCamelCase_ ,       UpperCamelCase_	)
    A__		 =					len(UpperCamelCase_	)

    # transposition
    A__		 =					(
        len([(ca, ca) for ca, ca in zip(UpperCamelCase_ ,       UpperCamelCase_	) if ca != ca]	) // 2
    )

    if not match_count:
        A__		 =					0.0
    else:
        A__		 =					(
            1
            / 3
            * (
                match_count / len(UpperCamelCase_	)
                + match_count / len(UpperCamelCase_	)
                + (match_count - transpositions) / match_count
            )
        )

    # common prefix up to 4 characters
    A__		 =					0
    for ca, ca in zip(stra[:4] ,       stra[:4]	):
        if ca == ca:
            prefix_len += 1
        else:
            break

    return jaro + 0.1 * prefix_len * (1 - jaro)


if __name__ == "__main__":
   import doctest

   doctest.testmod()
   print(jaro_winkler("hello", "world"))
