


'''simple docstring'''

import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, XLMRobertaTokenizer

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()





class 						__SCREAMING_SNAKE_CASE					(      lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__ ,			unittest.TestCase ):
     UpperCAmelCase 					=				AltDiffusionPipeline
     UpperCAmelCase 					=				TEXT_TO_IMAGE_PARAMS
     UpperCAmelCase 					=				TEXT_TO_IMAGE_BATCH_PARAMS
     UpperCAmelCase 					=				TEXT_TO_IMAGE_IMAGE_PARAMS
     UpperCAmelCase 					=				TEXT_TO_IMAGE_IMAGE_PARAMS






     def  a_					(			self		)       ->    Optional[int]:
          torch.manual_seed(0		)
          _a						=		UNetaDConditionModel(
              block_out_channels=(32, 64)     ,							layers_per_block=2     ,							sample_size=32     ,							in_channels=4     ,							out_channels=4     ,							down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")     ,							up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")     ,							cross_attention_dim=32     ,							)
          _a						=		DDIMScheduler(
              beta_start=0.0_00_85     ,							beta_end=0.0_12     ,							beta_schedule="scaled_linear"     ,							clip_sample=__UpperCamelCase     ,							set_alpha_to_one=__UpperCamelCase     ,							)
          torch.manual_seed(0		)
          _a						=		AutoencoderKL(
              block_out_channels=[32, 64]     ,							in_channels=3     ,							out_channels=3     ,							down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]     ,							up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]     ,							latent_channels=4     ,							)

          # TODO: address the non-deterministic text encoder (fails for save-load tests)
          # torch.manual_seed(0)
          # text_encoder_config = RobertaSeriesConfig(
          #     hidden_size=32,
          #     project_dim=32,
          #     intermediate_size=37,
          #     layer_norm_eps=1e-05,
          #     num_attention_heads=4,
          #     num_hidden_layers=5,
          #     vocab_size=5002,
          # )
          # text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

          torch.manual_seed(0		)
          _a						=		CLIPTextConfig(
              bos_token_id=0     ,							eos_token_id=2     ,							hidden_size=32     ,							projection_dim=32     ,							intermediate_size=37     ,							layer_norm_eps=1e-05     ,							num_attention_heads=4     ,							num_hidden_layers=5     ,							pad_token_id=1     ,							vocab_size=5_002     ,							)
          _a						=		CLIPTextModel(__UpperCamelCase		)

          _a						=		XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta"		)
          _a						=		77

          _a						=		{
              "unet": unet,
              "scheduler": scheduler,
              "vae": vae,
              "text_encoder": text_encoder,
              "tokenizer": tokenizer,
              "safety_checker": None,
              "feature_extractor": None,
          }
          return components






     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase=0		)       ->    List[Any]:
          if str(__UpperCamelCase		).startswith("mps"		):
               _a						=		torch.manual_seed(__UpperCamelCase		)
          else:
               _a						=		torch.Generator(device=__UpperCamelCase		).manual_seed(__UpperCamelCase		)
          _a						=		{
              "prompt": "A painting of a squirrel eating a burger",
              "generator": generator,
              "num_inference_steps": 2,
              "guidance_scale": 6.0,
              "output_type": "numpy",
          }
          return inputs






     def  a_					(			self		)       ->    Optional[Any]:
          super().test_attention_slicing_forward_pass(expected_max_diff=3e-3		)






     def  a_					(			self		)       ->    Optional[int]:
          super().test_inference_batch_single_identical(expected_max_diff=3e-3		)






     def  a_					(			self		)       ->    List[Any]:
          _a						=		"cpu"  # ensure determinism for the device-dependent torch.Generator

          _a						=		self.get_dummy_components()
          torch.manual_seed(0		)
          _a						=		RobertaSeriesConfig(
              hidden_size=32     ,							project_dim=32     ,							intermediate_size=37     ,							layer_norm_eps=1e-05     ,							num_attention_heads=4     ,							num_hidden_layers=5     ,							vocab_size=5_002     ,							)
          # TODO: remove after fixing the non-deterministic text encoder
          _a						=		RobertaSeriesModelWithTransformation(__UpperCamelCase		)
          _a						=		text_encoder

          _a						=		AltDiffusionPipeline(**__UpperCamelCase		)
          _a						=		alt_pipe.to(__UpperCamelCase		)
          alt_pipe.set_progress_bar_config(disable=__UpperCamelCase		)

          _a						=		self.get_dummy_inputs(__UpperCamelCase		)
          _a						=		"A photo of an astronaut"
          _a						=		alt_pipe(**__UpperCamelCase		)
          _a						=		output.images
          _a						=		image[0, -3:, -3:, -1]

          assert image.shape == (1, 64, 64, 3)
          _a						=		np.array(
              [0.5_74_81_62, 0.60_44_71_45, 0.48_82_12_17, 0.50_10_06_36, 0.5_43_11_85, 0.45_76_36_83, 0.49_65_76_96, 0.48_13_27_33, 0.47_57_30_93]		)

          assert np.abs(image_slice.flatten() - expected_slice		).max() < 1e-2






     def  a_					(			self		)       ->    Tuple:
          _a						=		"cpu"  # ensure determinism for the device-dependent torch.Generator

          _a						=		self.get_dummy_components()
          _a						=		PNDMScheduler(skip_prk_steps=__UpperCamelCase		)
          torch.manual_seed(0		)
          _a						=		RobertaSeriesConfig(
              hidden_size=32     ,							project_dim=32     ,							intermediate_size=37     ,							layer_norm_eps=1e-05     ,							num_attention_heads=4     ,							num_hidden_layers=5     ,							vocab_size=5_002     ,							)
          # TODO: remove after fixing the non-deterministic text encoder
          _a						=		RobertaSeriesModelWithTransformation(__UpperCamelCase		)
          _a						=		text_encoder
          _a						=		AltDiffusionPipeline(**__UpperCamelCase		)
          _a						=		alt_pipe.to(__UpperCamelCase		)
          alt_pipe.set_progress_bar_config(disable=__UpperCamelCase		)

          _a						=		self.get_dummy_inputs(__UpperCamelCase		)
          _a						=		alt_pipe(**__UpperCamelCase		)
          _a						=		output.images
          _a						=		image[0, -3:, -3:, -1]

          assert image.shape == (1, 64, 64, 3)
          _a						=		np.array(
              [0.51_60_50_93, 0.5_70_72_41, 0.47_36_55_07, 0.50_57_88_86, 0.5_63_38_77, 0.4_64_25_03, 0.5_18_20_81, 0.48_76_34_84, 0.49_08_42_37]		)

          assert np.abs(image_slice.flatten() - expected_slice		).max() < 1e-2







@slow
@require_torch_gpu
class 						__SCREAMING_SNAKE_CASE					(      unittest.TestCase ):






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






     def  a_					(			self		)       ->    List[str]:
          # make sure here that pndm scheduler skips prk
          _a						=		AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion"     ,							safety_checker=__UpperCamelCase		)
          _a						=		alt_pipe.to(__UpperCamelCase		)
          alt_pipe.set_progress_bar_config(disable=__UpperCamelCase		)

          _a						=		"A painting of a squirrel eating a burger"
          _a						=		torch.manual_seed(0		)
          _a						=		alt_pipe([prompt]     ,							generator=__UpperCamelCase     ,							guidance_scale=6.0     ,							num_inference_steps=20     ,							output_type="np"		)

          _a						=		output.images

          _a						=		image[0, -3:, -3:, -1]

          assert image.shape == (1, 512, 512, 3)
          _a						=		np.array([0.10_10, 0.08_00, 0.07_94, 0.08_85, 0.08_43, 0.07_62, 0.07_69, 0.07_29, 0.05_86]		)

          assert np.abs(image_slice.flatten() - expected_slice		).max() < 1e-2






     def  a_					(			self		)       ->    List[str]:
          _a						=		DDIMScheduler.from_pretrained("BAAI/AltDiffusion"     ,							subfolder="scheduler"		)

          _a						=		AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion"     ,							scheduler=__UpperCamelCase     ,							safety_checker=__UpperCamelCase		)
          _a						=		alt_pipe.to(__UpperCamelCase		)
          alt_pipe.set_progress_bar_config(disable=__UpperCamelCase		)

          _a						=		"A painting of a squirrel eating a burger"
          _a						=		torch.manual_seed(0		)

          _a						=		alt_pipe([prompt]     ,							generator=__UpperCamelCase     ,							num_inference_steps=2     ,							output_type="numpy"		)
          _a						=		output.images

          _a						=		image[0, -3:, -3:, -1]

          assert image.shape == (1, 512, 512, 3)
          _a						=		np.array([0.40_19, 0.40_52, 0.38_10, 0.41_19, 0.39_16, 0.39_82, 0.46_51, 0.41_95, 0.53_23]		)

          assert np.abs(image_slice.flatten() - expected_slice		).max() < 1e-2





'''simple docstring'''

import argparse
import os

import numpy as np
import tensorflow as tf
import torch

from transformers import BertModel






def __UpperCamelCase   (			__lowerCamelCase		: BertModel      ,  __lowerCamelCase		: str      ,  __lowerCamelCase		: str				)      ->     List[str]:







     '''simple docstring'''






     _a						=		("dense.weight", "attention.self.query", "attention.self.key", "attention.self.value")

     _a						=		(
         ("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				)

     _a						=		model.state_dict()

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

     def create_tf_var(__lowerCamelCase		: np.ndarray      ,  __lowerCamelCase		: str      ,  __lowerCamelCase		: tf.Session				):
          _a						=		tf.dtypes.as_dtype(tensor.dtype				)
          _a						=		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:
               _a						=		to_tf_var_name(__lowerCamelCase				)
               _a						=		state_dict[var_name].numpy()
               if any(x in var_name for x in tensors_to_transpose				):
                    _a						=		torch_tensor.T
               _a						=		create_tf_var(tensor=__lowerCamelCase      ,  name=__lowerCamelCase      ,  session=__lowerCamelCase				)
               tf.keras.backend.set_value(__lowerCamelCase      ,  __lowerCamelCase				)
               _a						=		session.run(__lowerCamelCase				)
               print(F"Successfully created {tf_name}: {np.allclose(__lowerCamelCase      ,  __lowerCamelCase				)}"				)

          _a						=		tf.train.Saver(tf.trainable_variables()				)
          saver.save(__lowerCamelCase      ,  os.path.join(__lowerCamelCase      ,  model_name.replace("-"      ,  "_"				) + ".ckpt"				)				)






def __UpperCamelCase   (			__lowerCamelCase		: str=None				)      ->     Optional[int]:







     '''simple docstring'''






     _a						=		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"				)
     _a						=		parser.parse_args(__lowerCamelCase				)

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

