





'''simple docstring'''




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 			_lowerCAmelCase		(			A__							,     A__							,     unittest.TestCase	):






     """simple docstring"""






     snake_case_						      =   AutoencoderKL
     snake_case_						      =   "sample"
     snake_case_						      =   1E-2



     @property
     def      lowerCAmelCase			(   self							:			Tuple						)->	Any:
        snake_case 						=   4
        snake_case 						=   3
        snake_case 						=   (32, 32)

        snake_case 						=   floats_tensor((batch_size, num_channels) + sizes						).to(__snake_case						)

        return {"sample": image}



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



     @property
     def      lowerCAmelCase			(   self							:			str						)->	Any:
        return (3, 32, 32)



     def      lowerCAmelCase			(   self							:			List[Any]						)->	List[str]:
        snake_case 						=   {
            """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 						=   self.dummy_input
        return init_dict, inputs_dict



     def      lowerCAmelCase			(   self							:			List[Any]						)->	Optional[int]:
        pass



     def      lowerCAmelCase			(   self							:			List[Any]						)->	Any:
        pass



     @unittest.skipIf(torch_device == """mps"""   ,       """Gradient checkpointing skipped on MPS"""						)
     def      lowerCAmelCase			(   self							:			List[Any]						)->	str:
        # enable deterministic behavior for gradient checkpointing
        snake_case						,    snake_case 						=   self.prepare_init_args_and_inputs_for_common()
        snake_case 						=   self.model_class(**__snake_case						)
        model.to(__snake_case						)

        assert not model.is_gradient_checkpointing and model.training

        snake_case 						=   model(**__snake_case						).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 						=   torch.randn_like(__snake_case						)
        snake_case 						=   (out - labels).mean()
        loss.backward()

        # re-instantiate the model now enabling gradient checkpointing
        snake_case 						=   self.model_class(**__snake_case						)
        # clone model
        model_a.load_state_dict(model.state_dict()						)
        model_a.to(__snake_case						)
        model_a.enable_gradient_checkpointing()

        assert model_a.is_gradient_checkpointing and model_a.training

        snake_case 						=   model_a(**__snake_case						).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 						=   (out_a - labels).mean()
        loss_a.backward()

        # compare the output and parameters gradients
        self.assertTrue((loss - loss_a).abs() < 1e-5						)
        snake_case 						=   dict(model.named_parameters()						)
        snake_case 						=   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      lowerCAmelCase			(   self							:			List[str]						)->	List[Any]:
        snake_case						,    snake_case 						=   AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""   ,       output_loading_info=__snake_case						)
        self.assertIsNotNone(__snake_case						)
        self.assertEqual(len(loading_info["""missing_keys"""]						)   ,       0						)

        model.to(__snake_case						)
        snake_case 						=   model(**self.dummy_input						)

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



     def      lowerCAmelCase			(   self							:			str						)->	Optional[int]:
        snake_case 						=   AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""						)
        snake_case 						=   model.to(__snake_case						)
        model.eval()

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

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

        snake_case 						=   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 						=   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 						=   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 						=   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(__snake_case   ,       __snake_case   ,       rtol=1e-2						)						)




@slow
class 			_lowerCAmelCase		(			unittest.TestCase	):






     """simple docstring"""






     def      lowerCAmelCase			(   self							:			Tuple   ,       __snake_case							:			Tuple   ,       __snake_case							:			int						)->	List[str]:
        return f'''gaussian_noise_s={seed}_shape={"_".join([str(__snake_case						) for s in shape]						)}.npy'''



     def      lowerCAmelCase			(   self							:			List[Any]						)->	Any:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()



     def      lowerCAmelCase			(   self							:			Optional[Any]   ,       __snake_case							:			List[str]=0   ,       __snake_case							:			Union[str, Any]=(4, 3, 5_12, 5_12)   ,       __snake_case							:			Any=False						)->	str:
        snake_case 						=   torch.floataa if fpaa else torch.floataa
        snake_case 						=   torch.from_numpy(load_hf_numpy(self.get_file_format(__snake_case   ,       __snake_case						)						)						).to(__snake_case						).to(__snake_case						)
        return image



     def      lowerCAmelCase			(   self							:			Optional[Any]   ,       __snake_case							:			Optional[int]="CompVis/stable-diffusion-v1-4"   ,       __snake_case							:			Optional[int]=False						)->	List[str]:
        snake_case 						=   """fp16""" if fpaa else None
        snake_case 						=   torch.floataa if fpaa else torch.floataa

        snake_case 						=   AutoencoderKL.from_pretrained(
            __snake_case   ,       subfolder="""vae"""   ,       torch_dtype=__snake_case   ,       revision=__snake_case   ,       )
        model.to(__snake_case						).eval()

        return model



     def      lowerCAmelCase			(   self							:			List[str]   ,       __snake_case							:			Union[str, Any]=0						)->	str:
        if torch_device == "mps":
           return torch.manual_seed(__snake_case						)
        return torch.Generator(device=__snake_case						).manual_seed(__snake_case						)



     @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      lowerCAmelCase			(   self							:			Optional[Any]   ,       __snake_case							:			Optional[int]   ,       __snake_case							:			Dict   ,       __snake_case							:			Dict						)->	Union[str, Any]:
        snake_case 						=   self.get_sd_vae_model()
        snake_case 						=   self.get_sd_image(__snake_case						)
        snake_case 						=   self.get_generator(__snake_case						)

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

        assert sample.shape == image.shape

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

        assert torch_all_close(__snake_case   ,       __snake_case   ,       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      lowerCAmelCase			(   self							:			List[str]   ,       __snake_case							:			Union[str, Any]   ,       __snake_case							:			Tuple						)->	List[Any]:
        snake_case 						=   self.get_sd_vae_model(fpaa=__snake_case						)
        snake_case 						=   self.get_sd_image(__snake_case   ,       fpaa=__snake_case						)
        snake_case 						=   self.get_generator(__snake_case						)

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

        assert sample.shape == image.shape

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

        assert torch_all_close(__snake_case   ,       __snake_case   ,       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      lowerCAmelCase			(   self							:			List[str]   ,       __snake_case							:			Any   ,       __snake_case							:			Tuple   ,       __snake_case							:			str						)->	int:
        snake_case 						=   self.get_sd_vae_model()
        snake_case 						=   self.get_sd_image(__snake_case						)

        with torch.no_grad():
           snake_case 						=   model(__snake_case						).sample

        assert sample.shape == image.shape

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

        assert torch_all_close(__snake_case   ,       __snake_case   ,       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      lowerCAmelCase			(   self							:			Tuple   ,       __snake_case							:			Optional[int]   ,       __snake_case							:			Any						)->	Tuple:
        snake_case 						=   self.get_sd_vae_model()
        snake_case 						=   self.get_sd_image(__snake_case   ,       shape=(3, 4, 64, 64)						)

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

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

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

        assert torch_all_close(__snake_case   ,       __snake_case   ,       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      lowerCAmelCase			(   self							:			str   ,       __snake_case							:			str   ,       __snake_case							:			str						)->	Union[str, Any]:
        snake_case 						=   self.get_sd_vae_model(fpaa=__snake_case						)
        snake_case 						=   self.get_sd_image(__snake_case   ,       shape=(3, 4, 64, 64)   ,       fpaa=__snake_case						)

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

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

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

        assert torch_all_close(__snake_case   ,       __snake_case   ,       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      lowerCAmelCase			(   self							:			Union[str, Any]   ,       __snake_case							:			int						)->	str:
        snake_case 						=   self.get_sd_vae_model(fpaa=__snake_case						)
        snake_case 						=   self.get_sd_image(__snake_case   ,       shape=(3, 4, 64, 64)   ,       fpaa=__snake_case						)

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

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

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

        assert torch_all_close(__snake_case   ,       __snake_case   ,       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      lowerCAmelCase			(   self							:			List[str]   ,       __snake_case							:			List[str]						)->	Any:
        snake_case 						=   self.get_sd_vae_model()
        snake_case 						=   self.get_sd_image(__snake_case   ,       shape=(3, 4, 64, 64)						)

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

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

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

        assert torch_all_close(__snake_case   ,       __snake_case   ,       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      lowerCAmelCase			(   self							:			Optional[int]   ,       __snake_case							:			Optional[Any]   ,       __snake_case							:			int						)->	Tuple:
        snake_case 						=   self.get_sd_vae_model()
        snake_case 						=   self.get_sd_image(__snake_case						)
        snake_case 						=   self.get_generator(__snake_case						)

        with torch.no_grad():
           snake_case 						=   model.encode(__snake_case						).latent_dist
           snake_case 						=   dist.sample(generator=__snake_case						)

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

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

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








'''simple docstring'''




import unittest

from transformers import is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, require_torch, slow


if is_flax_available():
 import optax
 from flax.training.common_utils import onehot

 from transformers import AutoTokenizer, FlaxMTaForConditionalGeneration
 from transformers.models.ta.modeling_flax_ta import shift_tokens_right
@require_torch
@require_sentencepiece
@require_tokenizers
@require_flax
class 			_lowerCAmelCase		(			unittest.TestCase	):






     """simple docstring"""






     @slow
     def      lowerCAmelCase			(   self							:			Union[str, Any]						)->	Optional[Any]:
        snake_case 						=   FlaxMTaForConditionalGeneration.from_pretrained("""google/mt5-small"""						)
        snake_case 						=   AutoTokenizer.from_pretrained("""google/mt5-small"""						)

        snake_case 						=   tokenizer("""Hello there"""   ,       return_tensors="""np"""						).input_ids
        snake_case 						=   tokenizer("""Hi I am"""   ,       return_tensors="""np"""						).input_ids

        snake_case 						=   shift_tokens_right(__snake_case   ,       model.config.pad_token_id   ,       model.config.decoder_start_token_id						)

        snake_case 						=   model(__snake_case   ,       decoder_input_ids=__snake_case						).logits
        snake_case 						=   optax.softmax_cross_entropy(__snake_case   ,       onehot(__snake_case   ,       logits.shape[-1]						)						).mean()

        snake_case 						=   -(labels.shape[-1] * loss.item())

        snake_case 						=   -84.91_27
        self.assertTrue(abs(mtf_score - EXPECTED_SCORE						) < 1e-4						)

