



'''simple docstring'''





import gc
import random
import unittest

import torch

from diffusers import (
    IFImgaImgPipeline,
    IFImgaImgSuperResolutionPipeline,
    IFInpaintingPipeline,
    IFInpaintingSuperResolutionPipeline,
    IFPipeline,
    IFSuperResolutionPipeline,
)
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import floats_tensor, load_numpy, require_torch_gpu, skip_mps, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
from . import IFPipelineTesterMixin






@skip_mps
class 					lowerCamelCase_    (				SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_   ,   unittest.TestCase     ):
  """simple docstring"""




  a_												=IFPipeline
  a_												=TEXT_TO_IMAGE_PARAMS - {"""width""", """height""", """latents"""}
  a_												=TEXT_TO_IMAGE_BATCH_PARAMS
  a_												=PipelineTesterMixin.required_optional_params - {"""latents"""}




  def 						_lowercase     ( self      :	int   )					->					int:
    return self._get_dummy_components()




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

    __lowerCamelCase       :       int           = {
        'prompt': 'A painting of a squirrel eating a burger',
        'generator': generator,
        'num_inference_steps': 2,
        'output_type': 'numpy',
    }

    return inputs




  def 						_lowercase     ( self      :	List[str]   )					->					str:
    self._test_save_load_optional_components()




  @unittest.skipIf(torch_device != 'cuda'			,	reason='float16 requires CUDA'   )
  def 						_lowercase     ( self      :	Optional[Any]   )					->					Union[str, Any]:
    # Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
    super().test_save_load_floataa(expected_max_diff=1e-1   )




  def 						_lowercase     ( self      :	Dict   )					->					Optional[Any]:
    self._test_attention_slicing_forward_pass(expected_max_diff=1e-2   )




  def 						_lowercase     ( self      :	Any   )					->					Union[str, Any]:
    self._test_save_load_local()




  def 						_lowercase     ( self      :	int   )					->					List[str]:
    self._test_inference_batch_single_identical(
        expected_max_diff=1e-2			,	)




  @unittest.skipIf(
      torch_device != 'cuda' or not is_xformers_available()			,	reason='XFormers attention is only available with CUDA and `xformers` installed'			,	)
  def 						_lowercase     ( self      :	Dict   )					->					List[str]:
    self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3   )






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




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




  def 						_lowercase     ( self      :	str   )					->					str:
    # if

    __lowerCamelCase       :       Any           = IFPipeline.from_pretrained('DeepFloyd/IF-I-XL-v1.0'			,	variant='fp16'			,	torch_dtype=torch.floataa   )

    __lowerCamelCase       :       int           = IFSuperResolutionPipeline.from_pretrained(
        'DeepFloyd/IF-II-L-v1.0'			,	variant='fp16'			,	torch_dtype=torch.floataa			,	text_encoder=_a			,	tokenizer=_a   )

    # pre compute text embeddings and remove T5 to save memory

    pipe_a.text_encoder.to('cuda'   )

    __lowerCamelCase       ,__lowerCamelCase       :       Optional[Any]           = pipe_a.encode_prompt('anime turtle'			,	device='cuda'   )

    del pipe_a.tokenizer
    del pipe_a.text_encoder
    gc.collect()

    __lowerCamelCase       :       List[str]           = None
    __lowerCamelCase       :       List[str]           = None

    pipe_a.enable_model_cpu_offload()
    pipe_a.enable_model_cpu_offload()

    pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()   )
    pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()   )

    self._test_if(_a			,	_a			,	_a			,	_a   )

    pipe_a.remove_all_hooks()
    pipe_a.remove_all_hooks()

    # img2img

    __lowerCamelCase       :       str           = IFImgaImgPipeline(**pipe_a.components   )
    __lowerCamelCase       :       Tuple           = IFImgaImgSuperResolutionPipeline(**pipe_a.components   )

    pipe_a.enable_model_cpu_offload()
    pipe_a.enable_model_cpu_offload()

    pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()   )
    pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()   )

    self._test_if_imgaimg(_a			,	_a			,	_a			,	_a   )

    pipe_a.remove_all_hooks()
    pipe_a.remove_all_hooks()

    # inpainting

    __lowerCamelCase       :       Dict           = IFInpaintingPipeline(**pipe_a.components   )
    __lowerCamelCase       :       Any           = IFInpaintingSuperResolutionPipeline(**pipe_a.components   )

    pipe_a.enable_model_cpu_offload()
    pipe_a.enable_model_cpu_offload()

    pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()   )
    pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()   )

    self._test_if_inpainting(_a			,	_a			,	_a			,	_a   )




  def 						_lowercase     ( self      :	List[Any]			,	_a      :	List[Any]			,	_a      :	Optional[int]			,	_a      :	List[Any]			,	_a      :	Tuple   )					->					List[str]:
    # pipeline 1

    _start_torch_memory_measurement()

    __lowerCamelCase       :       Optional[Any]           = torch.Generator(device='cpu'   ).manual_seed(0   )
    __lowerCamelCase       :       Optional[int]           = pipe_a(
        prompt_embeds=_a			,	negative_prompt_embeds=_a			,	num_inference_steps=2			,	generator=_a			,	output_type='np'			,	)

    __lowerCamelCase       :       Union[str, Any]           = output.images[0]

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

    __lowerCamelCase       :       int           = torch.cuda.max_memory_allocated()
    assert mem_bytes < 13 * 10**9

    __lowerCamelCase       :       Union[str, Any]           = load_numpy(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy'   )
    assert_mean_pixel_difference(_a			,	_a   )

    # pipeline 2

    _start_torch_memory_measurement()

    __lowerCamelCase       :       List[str]           = torch.Generator(device='cpu'   ).manual_seed(0   )

    __lowerCamelCase       :       Optional[int]           = floats_tensor((1, 3, 64, 64)			,	rng=random.Random(0   )   ).to(_a   )

    __lowerCamelCase       :       Optional[int]           = pipe_a(
        prompt_embeds=_a			,	negative_prompt_embeds=_a			,	image=_a			,	generator=_a			,	num_inference_steps=2			,	output_type='np'			,	)

    __lowerCamelCase       :       int           = output.images[0]

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

    __lowerCamelCase       :       Any           = torch.cuda.max_memory_allocated()
    assert mem_bytes < 4 * 10**9

    __lowerCamelCase       :       Optional[int]           = load_numpy(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy'   )
    assert_mean_pixel_difference(_a			,	_a   )




  def 						_lowercase     ( self      :	int			,	_a      :	Any			,	_a      :	Any			,	_a      :	Optional[int]			,	_a      :	List[Any]   )					->					List[str]:
    # pipeline 1

    _start_torch_memory_measurement()

    __lowerCamelCase       :       Optional[int]           = floats_tensor((1, 3, 64, 64)			,	rng=random.Random(0   )   ).to(_a   )

    __lowerCamelCase       :       Union[str, Any]           = torch.Generator(device='cpu'   ).manual_seed(0   )

    __lowerCamelCase       :       Optional[Any]           = pipe_a(
        prompt_embeds=_a			,	negative_prompt_embeds=_a			,	image=_a			,	num_inference_steps=2			,	generator=_a			,	output_type='np'			,	)

    __lowerCamelCase       :       List[str]           = output.images[0]

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

    __lowerCamelCase       :       List[str]           = torch.cuda.max_memory_allocated()
    assert mem_bytes < 10 * 10**9

    __lowerCamelCase       :       Union[str, Any]           = load_numpy(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy'   )
    assert_mean_pixel_difference(_a			,	_a   )

    # pipeline 2

    _start_torch_memory_measurement()

    __lowerCamelCase       :       str           = torch.Generator(device='cpu'   ).manual_seed(0   )

    __lowerCamelCase       :       Optional[Any]           = floats_tensor((1, 3, 256, 256)			,	rng=random.Random(0   )   ).to(_a   )
    __lowerCamelCase       :       Any           = floats_tensor((1, 3, 64, 64)			,	rng=random.Random(0   )   ).to(_a   )

    __lowerCamelCase       :       int           = pipe_a(
        prompt_embeds=_a			,	negative_prompt_embeds=_a			,	image=_a			,	original_image=_a			,	generator=_a			,	num_inference_steps=2			,	output_type='np'			,	)

    __lowerCamelCase       :       List[str]           = output.images[0]

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

    __lowerCamelCase       :       Union[str, Any]           = torch.cuda.max_memory_allocated()
    assert mem_bytes < 4 * 10**9

    __lowerCamelCase       :       Tuple           = load_numpy(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img_superresolution_stage_II.npy'   )
    assert_mean_pixel_difference(_a			,	_a   )




  def 						_lowercase     ( self      :	List[Any]			,	_a      :	Tuple			,	_a      :	Optional[Any]			,	_a      :	Union[str, Any]			,	_a      :	Any   )					->					Union[str, Any]:
    # pipeline 1

    _start_torch_memory_measurement()

    __lowerCamelCase       :       Optional[Any]           = floats_tensor((1, 3, 64, 64)			,	rng=random.Random(0   )   ).to(_a   )
    __lowerCamelCase       :       List[Any]           = floats_tensor((1, 3, 64, 64)			,	rng=random.Random(1   )   ).to(_a   )

    __lowerCamelCase       :       Any           = torch.Generator(device='cpu'   ).manual_seed(0   )
    __lowerCamelCase       :       str           = pipe_a(
        prompt_embeds=_a			,	negative_prompt_embeds=_a			,	image=_a			,	mask_image=_a			,	num_inference_steps=2			,	generator=_a			,	output_type='np'			,	)

    __lowerCamelCase       :       int           = output.images[0]

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

    __lowerCamelCase       :       Optional[Any]           = torch.cuda.max_memory_allocated()
    assert mem_bytes < 10 * 10**9

    __lowerCamelCase       :       List[Any]           = load_numpy(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy'   )
    assert_mean_pixel_difference(_a			,	_a   )

    # pipeline 2

    _start_torch_memory_measurement()

    __lowerCamelCase       :       Optional[Any]           = torch.Generator(device='cpu'   ).manual_seed(0   )

    __lowerCamelCase       :       str           = floats_tensor((1, 3, 64, 64)			,	rng=random.Random(0   )   ).to(_a   )
    __lowerCamelCase       :       Optional[Any]           = floats_tensor((1, 3, 256, 256)			,	rng=random.Random(0   )   ).to(_a   )
    __lowerCamelCase       :       int           = floats_tensor((1, 3, 256, 256)			,	rng=random.Random(1   )   ).to(_a   )

    __lowerCamelCase       :       List[Any]           = pipe_a(
        prompt_embeds=_a			,	negative_prompt_embeds=_a			,	image=_a			,	mask_image=_a			,	original_image=_a			,	generator=_a			,	num_inference_steps=2			,	output_type='np'			,	)

    __lowerCamelCase       :       Union[str, Any]           = output.images[0]

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

    __lowerCamelCase       :       Dict           = torch.cuda.max_memory_allocated()
    assert mem_bytes < 4 * 10**9

    __lowerCamelCase       :       Optional[Any]           = load_numpy(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting_superresolution_stage_II.npy'   )
    assert_mean_pixel_difference(_a			,	_a   )




def        a_     (							)      ->      Optional[int]:
  torch.cuda.empty_cache()
  torch.cuda.reset_max_memory_allocated()
  torch.cuda.reset_peak_memory_stats()











'''simple docstring'''





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
_UpperCamelCase											=						data_utils.TransfoXLTokenizer
_UpperCamelCase											=						data_utils.TransfoXLCorpus
_UpperCamelCase											=						data_utils
_UpperCamelCase											=						data_utils

def        a_     (							_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       )      ->      Optional[Any]:
  if transfo_xl_dataset_file:
    # Convert a pre-processed corpus (see original TensorFlow repo)
    with open(_lowerCAmelCase       ,'rb'       ) as fp:
      __lowerCamelCase       :       Optional[Any]           = pickle.load(_lowerCAmelCase       ,encoding='latin1'       )
    # Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
    __lowerCamelCase       :       Tuple           = pytorch_dump_folder_path + '/' + VOCAB_FILES_NAMES['pretrained_vocab_file']
    print(F'Save vocabulary to {pytorch_vocab_dump_path}'       )
    __lowerCamelCase       :       str           = corpus.vocab.__dict__
    torch.save(_lowerCAmelCase       ,_lowerCAmelCase       )

    __lowerCamelCase       :       Optional[Any]           = corpus.__dict__
    corpus_dict_no_vocab.pop('vocab'       ,_lowerCAmelCase       )
    __lowerCamelCase       :       Optional[Any]           = 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       :       int           = os.path.abspath(_lowerCAmelCase       )
    __lowerCamelCase       :       Any           = 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       :       Optional[int]           = TransfoXLConfig()
    else:
      __lowerCamelCase       :       Optional[int]           = TransfoXLConfig.from_json_file(_lowerCAmelCase       )
    print(F'Building PyTorch model from configuration: {config}'       )
    __lowerCamelCase       :       List[str]           = TransfoXLLMHeadModel(_lowerCAmelCase       )

    __lowerCamelCase       :       Dict           = load_tf_weights_in_transfo_xl(_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       )
    # Save pytorch-model
    __lowerCamelCase       :       List[str]           = os.path.join(_lowerCAmelCase       ,_lowerCAmelCase       )
    __lowerCamelCase       :       int           = 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__":
   _UpperCamelCase											=						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.',
   )
   _UpperCamelCase											=						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,
   )






