



'''simple docstring'''





import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModel,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableUnCLIPImgaImgPipeline, UNetaDConditionModel
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    floats_tensor,
    load_image,
    load_numpy,
    require_torch_gpu,
    skip_mps,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()






class    lowerCamelCase_       (							_a  ,  _a  ,  _a  ,  unittest.TestCase	):
  """simple docstring"""




  a_				   =StableUnCLIPImgaImgPipeline
  a_				   =TEXT_GUIDED_IMAGE_VARIATION_PARAMS
  a_				   =TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
  a_				   =frozenset(
      []	)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
  a_				   =frozenset([]	)




  def 		_lowercase	(					self :      List[Any]  )   ->   List[str]:
    __lowerCamelCase			:			Dict       			=					32
    __lowerCamelCase			:			Tuple       			=					embedder_hidden_size

    # image encoding components

    __lowerCamelCase			:			Any       			=					CLIPImageProcessor(crop_size=32     ,      size=32  )

    torch.manual_seed(0  )
    __lowerCamelCase			:			str       			=					CLIPVisionModelWithProjection(
        CLIPVisionConfig(
            hidden_size=__lowerCamelCase     ,      projection_dim=__lowerCamelCase     ,      num_hidden_layers=5     ,      num_attention_heads=4     ,      image_size=32     ,      intermediate_size=37     ,      patch_size=1     ,      )  )

    # regular denoising components

    torch.manual_seed(0  )
    __lowerCamelCase			:			int       			=					StableUnCLIPImageNormalizer(embedding_dim=__lowerCamelCase  )
    __lowerCamelCase			:			int       			=					DDPMScheduler(beta_schedule='squaredcos_cap_v2'  )

    torch.manual_seed(0  )
    __lowerCamelCase			:			Optional[Any]       			=					CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'  )

    torch.manual_seed(0  )
    __lowerCamelCase			:			Union[str, Any]       			=					CLIPTextModel(
        CLIPTextConfig(
            bos_token_id=0     ,      eos_token_id=2     ,      hidden_size=__lowerCamelCase     ,      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=1000     ,      )  )

    torch.manual_seed(0  )
    __lowerCamelCase			:			Tuple       			=					UNetaDConditionModel(
        sample_size=32     ,      in_channels=4     ,      out_channels=4     ,      down_block_types=('CrossAttnDownBlock2D', 'DownBlock2D')     ,      up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D')     ,      block_out_channels=(32, 64)     ,      attention_head_dim=(2, 4)     ,      class_embed_type='projection'     ,      projection_class_embeddings_input_dim=embedder_projection_dim * 2     ,      cross_attention_dim=__lowerCamelCase     ,      layers_per_block=1     ,      upcast_attention=__lowerCamelCase     ,      use_linear_projection=__lowerCamelCase     ,      )

    torch.manual_seed(0  )
    __lowerCamelCase			:			Optional[int]       			=					DDIMScheduler(
        beta_schedule='scaled_linear'     ,      beta_start=0.00085     ,      beta_end=0.012     ,      prediction_type='v_prediction'     ,      set_alpha_to_one=__lowerCamelCase     ,      steps_offset=1     ,      )

    torch.manual_seed(0  )
    __lowerCamelCase			:			Union[str, Any]       			=					AutoencoderKL()

    __lowerCamelCase			:			Any       			=					{
        # image encoding components
        'feature_extractor': feature_extractor,
        'image_encoder': image_encoder.eval(),
        # image noising components
        'image_normalizer': image_normalizer.eval(),
        'image_noising_scheduler': image_noising_scheduler,
        # regular denoising components
        'tokenizer': tokenizer,
        'text_encoder': text_encoder.eval(),
        'unet': unet.eval(),
        'scheduler': scheduler,
        'vae': vae.eval(),
    }

    return components




  def 		_lowercase	(					self :      str     ,      _a :      Optional[Any]     ,      _a :      Any=0     ,      _a :      List[str]=True  )   ->   List[Any]:
    if str(__lowerCamelCase  ).startswith('mps'  ):
      __lowerCamelCase			:			Dict       			=					torch.manual_seed(__lowerCamelCase  )
    else:
      __lowerCamelCase			:			Tuple       			=					torch.Generator(device=__lowerCamelCase  ).manual_seed(__lowerCamelCase  )

    __lowerCamelCase			:			str       			=					floats_tensor((1, 3, 32, 32)     ,      rng=random.Random(__lowerCamelCase  )  ).to(__lowerCamelCase  )

    if pil_image:
      __lowerCamelCase			:			Optional[Any]       			=					input_image * 0.5 + 0.5
      __lowerCamelCase			:			Any       			=					input_image.clamp(0     ,      1  )
      __lowerCamelCase			:			str       			=					input_image.cpu().permute(0     ,      2     ,      3     ,      1  ).float().numpy()
      __lowerCamelCase			:			Union[str, Any]       			=					DiffusionPipeline.numpy_to_pil(__lowerCamelCase  )[0]

    return {
        "prompt": "An anime racoon running a marathon",
        "image": input_image,
        "generator": generator,
        "num_inference_steps": 2,
        "output_type": "np",
    }




  @skip_mps
  def 		_lowercase	(					self :      Dict  )   ->   List[str]:
    __lowerCamelCase			:			List[Any]       			=					'cpu'  # ensure determinism for the device-dependent torch.Generator
    __lowerCamelCase			:			Optional[int]       			=					self.get_dummy_components()
    __lowerCamelCase			:			Optional[Any]       			=					StableUnCLIPImgaImgPipeline(**__lowerCamelCase  )
    __lowerCamelCase			:			List[Any]       			=					sd_pipe.to(__lowerCamelCase  )
    sd_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

    __lowerCamelCase			:			Any       			=					self.get_dummy_inputs(__lowerCamelCase  )
    inputs.update({'image_embeds': None}  )
    __lowerCamelCase			:			Tuple       			=					sd_pipe(**__lowerCamelCase  ).images
    __lowerCamelCase			:			Optional[Any]       			=					image[0, -3:, -3:, -1]

    assert image.shape == (1, 32, 32, 3)
    __lowerCamelCase			:			Any       			=					np.array([0.3872, 0.7224, 0.5601, 0.4741, 0.6872, 0.5814, 0.4636, 0.3867, 0.5078]  )

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




  def 		_lowercase	(					self :      List[Any]  )   ->   Union[str, Any]:
    __lowerCamelCase			:			str       			=					torch_device in ['cpu', 'mps']

    self._test_attention_slicing_forward_pass(test_max_difference=__lowerCamelCase  )




  def 		_lowercase	(					self :      Any  )   ->   Tuple:
    __lowerCamelCase			:			Any       			=					torch_device in ['cpu', 'mps']

    self._test_inference_batch_single_identical(test_max_difference=__lowerCamelCase  )




  @unittest.skipIf(
      torch_device != 'cuda' or not is_xformers_available()     ,      reason='XFormers attention is only available with CUDA and `xformers` installed'     ,      )
  def 		_lowercase	(					self :      Union[str, Any]  )   ->   str:
    self._test_xformers_attention_forwardGenerator_pass(test_max_difference=__lowerCamelCase  )






@slow
@require_torch_gpu
class    lowerCamelCase_       (							unittest.TestCase	):
  """simple docstring"""




  def 		_lowercase	(					self :      List[str]  )   ->   Any:
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()




  def 		_lowercase	(					self :      str  )   ->   Optional[int]:
    __lowerCamelCase			:			int       			=					load_image(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'  )

    __lowerCamelCase			:			Optional[Any]       			=					load_numpy(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_img2img_anime_turtle_fp16.npy'  )

    __lowerCamelCase			:			Optional[Any]       			=					StableUnCLIPImgaImgPipeline.from_pretrained(
        'fusing/stable-unclip-2-1-l-img2img'     ,      torch_dtype=torch.floataa  )
    pipe.to(__lowerCamelCase  )
    pipe.set_progress_bar_config(disable=__lowerCamelCase  )
    # stable unclip will oom when integration tests are run on a V100,
    # so turn on memory savings
    pipe.enable_attention_slicing()
    pipe.enable_sequential_cpu_offload()

    __lowerCamelCase			:			Any       			=					torch.Generator(device='cpu'  ).manual_seed(0  )
    __lowerCamelCase			:			Union[str, Any]       			=					pipe(__lowerCamelCase     ,      'anime turle'     ,      generator=__lowerCamelCase     ,      output_type='np'  )

    __lowerCamelCase			:			str       			=					output.images[0]

    assert image.shape == (768, 768, 3)

    assert_mean_pixel_difference(__lowerCamelCase     ,      __lowerCamelCase  )




  def 		_lowercase	(					self :      int  )   ->   Optional[int]:
    __lowerCamelCase			:			List[Any]       			=					load_image(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'  )

    __lowerCamelCase			:			Tuple       			=					load_numpy(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_h_img2img_anime_turtle_fp16.npy'  )

    __lowerCamelCase			:			int       			=					StableUnCLIPImgaImgPipeline.from_pretrained(
        'fusing/stable-unclip-2-1-h-img2img'     ,      torch_dtype=torch.floataa  )
    pipe.to(__lowerCamelCase  )
    pipe.set_progress_bar_config(disable=__lowerCamelCase  )
    # stable unclip will oom when integration tests are run on a V100,
    # so turn on memory savings
    pipe.enable_attention_slicing()
    pipe.enable_sequential_cpu_offload()

    __lowerCamelCase			:			Dict       			=					torch.Generator(device='cpu'  ).manual_seed(0  )
    __lowerCamelCase			:			Tuple       			=					pipe(__lowerCamelCase     ,      'anime turle'     ,      generator=__lowerCamelCase     ,      output_type='np'  )

    __lowerCamelCase			:			Dict       			=					output.images[0]

    assert image.shape == (768, 768, 3)

    assert_mean_pixel_difference(__lowerCamelCase     ,      __lowerCamelCase  )




  def 		_lowercase	(					self :      Optional[int]  )   ->   Optional[Any]:
    __lowerCamelCase			:			Optional[Any]       			=					load_image(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'  )

    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_peak_memory_stats()

    __lowerCamelCase			:			Optional[int]       			=					StableUnCLIPImgaImgPipeline.from_pretrained(
        'fusing/stable-unclip-2-1-h-img2img'     ,      torch_dtype=torch.floataa  )
    __lowerCamelCase			:			Tuple       			=					pipe.to(__lowerCamelCase  )
    pipe.set_progress_bar_config(disable=__lowerCamelCase  )
    pipe.enable_attention_slicing()
    pipe.enable_sequential_cpu_offload()

    __lowerCamelCase			:			Dict       			=					pipe(
        __lowerCamelCase     ,      'anime turtle'     ,      num_inference_steps=2     ,      output_type='np'     ,      )

    __lowerCamelCase			:			Tuple       			=					torch.cuda.max_memory_allocated()
    # make sure that less than 7 GB is allocated
    assert mem_bytes < 7 * 10**9












import argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

# We do this to be able to load python 2 datasets pickles
# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
__magic_name__  :  Tuple					=       data_utils.TransfoXLTokenizer
__magic_name__  :  int					=       data_utils.TransfoXLCorpus
__magic_name__  :  List[str]					=       data_utils
__magic_name__  :  str					=       data_utils
def 		a_							( __lowerCAmelCase   , __lowerCAmelCase   , __lowerCAmelCase   , __lowerCAmelCase				):
  if transfo_xl_dataset_file:
    # Convert a pre-processed corpus (see original TensorFlow repo)
    with open(__lowerCAmelCase   , '''rb'''				) as fp:
      lowerCAmelCase__										=      pickle.load(__lowerCAmelCase   , encoding='''latin1'''				)
    # Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
    lowerCAmelCase__										=      pytorch_dump_folder_path + '''/''' + VOCAB_FILES_NAMES['''pretrained_vocab_file''']
    print(F"""Save vocabulary to {pytorch_vocab_dump_path}"""				)
    lowerCAmelCase__										=      corpus.vocab.__dict__
    torch.save(__lowerCAmelCase   , __lowerCAmelCase				)

    lowerCAmelCase__										=      corpus.__dict__
    corpus_dict_no_vocab.pop('''vocab'''   , __lowerCAmelCase				)
    lowerCAmelCase__										=      pytorch_dump_folder_path + '''/''' + CORPUS_NAME
    print(F"""Save dataset to {pytorch_dataset_dump_path}"""				)
    torch.save(__lowerCAmelCase   , __lowerCAmelCase				)

  if tf_checkpoint_path:
    # Convert a pre-trained TensorFlow model
    lowerCAmelCase__										=      os.path.abspath(__lowerCAmelCase				)
    lowerCAmelCase__										=      os.path.abspath(__lowerCAmelCase				)

    print(F"""Converting Transformer XL checkpoint from {tf_path} with config at {config_path}."""				)
    # Initialise PyTorch model
    if transfo_xl_config_file == "":
      lowerCAmelCase__										=      TransfoXLConfig()
    else:
      lowerCAmelCase__										=      TransfoXLConfig.from_json_file(__lowerCAmelCase				)
    print(F"""Building PyTorch model from configuration: {config}"""				)
    lowerCAmelCase__										=      TransfoXLLMHeadModel(__lowerCAmelCase				)

    lowerCAmelCase__										=      load_tf_weights_in_transfo_xl(__lowerCAmelCase   , __lowerCAmelCase   , __lowerCAmelCase				)
    # Save pytorch-model
    lowerCAmelCase__										=      os.path.join(__lowerCAmelCase   , __lowerCAmelCase				)
    lowerCAmelCase__										=      os.path.join(__lowerCAmelCase   , __lowerCAmelCase				)
    print(F"""Save PyTorch model to {os.path.abspath(__lowerCAmelCase				)}"""				)
    torch.save(model.state_dict()   , __lowerCAmelCase				)
    print(F"""Save configuration file to {os.path.abspath(__lowerCAmelCase				)}"""				)
    with open(__lowerCAmelCase   , '''w'''   , encoding='''utf-8'''				) as f:
      f.write(config.to_json_string()				)


if __name__ == "__main__":
      __magic_name__  :  int					=       argparse.ArgumentParser()
      parser.add_argument(
          """--pytorch_dump_folder_path""",
          default=None,
          type=str,
          required=True,
          help="""Path to the folder to store the PyTorch model or dataset/vocab.""",
      )
      parser.add_argument(
          """--tf_checkpoint_path""",
          default="""""",
          type=str,
          help="""An optional path to a TensorFlow checkpoint path to be converted.""",
      )
      parser.add_argument(
          """--transfo_xl_config_file""",
          default="""""",
          type=str,
          help=(
              """An optional config json file corresponding to the pre-trained BERT model. \n"""
              """This specifies the model architecture."""
          ),
      )
      parser.add_argument(
          """--transfo_xl_dataset_file""",
          default="""""",
          type=str,
          help="""An optional dataset file to be converted in a vocabulary.""",
      )
      __magic_name__  :  Optional[Any]					=       parser.parse_args()
      convert_transfo_xl_checkpoint_to_pytorch(
          args.tf_checkpoint_path,
          args.transfo_xl_config_file,
          args.pytorch_dump_folder_path,
          args.transfo_xl_dataset_file,
      )

