
'''simple docstring'''







# This model implementation is heavily inspired by https://github.com/haofanwang/ControlNet-for-Diffusers/

import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    DDIMScheduler,
    StableDiffusionControlNetImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel
from diffusers.utils import floats_tensor, load_image, load_numpy, randn_tensor, slow, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

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


enable_full_determinism()







class UpperCAmelCase__	( lowercase__    ,				lowercase__    ,				lowercase__    ,				unittest.TestCase	):





  """simple docstring"""

  __UpperCAmelCase				:				List[Any]   =					StableDiffusionControlNetImgaImgPipeline
  __UpperCAmelCase				:				Any   =					TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
  __UpperCAmelCase				:				str   =					TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
  __UpperCAmelCase				:				Tuple   =					IMAGE_TO_IMAGE_IMAGE_PARAMS.union({'''control_image'''}	)
  __UpperCAmelCase				:				Tuple   =					IMAGE_TO_IMAGE_IMAGE_PARAMS


  def   __lowercase		(		self      :     List[str]       ):


     '''simple docstring'''






     torch.manual_seed(0       )
     _a				:	Any			=   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							,)
     torch.manual_seed(0       )
     _a				:	int			=   ControlNetModel(
         block_out_channels=(32, 64)							,layers_per_block=2							,in_channels=4							,down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')							,cross_attention_dim=32							,conditioning_embedding_out_channels=(16, 32)							,)
     torch.manual_seed(0       )
     _a				:	Any			=   DDIMScheduler(
         beta_start=0.0_0085							,beta_end=0.012							,beta_schedule='scaled_linear'							,clip_sample=_a							,set_alpha_to_one=_a							,)
     torch.manual_seed(0       )
     _a				:	List[str]			=   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							,)
     torch.manual_seed(0       )
     _a				:	Tuple			=   CLIPTextConfig(
         bos_token_id=0							,eos_token_id=2							,hidden_size=32							,intermediate_size=37							,layer_norm_eps=1E-05							,num_attention_heads=4							,num_hidden_layers=5							,pad_token_id=1							,vocab_size=1000							,)
     _a				:	int			=   CLIPTextModel(_a       )
     _a				:	Union[str, Any]			=   CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )

     _a				:	List[str]			=   {
         'unet': unet,
         'controlnet': controlnet,
         'scheduler': scheduler,
         'vae': vae,
         'text_encoder': text_encoder,
         'tokenizer': tokenizer,
         'safety_checker': None,
         'feature_extractor': None,
     }
     return components


  def   __lowercase		(		self      :     str							,_a      :     Optional[Any]							,_a      :     Union[str, Any]=0       ):


     '''simple docstring'''






     if str(_a       ).startswith('mps'       ):
        _a				:	Optional[Any]			=   torch.manual_seed(_a       )
     else:
        _a				:	Dict			=   torch.Generator(device=_a       ).manual_seed(_a       )

     _a				:	str			=   2
     _a				:	Union[str, Any]			=   randn_tensor(
         (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)							,generator=_a							,device=torch.device(_a       )							,)
     _a				:	Optional[Any]			=   floats_tensor(control_image.shape							,rng=random.Random(_a       )       ).to(_a       )
     _a				:	Any			=   image.cpu().permute(0							,2							,3							,1       )[0]
     _a				:	Dict			=   Image.fromarray(np.uinta(_a       )       ).convert('RGB'       ).resize((64, 64)       )
     _a				:	List[str]			=   {
         'prompt': 'A painting of a squirrel eating a burger',
         'generator': generator,
         'num_inference_steps': 2,
         'guidance_scale': 6.0,
         'output_type': 'numpy',
         'image': image,
         'control_image': control_image,
     }

     return inputs


  def   __lowercase		(		self      :     str       ):


     '''simple docstring'''






     return self._test_attention_slicing_forward_pass(expected_max_diff=2E-3       )


  @unittest.skipIf(
      torch_device != 'cuda' or not is_xformers_available()							,reason='XFormers attention is only available with CUDA and `xformers` installed'							,)
  def   __lowercase		(		self      :     Tuple       ):


     '''simple docstring'''






     self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3       )


  def   __lowercase		(		self      :     List[Any]       ):


     '''simple docstring'''






     self._test_inference_batch_single_identical(expected_max_diff=2E-3       )







class UpperCAmelCase__	( lowercase__    ,				lowercase__    ,				unittest.TestCase	):





  """simple docstring"""

  __UpperCAmelCase				:				Optional[int]   =					StableDiffusionControlNetImgaImgPipeline
  __UpperCAmelCase				:				Union[str, Any]   =					TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
  __UpperCAmelCase				:				Dict   =					TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
  __UpperCAmelCase				:				List[Any]   =					frozenset([]	)  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess


  def   __lowercase		(		self      :     Tuple       ):


     '''simple docstring'''






     torch.manual_seed(0       )
     _a				:	Union[str, Any]			=   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							,)
     torch.manual_seed(0       )

     def init_weights(_a      :     Any       ):
        if isinstance(_a							,torch.nn.Convad       ):
           torch.nn.init.normal(m.weight       )
           m.bias.data.fill_(1.0       )

     _a				:	List[str]			=   ControlNetModel(
         block_out_channels=(32, 64)							,layers_per_block=2							,in_channels=4							,down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')							,cross_attention_dim=32							,conditioning_embedding_out_channels=(16, 32)							,)
     controlneta.controlnet_down_blocks.apply(_a       )

     torch.manual_seed(0       )
     _a				:	Optional[Any]			=   ControlNetModel(
         block_out_channels=(32, 64)							,layers_per_block=2							,in_channels=4							,down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')							,cross_attention_dim=32							,conditioning_embedding_out_channels=(16, 32)							,)
     controlneta.controlnet_down_blocks.apply(_a       )

     torch.manual_seed(0       )
     _a				:	int			=   DDIMScheduler(
         beta_start=0.0_0085							,beta_end=0.012							,beta_schedule='scaled_linear'							,clip_sample=_a							,set_alpha_to_one=_a							,)
     torch.manual_seed(0       )
     _a				:	Tuple			=   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							,)
     torch.manual_seed(0       )
     _a				:	Dict			=   CLIPTextConfig(
         bos_token_id=0							,eos_token_id=2							,hidden_size=32							,intermediate_size=37							,layer_norm_eps=1E-05							,num_attention_heads=4							,num_hidden_layers=5							,pad_token_id=1							,vocab_size=1000							,)
     _a				:	int			=   CLIPTextModel(_a       )
     _a				:	str			=   CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )

     _a				:	Optional[int]			=   MultiControlNetModel([controlneta, controlneta]       )

     _a				:	List[Any]			=   {
         'unet': unet,
         'controlnet': controlnet,
         'scheduler': scheduler,
         'vae': vae,
         'text_encoder': text_encoder,
         'tokenizer': tokenizer,
         'safety_checker': None,
         'feature_extractor': None,
     }
     return components


  def   __lowercase		(		self      :     Any							,_a      :     Optional[Any]							,_a      :     str=0       ):


     '''simple docstring'''






     if str(_a       ).startswith('mps'       ):
        _a				:	Dict			=   torch.manual_seed(_a       )
     else:
        _a				:	Optional[Any]			=   torch.Generator(device=_a       ).manual_seed(_a       )

     _a				:	Dict			=   2

     _a				:	List[str]			=   [
         randn_tensor(
             (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)							,generator=_a							,device=torch.device(_a       )							,),
         randn_tensor(
             (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)							,generator=_a							,device=torch.device(_a       )							,),
     ]

     _a				:	List[Any]			=   floats_tensor(control_image[0].shape							,rng=random.Random(_a       )       ).to(_a       )
     _a				:	Dict			=   image.cpu().permute(0							,2							,3							,1       )[0]
     _a				:	Any			=   Image.fromarray(np.uinta(_a       )       ).convert('RGB'       ).resize((64, 64)       )
     _a				:	List[Any]			=   {
         'prompt': 'A painting of a squirrel eating a burger',
         'generator': generator,
         'num_inference_steps': 2,
         'guidance_scale': 6.0,
         'output_type': 'numpy',
         'image': image,
         'control_image': control_image,
     }

     return inputs


  def   __lowercase		(		self      :     List[Any]       ):


     '''simple docstring'''






     _a				:	Any			=   self.get_dummy_components()
     _a				:	Tuple			=   self.pipeline_class(**_a       )
     pipe.to(_a       )

     _a				:	str			=   10.0
     _a				:	List[str]			=   4

     _a				:	int			=   self.get_dummy_inputs(_a       )
     _a				:	Tuple			=   steps
     _a				:	Optional[int]			=   scale
     _a				:	Dict			=   pipe(**_a       )[0]

     _a				:	Tuple			=   self.get_dummy_inputs(_a       )
     _a				:	Tuple			=   steps
     _a				:	Dict			=   scale
     _a				:	List[Any]			=   pipe(**_a							,control_guidance_start=0.1							,control_guidance_end=0.2       )[0]

     _a				:	Optional[Any]			=   self.get_dummy_inputs(_a       )
     _a				:	int			=   steps
     _a				:	List[str]			=   scale
     _a				:	Any			=   pipe(**_a							,control_guidance_start=[0.1, 0.3]							,control_guidance_end=[0.2, 0.7]       )[0]

     _a				:	Union[str, Any]			=   self.get_dummy_inputs(_a       )
     _a				:	List[str]			=   steps
     _a				:	List[Any]			=   scale
     _a				:	Union[str, Any]			=   pipe(**_a							,control_guidance_start=0.4							,control_guidance_end=[0.5, 0.8]       )[0]

     # make sure that all outputs are different
     assert np.sum(np.abs(output_a - output_a       )       ) > 1E-3
     assert np.sum(np.abs(output_a - output_a       )       ) > 1E-3
     assert np.sum(np.abs(output_a - output_a       )       ) > 1E-3


  def   __lowercase		(		self      :     List[str]       ):


     '''simple docstring'''






     return self._test_attention_slicing_forward_pass(expected_max_diff=2E-3       )


  @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       ):


     '''simple docstring'''






     self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3       )


  def   __lowercase		(		self      :     List[Any]       ):


     '''simple docstring'''






     self._test_inference_batch_single_identical(expected_max_diff=2E-3       )


  def   __lowercase		(		self      :     Any       ):


     '''simple docstring'''






     _a				:	str			=   self.get_dummy_components()
     _a				:	List[str]			=   self.pipeline_class(**_a       )
     pipe.to(_a       )
     pipe.set_progress_bar_config(disable=_a       )
     with tempfile.TemporaryDirectory() as tmpdir:
        try:
           # save_pretrained is not implemented for Multi-ControlNet
           pipe.save_pretrained(_a       )
        except NotImplementedError:
           pass







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





  """simple docstring"""

  def   __lowercase		(		self      :     Tuple       ):


     '''simple docstring'''






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


  def   __lowercase		(		self      :     Any       ):


     '''simple docstring'''






     _a				:	Dict			=   ControlNetModel.from_pretrained('lllyasviel/sd-controlnet-canny'       )

     _a				:	Any			=   StableDiffusionControlNetImgaImgPipeline.from_pretrained(
         'runwayml/stable-diffusion-v1-5'							,safety_checker=_a							,controlnet=_a       )
     pipe.enable_model_cpu_offload()
     pipe.set_progress_bar_config(disable=_a       )

     _a				:	Tuple			=   torch.Generator(device='cpu'       ).manual_seed(0       )
     _a				:	str			=   'evil space-punk bird'
     _a				:	Optional[int]			=   load_image(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png'       ).resize((512, 512)       )
     _a				:	Union[str, Any]			=   load_image(
         'https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png'       ).resize((512, 512)       )

     _a				:	Any			=   pipe(
         _a							,_a							,control_image=_a							,generator=_a							,output_type='np'							,num_inference_steps=50							,strength=0.6							,)

     _a				:	int			=   output.images[0]

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

     _a				:	Dict			=   load_numpy(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/img2img.npy'       )

     assert np.abs(expected_image - image       ).max() < 9E-2





'''simple docstring'''







import argparse
import pickle

import numpy as np
import torch
from torch import nn

from transformers import ReformerConfig, ReformerModelWithLMHead
from transformers.utils import logging


logging.set_verbosity_info()





def 					UpperCAmelCase_		(__a      :    Optional[Any] ,	__a      :    str ,	__a      :    Optional[Any]=None ):






   """simple docstring"""

   assert torch_layer.weight.shape == weight.shape, f"""{torch_layer} layer.weight does not match"""
   _a				:	str			=   nn.Parameter(__a )
   if bias is not None:
      assert torch_layer.bias.shape == bias.shape, f"""{torch_layer} layer.bias does not match"""
      _a				:	Any			=   nn.Parameter(__a )





def 					UpperCAmelCase_		(__a      :    int ,	__a      :    Optional[Any] ,	__a      :    int ):






   """simple docstring"""

   _a				:	Tuple			=   np.asarray(weights[0] )
   _a				:	Union[str, Any]			=   np.asarray(weights[1] )
   _a				:	Dict			=   np.asarray(weights[2] )

   set_param(
       torch_layer.self_attention.query_key ,	torch.tensor(__a ).transpose(1 ,	2 ).contiguous().view(-1 ,	__a ) ,	)
   set_param(
       torch_layer.self_attention.value ,	torch.tensor(__a ).transpose(1 ,	2 ).contiguous().view(-1 ,	__a ) ,	)
   set_param(
       torch_layer.output.dense ,	torch.tensor(__a ).view(-1 ,	__a ).contiguous().transpose(0 ,	1 ) ,	)





def 					UpperCAmelCase_		(__a      :    Optional[Any] ,	__a      :    Optional[int] ,	__a      :    List[str] ):






   """simple docstring"""

   _a				:	Dict			=   np.asarray(weights[0] )
   _a				:	Union[str, Any]			=   np.asarray(weights[1] )
   _a				:	str			=   np.asarray(weights[2] )
   _a				:	int			=   np.asarray(weights[3] )

   set_param(
       torch_layer.self_attention.query ,	torch.tensor(__a ).transpose(1 ,	2 ).contiguous().view(-1 ,	__a ) ,	)
   set_param(
       torch_layer.self_attention.key ,	torch.tensor(__a ).transpose(1 ,	2 ).contiguous().view(-1 ,	__a ) ,	)
   set_param(
       torch_layer.self_attention.value ,	torch.tensor(__a ).transpose(1 ,	2 ).contiguous().view(-1 ,	__a ) ,	)
   set_param(
       torch_layer.output.dense ,	torch.tensor(__a ).view(-1 ,	__a ).contiguous().transpose(0 ,	1 ) ,	)





def 					UpperCAmelCase_		(__a      :    Any ,	__a      :    Any ,	__a      :    Optional[Any] ):






   """simple docstring"""

   _a				:	List[str]			=   weights[0][0][0]
   _a				:	List[Any]			=   np.asarray(layer_norm_a[0] )
   _a				:	List[str]			=   np.asarray(layer_norm_a[1] )
   set_param(
       torch_block.attention.layer_norm ,	torch.tensor(__a ) ,	torch.tensor(__a ) ,	)

   # lsh weights + output
   _a				:	List[str]			=   weights[0][1]
   if len(__a ) < 4:
      set_layer_weights_in_torch_lsh(__a ,	torch_block.attention ,	__a )
   else:
      set_layer_weights_in_torch_local(__a ,	torch_block.attention ,	__a )

   # intermediate weighs
   _a				:	Optional[Any]			=   weights[2][0][1][2]

   # Chunked Feed Forward
   if len(__a ) == 4:
      _a				:	Union[str, Any]			=   intermediate_weights[2]

   # layernorm 2
   _a				:	Any			=   np.asarray(intermediate_weights[0][0] )
   _a				:	List[Any]			=   np.asarray(intermediate_weights[0][1] )
   set_param(
       torch_block.feed_forward.layer_norm ,	torch.tensor(__a ) ,	torch.tensor(__a ) ,	)

   # intermediate dense
   _a				:	Any			=   np.asarray(intermediate_weights[1][0] )
   _a				:	Any			=   np.asarray(intermediate_weights[1][1] )
   set_param(
       torch_block.feed_forward.dense.dense ,	torch.tensor(__a ).transpose(0 ,	1 ).contiguous() ,	torch.tensor(__a ) ,	)

   # intermediate out
   _a				:	Optional[int]			=   np.asarray(intermediate_weights[4][0] )
   _a				:	int			=   np.asarray(intermediate_weights[4][1] )
   set_param(
       torch_block.feed_forward.output.dense ,	torch.tensor(__a ).transpose(0 ,	1 ).contiguous() ,	torch.tensor(__a ) ,	)





def 					UpperCAmelCase_		(__a      :    Dict ,	__a      :    Dict ,	__a      :    List[Any] ):






   """simple docstring"""

   _a				:	Optional[int]			=   torch_model.reformer

   # word embeds
   _a				:	Tuple			=   np.asarray(weights[1] )
   set_param(
       torch_model_reformer.embeddings.word_embeddings ,	torch.tensor(__a ) ,	)

   if isinstance(weights[3] ,	__a ):
      _a				:	Any			=   torch_model_reformer.embeddings.position_embeddings
      for emb_idx in range(len(position_embeddings.weights ) ):
         _a				:	List[Any]			=   np.asarray(weights[3][emb_idx][0] )
         assert (
             position_embeddings.weights[emb_idx].shape == emb_weights.shape
         ), f"""{position_embeddings[emb_idx]} emb does not match"""
         _a				:	Any			=   nn.Parameter(torch.tensor(__a ) )

   _a				:	List[str]			=   weights[5]
   assert len(torch_model_reformer.encoder.layers ) * 4 == len(
       __a ), "HF and trax model do not have the same number of layers"
   for layer_idx, layer in enumerate(torch_model_reformer.encoder.layers ):
      _a				:	Tuple			=   trax_layer_weights[4 * layer_idx : 4 * (layer_idx + 1)]
      set_block_weights_in_torch(__a ,	__a ,	__a )

   # output layer norm
   _a				:	Optional[Any]			=   np.asarray(weights[7][0] )
   _a				:	int			=   np.asarray(weights[7][1] )
   set_param(
       torch_model_reformer.encoder.layer_norm ,	torch.tensor(__a ) ,	torch.tensor(__a ) ,	)

   # output embeddings
   _a				:	List[str]			=   np.asarray(weights[9][0] )
   _a				:	int			=   np.asarray(weights[9][1] )
   set_param(
       torch_model.lm_head.decoder ,	torch.tensor(__a ).transpose(0 ,	1 ).contiguous() ,	torch.tensor(__a ) ,	)





def 					UpperCAmelCase_		(__a      :    Tuple ,	__a      :    Optional[Any] ,	__a      :    Dict ):






   """simple docstring"""

   _a				:	List[Any]			=   ReformerConfig.from_json_file(__a )
   print(f"""Building PyTorch model from configuration: {config}""" )
   _a				:	int			=   ReformerModelWithLMHead(__a )

   with open(__a ,	'rb' ) as f:
      _a				:	Optional[Any]			=   pickle.load(__a )['weights']

   set_model_weights_in_torch(__a ,	__a ,	config.hidden_size )

   # Save pytorch-model
   print(f"""Save PyTorch model to {pytorch_dump_path}""" )
   torch.save(model.state_dict() ,	__a )


if __name__ == "__main__":
  __lowerCAmelCase         =    argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--trax_model_pkl_path""", default=None, type=str, required=True, help="""Path to the TensorFlow checkpoint path."""
  )
  parser.add_argument(
      """--config_file""",
      default=None,
      type=str,
      required=True,
      help=(
          """The config json file corresponding to the pre-trained Reformer model. \n"""
          """This specifies the model architecture."""
      ),
  )
  parser.add_argument(
      """--pytorch_dump_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
  )
  __lowerCAmelCase         =    parser.parse_args()
  convert_trax_checkpoint_to_pytorch(args.trax_model_pkl_path, args.config_file, args.pytorch_dump_path)



