



import argparse
import json

import torch

from diffusers import AutoencoderKL, DDPMPipeline, DDPMScheduler, UNetaDModel, VQModel
def   UpperCAmelCase							(		A__				,  A__=1     )							->					Dict:
	if n_shave_prefix_segments >= 0:
		return ".".join(path.split("""."""     )[n_shave_prefix_segments:]     )
	else:
		return ".".join(path.split("""."""     )[:n_shave_prefix_segments]     )
def   UpperCAmelCase							(		A__				,  A__=0     )							->					Dict:
	_snake_case					:							Optional[Any]					     =   []
	for old_item in old_list:
		_snake_case					:							List[str]					     =   old_item
		_snake_case					:							List[Any]					     =   new_item.replace("""block."""				,  """resnets."""     )
		_snake_case					:							Optional[int]					     =   new_item.replace("""conv_shorcut"""				,  """conv1"""     )
		_snake_case					:							Optional[int]					     =   new_item.replace("""in_shortcut"""				,  """conv_shortcut"""     )
		_snake_case					:							int					     =   new_item.replace("""temb_proj"""				,  """time_emb_proj"""     )

		_snake_case					:							Any					     =   shave_segments(A__				,  n_shave_prefix_segments=A__     )

		mapping.append({"""old""": old_item, """new""": new_item}     )

	return mapping
def   UpperCAmelCase							(		A__				,  A__=0				,  A__=False     )							->					int:
	_snake_case					:							List[str]					     =   []
	for old_item in old_list:
		_snake_case					:							Tuple					     =   old_item

		# In `model.mid`, the layer is called `attn`.
		if not in_mid:
			_snake_case					:							List[str]					     =   new_item.replace("""attn"""				,  """attentions"""     )
		_snake_case					:							List[str]					     =   new_item.replace(""".k."""				,  """.key."""     )
		_snake_case					:							Optional[Any]					     =   new_item.replace(""".v."""				,  """.value."""     )
		_snake_case					:							Optional[Any]					     =   new_item.replace(""".q."""				,  """.query."""     )

		_snake_case					:							Tuple					     =   new_item.replace("""proj_out"""				,  """proj_attn"""     )
		_snake_case					:							Optional[int]					     =   new_item.replace("""norm"""				,  """group_norm"""     )

		_snake_case					:							List[str]					     =   shave_segments(A__				,  n_shave_prefix_segments=A__     )
		mapping.append({"""old""": old_item, """new""": new_item}     )

	return mapping
def   UpperCAmelCase							(		A__				,  A__				,  A__				,  A__=None				,  A__=None				,  A__=None     )							->					List[str]:
	assert isinstance(A__				,  A__     ), "Paths should be a list of dicts containing 'old' and 'new' keys."

	if attention_paths_to_split is not None:
		if config is None:
			raise ValueError("""Please specify the config if setting 'attention_paths_to_split' to 'True'."""     )

		for path, path_map in attention_paths_to_split.items():
			_snake_case					:							Dict					     =   old_checkpoint[path]
			_snake_case					:							Optional[int]					     =   old_tensor.shape[0] // 3

			_snake_case					:							int					     =   (-1, channels) if len(old_tensor.shape     ) == 3 else (-1)

			_snake_case					:							Optional[Any]					     =   old_tensor.shape[0] // config.get("""num_head_channels"""				,  1     ) // 3

			_snake_case					:							Dict					     =   old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]     )
			_snake_case     , _snake_case     , _snake_case					:							Optional[Any]					     =   old_tensor.split(channels // num_heads				,  dim=1     )

			_snake_case					:							Tuple					     =   query.reshape(A__     ).squeeze()
			_snake_case					:							Any					     =   key.reshape(A__     ).squeeze()
			_snake_case					:							int					     =   value.reshape(A__     ).squeeze()

	for path in paths:
		_snake_case					:							Optional[int]					     =   path["""new"""]

		if attention_paths_to_split is not None and new_path in attention_paths_to_split:
			continue

		_snake_case					:							List[Any]					     =   new_path.replace("""down."""				,  """down_blocks."""     )
		_snake_case					:							int					     =   new_path.replace("""up."""				,  """up_blocks."""     )

		if additional_replacements is not None:
			for replacement in additional_replacements:
				_snake_case					:							List[Any]					     =   new_path.replace(replacement["""old"""]				,  replacement["""new"""]     )

		if "attentions" in new_path:
			_snake_case					:							str					     =   old_checkpoint[path["""old"""]].squeeze()
		else:
			_snake_case					:							Union[str, Any]					     =   old_checkpoint[path["""old"""]]
def   UpperCAmelCase							(		A__				,  A__     )							->					Any:
	_snake_case					:							Optional[int]					     =   {}

	_snake_case					:							Tuple					     =   checkpoint["""temb.dense.0.weight"""]
	_snake_case					:							Tuple					     =   checkpoint["""temb.dense.0.bias"""]
	_snake_case					:							Union[str, Any]					     =   checkpoint["""temb.dense.1.weight"""]
	_snake_case					:							str					     =   checkpoint["""temb.dense.1.bias"""]

	_snake_case					:							List[Any]					     =   checkpoint["""norm_out.weight"""]
	_snake_case					:							List[str]					     =   checkpoint["""norm_out.bias"""]

	_snake_case					:							List[Any]					     =   checkpoint["""conv_in.weight"""]
	_snake_case					:							Optional[int]					     =   checkpoint["""conv_in.bias"""]
	_snake_case					:							Union[str, Any]					     =   checkpoint["""conv_out.weight"""]
	_snake_case					:							Optional[int]					     =   checkpoint["""conv_out.bias"""]

	_snake_case					:							Optional[int]					     =   len({""".""".join(layer.split("""."""     )[:2]     ) for layer in checkpoint if """down""" in layer}     )
	_snake_case					:							List[Any]					     =   {
	    layer_id: [key for key in checkpoint if f'''down.{layer_id}''' in key] for layer_id in range(A__     )
	}

	_snake_case					:							str					     =   len({""".""".join(layer.split("""."""     )[:2]     ) for layer in checkpoint if """up""" in layer}     )
	_snake_case					:							Dict					     =   {layer_id: [key for key in checkpoint if f'''up.{layer_id}''' in key] for layer_id in range(A__     )}

	for i in range(A__     ):
		_snake_case					:							List[str]					     =   (i - 1) // (config["""layers_per_block"""] + 1)

		if any("""downsample""" in layer for layer in down_blocks[i]     ):
			_snake_case					:							List[Any]					     =   checkpoint[
			    f'''down.{i}.downsample.op.weight'''
			]
			_snake_case					:							Union[str, Any]					     =   checkpoint[f'''down.{i}.downsample.op.bias''']
		#            new_checkpoint[f'down_blocks.{i}.downsamplers.0.op.weight'] = checkpoint[f'down.{i}.downsample.conv.weight']
		#            new_checkpoint[f'down_blocks.{i}.downsamplers.0.op.bias'] = checkpoint[f'down.{i}.downsample.conv.bias']

		if any("""block""" in layer for layer in down_blocks[i]     ):
			_snake_case					:							int					     =   len(
			    {""".""".join(shave_segments(A__				,  2     ).split("""."""     )[:2]     ) for layer in down_blocks[i] if """block""" in layer}     )
			_snake_case					:							Any					     =   {
			    layer_id: [key for key in down_blocks[i] if f'''block.{layer_id}''' in key]
			    for layer_id in range(A__     )
			}

			if num_blocks > 0:
				for j in range(config["""layers_per_block"""]     ):
					_snake_case					:							List[str]					     =   renew_resnet_paths(blocks[j]     )
					assign_to_checkpoint(A__				,  A__				,  A__     )

		if any("""attn""" in layer for layer in down_blocks[i]     ):
			_snake_case					:							int					     =   len(
			    {""".""".join(shave_segments(A__				,  2     ).split("""."""     )[:2]     ) for layer in down_blocks[i] if """attn""" in layer}     )
			_snake_case					:							Optional[Any]					     =   {
			    layer_id: [key for key in down_blocks[i] if f'''attn.{layer_id}''' in key]
			    for layer_id in range(A__     )
			}

			if num_attn > 0:
				for j in range(config["""layers_per_block"""]     ):
					_snake_case					:							Dict					     =   renew_attention_paths(attns[j]     )
					assign_to_checkpoint(A__				,  A__				,  A__				,  config=A__     )

	_snake_case					:							Optional[Any]					     =   [key for key in checkpoint if """mid.block_1""" in key]
	_snake_case					:							List[str]					     =   [key for key in checkpoint if """mid.block_2""" in key]
	_snake_case					:							Dict					     =   [key for key in checkpoint if """mid.attn_1""" in key]

	# Mid new 2
	_snake_case					:							List[Any]					     =   renew_resnet_paths(A__     )
	assign_to_checkpoint(
	    A__				,  A__				,  A__				,  additional_replacements=[{"""old""": """mid.""", """new""": """mid_new_2."""}, {"""old""": """block_1""", """new""": """resnets.0"""}]				,  )

	_snake_case					:							Union[str, Any]					     =   renew_resnet_paths(A__     )
	assign_to_checkpoint(
	    A__				,  A__				,  A__				,  additional_replacements=[{"""old""": """mid.""", """new""": """mid_new_2."""}, {"""old""": """block_2""", """new""": """resnets.1"""}]				,  )

	_snake_case					:							List[Any]					     =   renew_attention_paths(A__				,  in_mid=A__     )
	assign_to_checkpoint(
	    A__				,  A__				,  A__				,  additional_replacements=[{"""old""": """mid.""", """new""": """mid_new_2."""}, {"""old""": """attn_1""", """new""": """attentions.0"""}]				,  )

	for i in range(A__     ):
		_snake_case					:							List[str]					     =   num_up_blocks - 1 - i

		if any("""upsample""" in layer for layer in up_blocks[i]     ):
			_snake_case					:							List[Any]					     =   checkpoint[
			    f'''up.{i}.upsample.conv.weight'''
			]
			_snake_case					:							int					     =   checkpoint[f'''up.{i}.upsample.conv.bias''']

		if any("""block""" in layer for layer in up_blocks[i]     ):
			_snake_case					:							Dict					     =   len(
			    {""".""".join(shave_segments(A__				,  2     ).split("""."""     )[:2]     ) for layer in up_blocks[i] if """block""" in layer}     )
			_snake_case					:							Tuple					     =   {
			    layer_id: [key for key in up_blocks[i] if f'''block.{layer_id}''' in key] for layer_id in range(A__     )
			}

			if num_blocks > 0:
				for j in range(config["""layers_per_block"""] + 1     ):
					_snake_case					:							str					     =   {"""old""": f'''up_blocks.{i}''', """new""": f'''up_blocks.{block_id}'''}
					_snake_case					:							Union[str, Any]					     =   renew_resnet_paths(blocks[j]     )
					assign_to_checkpoint(A__				,  A__				,  A__				,  additional_replacements=[replace_indices]     )

		if any("""attn""" in layer for layer in up_blocks[i]     ):
			_snake_case					:							Dict					     =   len(
			    {""".""".join(shave_segments(A__				,  2     ).split("""."""     )[:2]     ) for layer in up_blocks[i] if """attn""" in layer}     )
			_snake_case					:							Dict					     =   {
			    layer_id: [key for key in up_blocks[i] if f'''attn.{layer_id}''' in key] for layer_id in range(A__     )
			}

			if num_attn > 0:
				for j in range(config["""layers_per_block"""] + 1     ):
					_snake_case					:							Tuple					     =   {"""old""": f'''up_blocks.{i}''', """new""": f'''up_blocks.{block_id}'''}
					_snake_case					:							int					     =   renew_attention_paths(attns[j]     )
					assign_to_checkpoint(A__				,  A__				,  A__				,  additional_replacements=[replace_indices]     )

	_snake_case					:							Optional[int]					     =   {k.replace("""mid_new_2"""				,  """mid_block"""     ): v for k, v in new_checkpoint.items()}
	return new_checkpoint
def   UpperCAmelCase							(		A__				,  A__     )							->					Optional[Any]:
	_snake_case					:							Tuple					     =   {}

	_snake_case					:							int					     =   checkpoint["""encoder.norm_out.weight"""]
	_snake_case					:							int					     =   checkpoint["""encoder.norm_out.bias"""]

	_snake_case					:							Optional[Any]					     =   checkpoint["""encoder.conv_in.weight"""]
	_snake_case					:							Optional[Any]					     =   checkpoint["""encoder.conv_in.bias"""]
	_snake_case					:							List[str]					     =   checkpoint["""encoder.conv_out.weight"""]
	_snake_case					:							int					     =   checkpoint["""encoder.conv_out.bias"""]

	_snake_case					:							int					     =   checkpoint["""decoder.norm_out.weight"""]
	_snake_case					:							int					     =   checkpoint["""decoder.norm_out.bias"""]

	_snake_case					:							str					     =   checkpoint["""decoder.conv_in.weight"""]
	_snake_case					:							Any					     =   checkpoint["""decoder.conv_in.bias"""]
	_snake_case					:							str					     =   checkpoint["""decoder.conv_out.weight"""]
	_snake_case					:							Dict					     =   checkpoint["""decoder.conv_out.bias"""]

	_snake_case					:							Union[str, Any]					     =   len({""".""".join(layer.split("""."""     )[:3]     ) for layer in checkpoint if """down""" in layer}     )
	_snake_case					:							Tuple					     =   {
	    layer_id: [key for key in checkpoint if f'''down.{layer_id}''' in key] for layer_id in range(A__     )
	}

	_snake_case					:							str					     =   len({""".""".join(layer.split("""."""     )[:3]     ) for layer in checkpoint if """up""" in layer}     )
	_snake_case					:							Any					     =   {layer_id: [key for key in checkpoint if f'''up.{layer_id}''' in key] for layer_id in range(A__     )}

	for i in range(A__     ):
		_snake_case					:							Any					     =   (i - 1) // (config["""layers_per_block"""] + 1)

		if any("""downsample""" in layer for layer in down_blocks[i]     ):
			_snake_case					:							str					     =   checkpoint[
			    f'''encoder.down.{i}.downsample.conv.weight'''
			]
			_snake_case					:							Optional[int]					     =   checkpoint[
			    f'''encoder.down.{i}.downsample.conv.bias'''
			]

		if any("""block""" in layer for layer in down_blocks[i]     ):
			_snake_case					:							Optional[Any]					     =   len(
			    {""".""".join(shave_segments(A__				,  3     ).split("""."""     )[:3]     ) for layer in down_blocks[i] if """block""" in layer}     )
			_snake_case					:							Optional[Any]					     =   {
			    layer_id: [key for key in down_blocks[i] if f'''block.{layer_id}''' in key]
			    for layer_id in range(A__     )
			}

			if num_blocks > 0:
				for j in range(config["""layers_per_block"""]     ):
					_snake_case					:							Tuple					     =   renew_resnet_paths(blocks[j]     )
					assign_to_checkpoint(A__				,  A__				,  A__     )

		if any("""attn""" in layer for layer in down_blocks[i]     ):
			_snake_case					:							str					     =   len(
			    {""".""".join(shave_segments(A__				,  3     ).split("""."""     )[:3]     ) for layer in down_blocks[i] if """attn""" in layer}     )
			_snake_case					:							List[str]					     =   {
			    layer_id: [key for key in down_blocks[i] if f'''attn.{layer_id}''' in key]
			    for layer_id in range(A__     )
			}

			if num_attn > 0:
				for j in range(config["""layers_per_block"""]     ):
					_snake_case					:							Union[str, Any]					     =   renew_attention_paths(attns[j]     )
					assign_to_checkpoint(A__				,  A__				,  A__				,  config=A__     )

	_snake_case					:							Optional[Any]					     =   [key for key in checkpoint if """mid.block_1""" in key]
	_snake_case					:							Dict					     =   [key for key in checkpoint if """mid.block_2""" in key]
	_snake_case					:							str					     =   [key for key in checkpoint if """mid.attn_1""" in key]

	# Mid new 2
	_snake_case					:							Union[str, Any]					     =   renew_resnet_paths(A__     )
	assign_to_checkpoint(
	    A__				,  A__				,  A__				,  additional_replacements=[{"""old""": """mid.""", """new""": """mid_new_2."""}, {"""old""": """block_1""", """new""": """resnets.0"""}]				,  )

	_snake_case					:							Any					     =   renew_resnet_paths(A__     )
	assign_to_checkpoint(
	    A__				,  A__				,  A__				,  additional_replacements=[{"""old""": """mid.""", """new""": """mid_new_2."""}, {"""old""": """block_2""", """new""": """resnets.1"""}]				,  )

	_snake_case					:							Optional[Any]					     =   renew_attention_paths(A__				,  in_mid=A__     )
	assign_to_checkpoint(
	    A__				,  A__				,  A__				,  additional_replacements=[{"""old""": """mid.""", """new""": """mid_new_2."""}, {"""old""": """attn_1""", """new""": """attentions.0"""}]				,  )

	for i in range(A__     ):
		_snake_case					:							List[Any]					     =   num_up_blocks - 1 - i

		if any("""upsample""" in layer for layer in up_blocks[i]     ):
			_snake_case					:							Optional[int]					     =   checkpoint[
			    f'''decoder.up.{i}.upsample.conv.weight'''
			]
			_snake_case					:							List[str]					     =   checkpoint[
			    f'''decoder.up.{i}.upsample.conv.bias'''
			]

		if any("""block""" in layer for layer in up_blocks[i]     ):
			_snake_case					:							int					     =   len(
			    {""".""".join(shave_segments(A__				,  3     ).split("""."""     )[:3]     ) for layer in up_blocks[i] if """block""" in layer}     )
			_snake_case					:							List[Any]					     =   {
			    layer_id: [key for key in up_blocks[i] if f'''block.{layer_id}''' in key] for layer_id in range(A__     )
			}

			if num_blocks > 0:
				for j in range(config["""layers_per_block"""] + 1     ):
					_snake_case					:							List[Any]					     =   {"""old""": f'''up_blocks.{i}''', """new""": f'''up_blocks.{block_id}'''}
					_snake_case					:							Optional[int]					     =   renew_resnet_paths(blocks[j]     )
					assign_to_checkpoint(A__				,  A__				,  A__				,  additional_replacements=[replace_indices]     )

		if any("""attn""" in layer for layer in up_blocks[i]     ):
			_snake_case					:							Union[str, Any]					     =   len(
			    {""".""".join(shave_segments(A__				,  3     ).split("""."""     )[:3]     ) for layer in up_blocks[i] if """attn""" in layer}     )
			_snake_case					:							Dict					     =   {
			    layer_id: [key for key in up_blocks[i] if f'''attn.{layer_id}''' in key] for layer_id in range(A__     )
			}

			if num_attn > 0:
				for j in range(config["""layers_per_block"""] + 1     ):
					_snake_case					:							str					     =   {"""old""": f'''up_blocks.{i}''', """new""": f'''up_blocks.{block_id}'''}
					_snake_case					:							List[Any]					     =   renew_attention_paths(attns[j]     )
					assign_to_checkpoint(A__				,  A__				,  A__				,  additional_replacements=[replace_indices]     )

	_snake_case					:							Dict					     =   {k.replace("""mid_new_2"""				,  """mid_block"""     ): v for k, v in new_checkpoint.items()}
	_snake_case					:							Any					     =   checkpoint["""quant_conv.weight"""]
	_snake_case					:							Dict					     =   checkpoint["""quant_conv.bias"""]
	if "quantize.embedding.weight" in checkpoint:
		_snake_case					:							Tuple					     =   checkpoint["""quantize.embedding.weight"""]
	_snake_case					:							str					     =   checkpoint["""post_quant_conv.weight"""]
	_snake_case					:							str					     =   checkpoint["""post_quant_conv.bias"""]

	return new_checkpoint


if __name__ == "__main__":
						UpperCAmelCase_			      =  argparse.ArgumentParser()

						parser.add_argument(
						    '''--checkpoint_path''', default=None, type=str, required=True, help='''Path to the checkpoint to convert.'''
						)

						parser.add_argument(
						    '''--config_file''',
						    default=None,
						    type=str,
						    required=True,
						    help='''The config json file corresponding to the architecture.''',
						)

						parser.add_argument('''--dump_path''', default=None, type=str, required=True, help='''Path to the output model.''')

						UpperCAmelCase_			      =  parser.parse_args()
						UpperCAmelCase_			      =  torch.load(args.checkpoint_path)

						with open(args.config_file) as f:
												UpperCAmelCase_			      =  json.loads(f.read())

						# unet case
						UpperCAmelCase_			      =  {key.split('''.''')[0] for key in checkpoint.keys()}
						if "encoder" in key_prefix_set and "decoder" in key_prefix_set:
												UpperCAmelCase_			      =  convert_vq_autoenc_checkpoint(checkpoint, config)
						else:
												UpperCAmelCase_			      =  convert_ddpm_checkpoint(checkpoint, config)

						if "ddpm" in config:
												del config["ddpm"]

						if config["_class_name"] == "VQModel":
												UpperCAmelCase_			      =  VQModel(**config)
												model.load_state_dict(converted_checkpoint)
												model.save_pretrained(args.dump_path)
						elif config["_class_name"] == "AutoencoderKL":
												UpperCAmelCase_			      =  AutoencoderKL(**config)
												model.load_state_dict(converted_checkpoint)
												model.save_pretrained(args.dump_path)
						else:
												UpperCAmelCase_			      =  UNetaDModel(**config)
												model.load_state_dict(converted_checkpoint)

												UpperCAmelCase_			      =  DDPMScheduler.from_config('''/'''.join(args.checkpoint_path.split('''/''')[:-1]))

												UpperCAmelCase_			      =  DDPMPipeline(unet=model, scheduler=scheduler)
												pipe.save_pretrained(args.dump_path)










import tempfile
import unittest

import numpy as np
import torch

from diffusers import DiffusionPipeline
from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor





class 		__SCREAMING_SNAKE_CASE      (				unittest.TestCase      ):






				"""simple docstring"""

				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__ = False      ):






					"""simple docstring"""

					_snake_case					:							Any					     =   10

					if only_cross_attention:
						_snake_case					:							Tuple					     =   12
					else:
						# when only cross attention is not set, the cross attention dim must be the same as the query dim
						_snake_case					:							Any					     =   query_dim

					return {
					    "query_dim": query_dim,
					    "cross_attention_dim": cross_attention_dim,
					    "heads": 2,
					    "dim_head": 4,
					    "added_kv_proj_dim": 6,
					    "norm_num_groups": 1,
					    "only_cross_attention": only_cross_attention,
					    "processor": AttnAddedKVProcessor(),
					}







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					_snake_case					:							Any					     =   2

					_snake_case					:							List[str]					     =   torch.rand(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	3						,	2      )
					_snake_case					:							List[str]					     =   torch.rand(SCREAMING_SNAKE_CASE__						,	4						,	SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Optional[int]					     =   None

					return {
					    "hidden_states": hidden_states,
					    "encoder_hidden_states": encoder_hidden_states,
					    "attention_mask": attention_mask,
					}







				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					torch.manual_seed(0      )

					_snake_case					:							Optional[int]					     =   self.get_constructor_arguments(only_cross_attention=SCREAMING_SNAKE_CASE__      )
					_snake_case					:							List[str]					     =   Attention(**SCREAMING_SNAKE_CASE__      )

					self.assertTrue(attn.to_k is not None      )
					self.assertTrue(attn.to_v is not None      )

					_snake_case					:							Dict					     =   self.get_forward_arguments(
					    query_dim=constructor_args["""query_dim"""]						,	added_kv_proj_dim=constructor_args["""added_kv_proj_dim"""]      )

					_snake_case					:							Any					     =   attn(**SCREAMING_SNAKE_CASE__      )

					# only self attention

					torch.manual_seed(0      )

					_snake_case					:							Any					     =   self.get_constructor_arguments(only_cross_attention=SCREAMING_SNAKE_CASE__      )
					_snake_case					:							List[str]					     =   Attention(**SCREAMING_SNAKE_CASE__      )

					self.assertTrue(attn.to_k is None      )
					self.assertTrue(attn.to_v is None      )

					_snake_case					:							str					     =   self.get_forward_arguments(
					    query_dim=constructor_args["""query_dim"""]						,	added_kv_proj_dim=constructor_args["""added_kv_proj_dim"""]      )

					_snake_case					:							Any					     =   attn(**SCREAMING_SNAKE_CASE__      )

					self.assertTrue((only_cross_attn_out != self_and_cross_attn_out).all()      )





class 		__SCREAMING_SNAKE_CASE      (				unittest.TestCase      ):






				"""simple docstring"""

				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					_snake_case					:							str					     =   DiffusionPipeline.from_pretrained("""hf-internal-testing/tiny-stable-diffusion-pipe"""						,	safety_checker=SCREAMING_SNAKE_CASE__      )

					_snake_case					:							Tuple					     =   pipe(
					    """foo"""						,	num_inference_steps=2						,	generator=torch.Generator("""cpu"""      ).manual_seed(0      )						,	output_type="""np"""						,	).images

					# the initial conversion succeeds
					_snake_case					:							List[str]					     =   DiffusionPipeline.from_pretrained(
					    """hf-internal-testing/tiny-stable-diffusion-pipe"""						,	device_map="""sequential"""						,	safety_checker=SCREAMING_SNAKE_CASE__      )

					_snake_case					:							Union[str, Any]					     =   pipe(
					    """foo"""						,	num_inference_steps=2						,	generator=torch.Generator("""cpu"""      ).manual_seed(0      )						,	output_type="""np"""						,	).images

					with tempfile.TemporaryDirectory() as tmpdir:
						# save the converted model
						pipe.save_pretrained(SCREAMING_SNAKE_CASE__      )

						# can also load the converted weights
						_snake_case					:							str					     =   DiffusionPipeline.from_pretrained(SCREAMING_SNAKE_CASE__						,	device_map="""sequential"""						,	safety_checker=SCREAMING_SNAKE_CASE__      )

					_snake_case					:							int					     =   pipe(
					    """foo"""						,	num_inference_steps=2						,	generator=torch.Generator("""cpu"""      ).manual_seed(0      )						,	output_type="""np"""						,	).images

					self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	atol=1e-5      )      )
					self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	atol=1e-5      )      )





