



"""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  lowercase__  (       __snake_case							,   __snake_case							,   unittest.TestCase ):







       '''simple docstring'''


       UpperCamelCase		=	AutoencoderKL
       UpperCamelCase		=	'''sample'''
       UpperCamelCase		=	1E-2



       @property
       def 	lowercase__							(				self     :    Dict       )					->		Union[str, Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				4
             UpperCAmelCase_     					=				3
             UpperCAmelCase_     					=				(32, 32)

             UpperCAmelCase_     					=				floats_tensor((batch_size, num_channels) + sizes       ).to(__lowerCamelCase       )

             return {"sample": image}



       @property
       def 	lowercase__							(				self     :    str       )					->		Tuple:


             '''simple docstring'''




             return (3, 32, 32)



       @property
       def 	lowercase__							(				self     :    List[str]       )					->		Optional[int]:


             '''simple docstring'''




             return (3, 32, 32)



       def 	lowercase__							(				self     :    Any       )					->		Dict:


             '''simple docstring'''




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



       def 	lowercase__							(				self     :    Union[str, Any]       )					->		Tuple:


             '''simple docstring'''




             pass



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


             '''simple docstring'''




             pass



       @unittest.skipIf(torch_device == "mps"						,							"Gradient checkpointing skipped on MPS"       )
       def 	lowercase__							(				self     :    List[str]       )					->		Union[str, Any]:


             '''simple docstring'''




             UpperCAmelCase_					,       UpperCAmelCase_     					=				self.prepare_init_args_and_inputs_for_common()
             UpperCAmelCase_     					=				self.model_class(**__lowerCamelCase       )
             model.to(__lowerCamelCase       )

             assert not model.is_gradient_checkpointing and model.training

             UpperCAmelCase_     					=				model(**__lowerCamelCase       ).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()

             UpperCAmelCase_     					=				torch.randn_like(__lowerCamelCase       )
             UpperCAmelCase_     					=				(out - labels).mean()
             loss.backward()

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

             assert model_a.is_gradient_checkpointing and model_a.training

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

             # compare the output and parameters gradients
             self.assertTrue((loss - loss_a).abs() < 1e-5       )
             UpperCAmelCase_     					=				dict(model.named_parameters()       )
             UpperCAmelCase_     					=				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     :    int       )					->		Tuple:


             '''simple docstring'''




             UpperCAmelCase_					,       UpperCAmelCase_     					=				AutoencoderKL.from_pretrained("fusing/autoencoder-kl-dummy"						,							output_loading_info=__lowerCamelCase       )
             self.assertIsNotNone(__lowerCamelCase       )
             self.assertEqual(len(loading_info["missing_keys"]       )						,							0       )

             model.to(__lowerCamelCase       )
             UpperCAmelCase_     					=				model(**self.dummy_input       )

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







       def 	lowercase__							(				self     :    Optional[Any]       )					->		List[Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				AutoencoderKL.from_pretrained("fusing/autoencoder-kl-dummy"       )
             UpperCAmelCase_     					=				model.to(__lowerCamelCase       )
             model.eval()

             if torch_device == "mps":
                   UpperCAmelCase_     					=				torch.manual_seed(0       )
             else:
                   UpperCAmelCase_     					=				torch.Generator(device=__lowerCamelCase       ).manual_seed(0       )

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

             UpperCAmelCase_     					=				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":
                   UpperCAmelCase_     					=				torch.tensor(
                       [
                           -4.0078e-01,
                           -3.8323e-04,
                           -1.2681e-01,
                           -1.1462e-01,
                           2.0095e-01,
                           1.0893e-01,
                           -8.8247e-02,
                           -3.0361e-01,
                           -9.8644e-03,
                       ]       )
             elif torch_device == "cpu":
                   UpperCAmelCase_     					=				torch.tensor(
                       [-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026]       )
             else:
                   UpperCAmelCase_     					=				torch.tensor(
                       [-0.2421, 0.4642, 0.2507, -0.0438, 0.0682, 0.3160, -0.2018, -0.0727, 0.2485]       )

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



@slow
class  lowercase__  (       unittest.TestCase ):







       '''simple docstring'''


       def 	lowercase__							(				self     :    int						,							_UpperCAmelCase     :    Union[str, Any]						,							_UpperCAmelCase     :    int       )					->		Optional[Any]:


             '''simple docstring'''




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



       def 	lowercase__							(				self     :    Union[str, Any]       )					->		Dict:


             '''simple docstring'''




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



       def 	lowercase__							(				self     :    Optional[int]						,							_UpperCAmelCase     :    Tuple=0						,							_UpperCAmelCase     :    List[Any]=(4, 3, 512, 512)						,							_UpperCAmelCase     :    List[str]=False       )					->		str:


             '''simple docstring'''




             UpperCAmelCase_     					=				torch.floataa if fpaa else torch.floataa
             UpperCAmelCase_     					=				torch.from_numpy(load_hf_numpy(self.get_file_format(__lowerCamelCase						,							__lowerCamelCase       )       )       ).to(__lowerCamelCase       ).to(__lowerCamelCase       )
             return image



       def 	lowercase__							(				self     :    Optional[Any]						,							_UpperCAmelCase     :    Optional[int]="CompVis/stable-diffusion-v1-4"						,							_UpperCAmelCase     :    Any=False       )					->		int:


             '''simple docstring'''




             UpperCAmelCase_     					=				"fp16" if fpaa else None
             UpperCAmelCase_     					=				torch.floataa if fpaa else torch.floataa

             UpperCAmelCase_     					=				AutoencoderKL.from_pretrained(
                 __lowerCamelCase						,							subfolder="vae"						,							torch_dtype=__lowerCamelCase						,							revision=__lowerCamelCase						,							)
             model.to(__lowerCamelCase       ).eval()

             return model



       def 	lowercase__							(				self     :    List[Any]						,							_UpperCAmelCase     :    Union[str, Any]=0       )					->		Any:


             '''simple docstring'''




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



       @parameterized.expand(
           [
               # fmt: off
               [33, [-0.1603, 0.9878, -0.0495, -0.0790, -0.2709, 0.8375, -0.2060, -0.0824], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
               [47, [-0.2376, 0.1168, 0.1332, -0.4840, -0.2508, -0.0791, -0.0493, -0.4089], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
               # fmt: on
           ]       )
       def 	lowercase__							(				self     :    Tuple						,							_UpperCAmelCase     :    Tuple						,							_UpperCAmelCase     :    Optional[Any]						,							_UpperCAmelCase     :    Optional[Any]       )					->		Union[str, Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.get_sd_vae_model()
             UpperCAmelCase_     					=				self.get_sd_image(__lowerCamelCase       )
             UpperCAmelCase_     					=				self.get_generator(__lowerCamelCase       )

             with torch.no_grad():
                   UpperCAmelCase_     					=				model(__lowerCamelCase						,							generator=__lowerCamelCase						,							sample_posterior=__lowerCamelCase       ).sample

             assert sample.shape == image.shape

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

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



       @parameterized.expand(
           [
               # fmt: off
               [33, [-0.0513, 0.0289, 1.3799, 0.2166, -0.2573, -0.0871, 0.5103, -0.0999]],
               [47, [-0.4128, -0.1320, -0.3704, 0.1965, -0.4116, -0.2332, -0.3340, 0.2247]],
               # fmt: on
           ]       )
       @require_torch_gpu
       def 	lowercase__							(				self     :    int						,							_UpperCAmelCase     :    int						,							_UpperCAmelCase     :    Tuple       )					->		Tuple:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.get_sd_vae_model(fpaa=__lowerCamelCase       )
             UpperCAmelCase_     					=				self.get_sd_image(__lowerCamelCase						,							fpaa=__lowerCamelCase       )
             UpperCAmelCase_     					=				self.get_generator(__lowerCamelCase       )

             with torch.no_grad():
                   UpperCAmelCase_     					=				model(__lowerCamelCase						,							generator=__lowerCamelCase						,							sample_posterior=__lowerCamelCase       ).sample

             assert sample.shape == image.shape

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

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



       @parameterized.expand(
           [
               # fmt: off
               [33, [-0.1609, 0.9866, -0.0487, -0.0777, -0.2716, 0.8368, -0.2055, -0.0814], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
               [47, [-0.2377, 0.1147, 0.1333, -0.4841, -0.2506, -0.0805, -0.0491, -0.4085], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
               # fmt: on
           ]       )
       def 	lowercase__							(				self     :    Any						,							_UpperCAmelCase     :    Any						,							_UpperCAmelCase     :    List[Any]						,							_UpperCAmelCase     :    List[Any]       )					->		List[Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.get_sd_vae_model()
             UpperCAmelCase_     					=				self.get_sd_image(__lowerCamelCase       )

             with torch.no_grad():
                   UpperCAmelCase_     					=				model(__lowerCamelCase       ).sample

             assert sample.shape == image.shape

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

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



       @parameterized.expand(
           [
               # fmt: off
               [13, [-0.2051, -0.1803, -0.2311, -0.2114, -0.3292, -0.3574, -0.2953, -0.3323]],
               [37, [-0.2632, -0.2625, -0.2199, -0.2741, -0.4539, -0.4990, -0.3720, -0.4925]],
               # fmt: on
           ]       )
       @require_torch_gpu
       def 	lowercase__							(				self     :    List[str]						,							_UpperCAmelCase     :    Optional[int]						,							_UpperCAmelCase     :    Any       )					->		Tuple:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.get_sd_vae_model()
             UpperCAmelCase_     					=				self.get_sd_image(__lowerCamelCase						,							shape=(3, 4, 64, 64)       )

             with torch.no_grad():
                   UpperCAmelCase_     					=				model.decode(__lowerCamelCase       ).sample

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

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

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



       @parameterized.expand(
           [
               # fmt: off
               [27, [-0.0369, 0.0207, -0.0776, -0.0682, -0.1747, -0.1930, -0.1465, -0.2039]],
               [16, [-0.1628, -0.2134, -0.2747, -0.2642, -0.3774, -0.4404, -0.3687, -0.4277]],
               # fmt: on
           ]       )
       @require_torch_gpu
       def 	lowercase__							(				self     :    Any						,							_UpperCAmelCase     :    Union[str, Any]						,							_UpperCAmelCase     :    str       )					->		str:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.get_sd_vae_model(fpaa=__lowerCamelCase       )
             UpperCAmelCase_     					=				self.get_sd_image(__lowerCamelCase						,							shape=(3, 4, 64, 64)						,							fpaa=__lowerCamelCase       )

             with torch.no_grad():
                   UpperCAmelCase_     					=				model.decode(__lowerCamelCase       ).sample

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

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

             assert torch_all_close(__lowerCamelCase						,							__lowerCamelCase						,							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     :    str						,							_UpperCAmelCase     :    List[Any]       )					->		Tuple:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.get_sd_vae_model(fpaa=__lowerCamelCase       )
             UpperCAmelCase_     					=				self.get_sd_image(__lowerCamelCase						,							shape=(3, 4, 64, 64)						,							fpaa=__lowerCamelCase       )

             with torch.no_grad():
                   UpperCAmelCase_     					=				model.decode(__lowerCamelCase       ).sample

             model.enable_xformers_memory_efficient_attention()
             with torch.no_grad():
                   UpperCAmelCase_     					=				model.decode(__lowerCamelCase       ).sample

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

             assert torch_all_close(__lowerCamelCase						,							__lowerCamelCase						,							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     :    Optional[int]						,							_UpperCAmelCase     :    Any       )					->		Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.get_sd_vae_model()
             UpperCAmelCase_     					=				self.get_sd_image(__lowerCamelCase						,							shape=(3, 4, 64, 64)       )

             with torch.no_grad():
                   UpperCAmelCase_     					=				model.decode(__lowerCamelCase       ).sample

             model.enable_xformers_memory_efficient_attention()
             with torch.no_grad():
                   UpperCAmelCase_     					=				model.decode(__lowerCamelCase       ).sample

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

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







       @parameterized.expand(
           [
               # fmt: off
               [33, [-0.3001, 0.0918, -2.6984, -3.9720, -3.2099, -5.0353, 1.7338, -0.2065, 3.4267]],
               [47, [-1.5030, -4.3871, -6.0355, -9.1157, -1.6661, -2.7853, 2.1607, -5.0823, 2.5633]],
               # fmt: on
           ]       )
       def 	lowercase__							(				self     :    Dict						,							_UpperCAmelCase     :    Union[str, Any]						,							_UpperCAmelCase     :    List[Any]       )					->		Any:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.get_sd_vae_model()
             UpperCAmelCase_     					=				self.get_sd_image(__lowerCamelCase       )
             UpperCAmelCase_     					=				self.get_generator(__lowerCamelCase       )

             with torch.no_grad():
                   UpperCAmelCase_     					=				model.encode(__lowerCamelCase       ).latent_dist
                   UpperCAmelCase_     					=				dist.sample(generator=__lowerCamelCase       )

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

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

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








import argparse
import os

import numpy as np
import tensorflow as tf
import torch

from transformers import BertModel
def 	_UpperCamelCase					(    lowerCAmelCase_	,    lowerCAmelCase_	,    lowerCAmelCase_ )    ->Optional[Any]:
   UpperCAmelCase      					=  ("""dense.weight""", """attention.self.query""", """attention.self.key""", """attention.self.value""")

   UpperCAmelCase      					=  (
       ("""layer.""", """layer_"""),
       ("""word_embeddings.weight""", """word_embeddings"""),
       ("""position_embeddings.weight""", """position_embeddings"""),
       ("""token_type_embeddings.weight""", """token_type_embeddings"""),
       (""".""", """/"""),
       ("""LayerNorm/weight""", """LayerNorm/gamma"""),
       ("""LayerNorm/bias""", """LayerNorm/beta"""),
       ("""weight""", """kernel"""),
   )

   if not os.path.isdir(lowerCAmelCase_ ):
      os.makedirs(lowerCAmelCase_ )

   UpperCAmelCase      					=  model.state_dict()

   def to_tf_var_name(lowerCAmelCase_ ):
      for patt, repl in iter(lowerCAmelCase_ ):
         UpperCAmelCase      					=  name.replace(lowerCAmelCase_	,    lowerCAmelCase_ )
      return F"""bert/{name}"""

   def create_tf_var(lowerCAmelCase_	,    lowerCAmelCase_	,    lowerCAmelCase_ ):
      UpperCAmelCase      					=  tf.dtypes.as_dtype(tensor.dtype )
      UpperCAmelCase      					=  tf.get_variable(dtype=lowerCAmelCase_	,    shape=tensor.shape	,    name=lowerCAmelCase_	,    initializer=tf.zeros_initializer() )
      session.run(tf.variables_initializer([tf_var] ) )
      session.run(lowerCAmelCase_ )
      return tf_var

   tf.reset_default_graph()
   with tf.Session() as session:
      for var_name in state_dict:
         UpperCAmelCase      					=  to_tf_var_name(lowerCAmelCase_ )
         UpperCAmelCase      					=  state_dict[var_name].numpy()
         if any(x in var_name for x in tensors_to_transpose ):
            UpperCAmelCase      					=  torch_tensor.T
         UpperCAmelCase      					=  create_tf_var(tensor=lowerCAmelCase_	,    name=lowerCAmelCase_	,    session=lowerCAmelCase_ )
         tf.keras.backend.set_value(lowerCAmelCase_	,    lowerCAmelCase_ )
         UpperCAmelCase      					=  session.run(lowerCAmelCase_ )
         print(F"""Successfully created {tf_name}: {np.allclose(lowerCAmelCase_	,    lowerCAmelCase_ )}""" )

      UpperCAmelCase      					=  tf.train.Saver(tf.trainable_variables() )
      saver.save(lowerCAmelCase_	,    os.path.join(lowerCAmelCase_	,    model_name.replace("""-"""	,    """_""" ) + """.ckpt""" ) )
def 	_UpperCamelCase					(    lowerCAmelCase_=None )    ->Union[str, Any]:
   UpperCAmelCase      					=  argparse.ArgumentParser()
   parser.add_argument("""--model_name"""	,    type=lowerCAmelCase_	,    required=lowerCAmelCase_	,    help="""model name e.g. bert-base-uncased""" )
   parser.add_argument(
       """--cache_dir"""	,    type=lowerCAmelCase_	,    default=lowerCAmelCase_	,    required=lowerCAmelCase_	,    help="""Directory containing pytorch model""" )
   parser.add_argument("""--pytorch_model_path"""	,    type=lowerCAmelCase_	,    required=lowerCAmelCase_	,    help="""/path/to/<pytorch-model-name>.bin""" )
   parser.add_argument("""--tf_cache_dir"""	,    type=lowerCAmelCase_	,    required=lowerCAmelCase_	,    help="""Directory in which to save tensorflow model""" )
   UpperCAmelCase      					=  parser.parse_args(lowerCAmelCase_ )

   UpperCAmelCase      					=  BertModel.from_pretrained(
       pretrained_model_name_or_path=args.model_name	,    state_dict=torch.load(args.pytorch_model_path )	,    cache_dir=args.cache_dir	,    )

   convert_pytorch_checkpoint_to_tf(model=lowerCAmelCase_	,    ckpt_dir=args.tf_cache_dir	,    model_name=args.model_name )


if __name__ == "__main__":
 main()



