



import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import SegformerImageProcessor, SwinConfig, UperNetConfig, UperNetForSemanticSegmentation






def 			lowerCAmelCase__	(					lowerCamelCase_	:	Any):





				'''simple docstring'''

				lowerCAmelCase__  :       Optional[int]				=  384
				lowerCAmelCase__  :       int				=  7
				if "tiny" in model_name:
								lowerCAmelCase__  :       Dict				=  96
								lowerCAmelCase__  :       int				=  (2, 2, 6, 2)
								lowerCAmelCase__  :       List[str]				=  (3, 6, 12, 24)
				elif "small" in model_name:
								lowerCAmelCase__  :       Optional[int]				=  96
								lowerCAmelCase__  :       Union[str, Any]				=  (2, 2, 18, 2)
								lowerCAmelCase__  :       Tuple				=  (3, 6, 12, 24)
				elif "base" in model_name:
								lowerCAmelCase__  :       Any				=  128
								lowerCAmelCase__  :       int				=  (2, 2, 18, 2)
								lowerCAmelCase__  :       str				=  (4, 8, 16, 32)
								lowerCAmelCase__  :       List[str]				=  12
								lowerCAmelCase__  :       List[Any]				=  512
				elif "large" in model_name:
								lowerCAmelCase__  :       Tuple				=  192
								lowerCAmelCase__  :       Optional[Any]				=  (2, 2, 18, 2)
								lowerCAmelCase__  :       Tuple				=  (6, 12, 24, 48)
								lowerCAmelCase__  :       Tuple				=  12
								lowerCAmelCase__  :       Optional[int]				=  768

				# set label information
				lowerCAmelCase__  :       Tuple				=  150
				lowerCAmelCase__  :       Optional[int]				=  '''huggingface/label-files'''
				lowerCAmelCase__  :       str				=  '''ade20k-id2label.json'''
				lowerCAmelCase__  :       List[Any]				=  json.load(open(hf_hub_download(lowerCamelCase_ ,lowerCamelCase_ ,repo_type='''dataset''') ,'''r'''))
				lowerCAmelCase__  :       Tuple				=  {int(lowerCamelCase_): v for k, v in idalabel.items()}
				lowerCAmelCase__  :       Tuple				=  {v: k for k, v in idalabel.items()}

				lowerCAmelCase__  :       Dict				=  SwinConfig(
				    embed_dim=lowerCamelCase_ ,depths=lowerCamelCase_ ,num_heads=lowerCamelCase_ ,window_size=lowerCamelCase_ ,out_features=['''stage1''', '''stage2''', '''stage3''', '''stage4'''] ,)
				lowerCAmelCase__  :       Dict				=  UperNetConfig(
				    backbone_config=lowerCamelCase_ ,auxiliary_in_channels=lowerCamelCase_ ,num_labels=lowerCamelCase_ ,idalabel=lowerCamelCase_ ,labelaid=lowerCamelCase_ ,)

				return config






def 			lowerCAmelCase__	(					lowerCamelCase_	:	int):





				'''simple docstring'''

				lowerCAmelCase__  :       Optional[int]				=  []

				# fmt: off
				# stem
				rename_keys.append(('''backbone.patch_embed.projection.weight''', '''backbone.embeddings.patch_embeddings.projection.weight'''))
				rename_keys.append(('''backbone.patch_embed.projection.bias''', '''backbone.embeddings.patch_embeddings.projection.bias'''))
				rename_keys.append(('''backbone.patch_embed.norm.weight''', '''backbone.embeddings.norm.weight'''))
				rename_keys.append(('''backbone.patch_embed.norm.bias''', '''backbone.embeddings.norm.bias'''))
				# stages
				for i in range(len(config.backbone_config.depths)):
								for j in range(config.backbone_config.depths[i]):
												rename_keys.append((f"""backbone.stages.{i}.blocks.{j}.norm1.weight""", f"""backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.weight"""))
												rename_keys.append((f"""backbone.stages.{i}.blocks.{j}.norm1.bias""", f"""backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.bias"""))
												rename_keys.append((f"""backbone.stages.{i}.blocks.{j}.attn.w_msa.relative_position_bias_table""", f"""backbone.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table"""))
												rename_keys.append((f"""backbone.stages.{i}.blocks.{j}.attn.w_msa.relative_position_index""", f"""backbone.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index"""))
												rename_keys.append((f"""backbone.stages.{i}.blocks.{j}.attn.w_msa.proj.weight""", f"""backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight"""))
												rename_keys.append((f"""backbone.stages.{i}.blocks.{j}.attn.w_msa.proj.bias""", f"""backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias"""))
												rename_keys.append((f"""backbone.stages.{i}.blocks.{j}.norm2.weight""", f"""backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.weight"""))
												rename_keys.append((f"""backbone.stages.{i}.blocks.{j}.norm2.bias""", f"""backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.bias"""))
												rename_keys.append((f"""backbone.stages.{i}.blocks.{j}.ffn.layers.0.0.weight""", f"""backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight"""))
												rename_keys.append((f"""backbone.stages.{i}.blocks.{j}.ffn.layers.0.0.bias""", f"""backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias"""))
												rename_keys.append((f"""backbone.stages.{i}.blocks.{j}.ffn.layers.1.weight""", f"""backbone.encoder.layers.{i}.blocks.{j}.output.dense.weight"""))
												rename_keys.append((f"""backbone.stages.{i}.blocks.{j}.ffn.layers.1.bias""", f"""backbone.encoder.layers.{i}.blocks.{j}.output.dense.bias"""))

								if i < 3:
												rename_keys.append((f"""backbone.stages.{i}.downsample.reduction.weight""", f"""backbone.encoder.layers.{i}.downsample.reduction.weight"""))
												rename_keys.append((f"""backbone.stages.{i}.downsample.norm.weight""", f"""backbone.encoder.layers.{i}.downsample.norm.weight"""))
												rename_keys.append((f"""backbone.stages.{i}.downsample.norm.bias""", f"""backbone.encoder.layers.{i}.downsample.norm.bias"""))
								rename_keys.append((f"""backbone.norm{i}.weight""", f"""backbone.hidden_states_norms.stage{i+1}.weight"""))
								rename_keys.append((f"""backbone.norm{i}.bias""", f"""backbone.hidden_states_norms.stage{i+1}.bias"""))

				# decode head
				rename_keys.extend(
				    [
				        ('''decode_head.conv_seg.weight''', '''decode_head.classifier.weight'''),
				        ('''decode_head.conv_seg.bias''', '''decode_head.classifier.bias'''),
				        ('''auxiliary_head.conv_seg.weight''', '''auxiliary_head.classifier.weight'''),
				        ('''auxiliary_head.conv_seg.bias''', '''auxiliary_head.classifier.bias'''),
				    ])
				# fmt: on

				return rename_keys






def 			lowerCAmelCase__	(					lowerCamelCase_	:	List[Any] ,lowerCamelCase_	:	Any ,lowerCamelCase_	:	Dict):





				'''simple docstring'''

				lowerCAmelCase__  :       List[str]				=  dct.pop(lowerCamelCase_)
				lowerCAmelCase__  :       Tuple				=  val






def 			lowerCAmelCase__	(					lowerCamelCase_	:	List[Any] ,lowerCamelCase_	:	str):





				'''simple docstring'''

				lowerCAmelCase__  :       Tuple				=  [int(backbone_config.embed_dim * 2**i) for i in range(len(backbone_config.depths))]
				for i in range(len(backbone_config.depths)):
								lowerCAmelCase__  :       Optional[Any]				=  num_features[i]
								for j in range(backbone_config.depths[i]):
												# fmt: off
												# read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
												lowerCAmelCase__  :       Optional[int]				=  state_dict.pop(f"""backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.weight""")
												lowerCAmelCase__  :       Any				=  state_dict.pop(f"""backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.bias""")
												# next, add query, keys and values (in that order) to the state dict
												lowerCAmelCase__  :       str				=  in_proj_weight[:dim, :]
												lowerCAmelCase__  :       List[str]				=  in_proj_bias[: dim]
												lowerCAmelCase__  :       Any				=  in_proj_weight[
												    dim : dim * 2, :
												]
												lowerCAmelCase__  :       Tuple				=  in_proj_bias[
												    dim : dim * 2
												]
												lowerCAmelCase__  :       List[Any]				=  in_proj_weight[
												    -dim :, :
												]
												lowerCAmelCase__  :       Union[str, Any]				=  in_proj_bias[-dim :]
												# fmt: on






def 			lowerCAmelCase__	(					lowerCamelCase_	:	Optional[int]):





				'''simple docstring'''

				lowerCAmelCase__  :       Any				=  x.shape
				lowerCAmelCase__  :       Any				=  x.reshape(lowerCamelCase_ ,4 ,in_channel // 4)
				lowerCAmelCase__  :       List[Any]				=  x[:, [0, 2, 1, 3], :].transpose(1 ,2).reshape(lowerCamelCase_ ,lowerCamelCase_)
				return x






def 			lowerCAmelCase__	(					lowerCamelCase_	:	List[str]):





				'''simple docstring'''

				lowerCAmelCase__  :       Union[str, Any]				=  x.shape
				lowerCAmelCase__  :       List[Any]				=  x.reshape(lowerCamelCase_ ,in_channel // 4 ,4)
				lowerCAmelCase__  :       List[str]				=  x[:, :, [0, 2, 1, 3]].transpose(1 ,2).reshape(lowerCamelCase_ ,lowerCamelCase_)

				return x






def 			lowerCAmelCase__	(					lowerCamelCase_	:	int):





				'''simple docstring'''

				lowerCAmelCase__  :       int				=  x.shape[0]
				lowerCAmelCase__  :       List[Any]				=  x.reshape(4 ,in_channel // 4)
				lowerCAmelCase__  :       Tuple				=  x[[0, 2, 1, 3], :].transpose(0 ,1).reshape(lowerCamelCase_)
				return x






def 			lowerCAmelCase__	(					lowerCamelCase_	:	Optional[Any]):





				'''simple docstring'''

				lowerCAmelCase__  :       Tuple				=  x.shape[0]
				lowerCAmelCase__  :       int				=  x.reshape(in_channel // 4 ,4)
				lowerCAmelCase__  :       Optional[int]				=  x[:, [0, 2, 1, 3]].transpose(0 ,1).reshape(lowerCamelCase_)
				return x






def 			lowerCAmelCase__	(					lowerCamelCase_	:	Any ,lowerCamelCase_	:	int ,lowerCamelCase_	:	Any):





				'''simple docstring'''

				lowerCAmelCase__  :       Optional[int]				=  {
				    '''upernet-swin-tiny''': '''https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K_20210531_112542-e380ad3e.pth''',
				    '''upernet-swin-small''': '''https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_small_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K/upernet_swin_small_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K_20210526_192015-ee2fff1c.pth''',
				    '''upernet-swin-base''': '''https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_base_patch4_window12_512x512_160k_ade20k_pretrain_384x384_22K/upernet_swin_base_patch4_window12_512x512_160k_ade20k_pretrain_384x384_22K_20210531_125459-429057bf.pth''',
				    '''upernet-swin-large''': '''https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_large_patch4_window12_512x512_pretrain_384x384_22K_160k_ade20k/upernet_swin_large_patch4_window12_512x512_pretrain_384x384_22K_160k_ade20k_20220318_091743-9ba68901.pth''',
				}
				lowerCAmelCase__  :       Any				=  model_name_to_url[model_name]
				lowerCAmelCase__  :       List[Any]				=  torch.hub.load_state_dict_from_url(lowerCamelCase_ ,map_location='''cpu''' ,file_name=lowerCamelCase_)[
				    '''state_dict'''
				]

				for name, param in state_dict.items():
								print(lowerCamelCase_ ,param.shape)

				lowerCAmelCase__  :       Union[str, Any]				=  get_upernet_config(lowerCamelCase_)
				lowerCAmelCase__  :       Tuple				=  UperNetForSemanticSegmentation(lowerCamelCase_)
				model.eval()

				# replace "bn" => "batch_norm"
				for key in state_dict.copy().keys():
								lowerCAmelCase__  :       List[Any]				=  state_dict.pop(lowerCamelCase_)
								if "bn" in key:
												lowerCAmelCase__  :       int				=  key.replace('''bn''' ,'''batch_norm''')
								lowerCAmelCase__  :       Union[str, Any]				=  val

				# rename keys
				lowerCAmelCase__  :       str				=  create_rename_keys(lowerCamelCase_)
				for src, dest in rename_keys:
								rename_key(lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_)
				read_in_q_k_v(lowerCamelCase_ ,config.backbone_config)

				# fix downsample parameters
				for key, value in state_dict.items():
								if "downsample" in key:
												if "reduction" in key:
																lowerCAmelCase__  :       int				=  reverse_correct_unfold_reduction_order(lowerCamelCase_)
												if "norm" in key:
																lowerCAmelCase__  :       Union[str, Any]				=  reverse_correct_unfold_norm_order(lowerCamelCase_)

				model.load_state_dict(lowerCamelCase_)

				# verify on image
				lowerCAmelCase__  :       Dict				=  '''https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg'''
				lowerCAmelCase__  :       Dict				=  Image.open(requests.get(lowerCamelCase_ ,stream=lowerCamelCase_).raw).convert('''RGB''')

				lowerCAmelCase__  :       Tuple				=  SegformerImageProcessor()
				lowerCAmelCase__  :       int				=  processor(lowerCamelCase_ ,return_tensors='''pt''').pixel_values

				with torch.no_grad():
								lowerCAmelCase__  :       List[Any]				=  model(lowerCamelCase_)
								lowerCAmelCase__  :       List[Any]				=  outputs.logits

				print(logits.shape)
				print('''First values of logits:''' ,logits[0, 0, :3, :3])
				# assert values
				if model_name == "upernet-swin-tiny":
								lowerCAmelCase__  :       Tuple				=  torch.tensor(
								    [[-7.5958, -7.5958, -7.4302], [-7.5958, -7.5958, -7.4302], [-7.4797, -7.4797, -7.3068]])
				elif model_name == "upernet-swin-small":
								lowerCAmelCase__  :       int				=  torch.tensor(
								    [[-7.1921, -7.1921, -6.9532], [-7.1921, -7.1921, -6.9532], [-7.0908, -7.0908, -6.8534]])
				elif model_name == "upernet-swin-base":
								lowerCAmelCase__  :       Optional[Any]				=  torch.tensor(
								    [[-6.5851, -6.5851, -6.4330], [-6.5851, -6.5851, -6.4330], [-6.4763, -6.4763, -6.3254]])
				elif model_name == "upernet-swin-large":
								lowerCAmelCase__  :       List[Any]				=  torch.tensor(
								    [[-7.5297, -7.5297, -7.3802], [-7.5297, -7.5297, -7.3802], [-7.4044, -7.4044, -7.2586]])
				print('''Logits:''' ,outputs.logits[0, 0, :3, :3])
				assert torch.allclose(outputs.logits[0, 0, :3, :3] ,lowerCamelCase_ ,atol=1E-4)
				print('''Looks ok!''')

				if pytorch_dump_folder_path is not None:
								print(f"""Saving model {model_name} to {pytorch_dump_folder_path}""")
								model.save_pretrained(lowerCamelCase_)
								print(f"""Saving processor to {pytorch_dump_folder_path}""")
								processor.save_pretrained(lowerCamelCase_)

				if push_to_hub:
								print(f"""Pushing model and processor for {model_name} to hub""")
								model.push_to_hub(f"""openmmlab/{model_name}""")
								processor.push_to_hub(f"""openmmlab/{model_name}""")


if __name__ == "__main__":
							__snake_case       :	str            =argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '--model_name',
							    default='upernet-swin-tiny',
							    type=str,
							    choices=[f"""upernet-swin-{size}""" for size in ['tiny', 'small', 'base', 'large']],
							    help='Name of the Swin + UperNet model you\'d like to convert.',
							)
							parser.add_argument(
							    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
							)
							parser.add_argument(
							    '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
							)

							__snake_case       :	str            =parser.parse_args()
							convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)






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  lowerCamelCase__    (							lowerCamelCase__   ,       lowerCamelCase__   ,       lowerCamelCase__   ,       unittest.TestCase):




		'''simple docstring'''



		snake_case_			=AltDiffusionPipeline
		snake_case_			=TEXT_TO_IMAGE_PARAMS
		snake_case_			=TEXT_TO_IMAGE_BATCH_PARAMS
		snake_case_			=TEXT_TO_IMAGE_IMAGE_PARAMS
		snake_case_			=TEXT_TO_IMAGE_IMAGE_PARAMS



		def 	lowerCAmelCase__   (self				)  ->    Dict:

						"""simple docstring"""






						torch.manual_seed(0				)
						lowerCAmelCase__  :       int				=  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	,)
						lowerCAmelCase__  :       List[str]				=  DDIMScheduler(
						    beta_start=0.0_0085	,beta_end=0.012	,beta_schedule='''scaled_linear'''	,clip_sample=__lowerCamelCase	,set_alpha_to_one=__lowerCamelCase	,)
						torch.manual_seed(0				)
						lowerCAmelCase__  :       Optional[Any]				=  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				)
						lowerCAmelCase__  :       Dict				=  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=50_02	,)
						lowerCAmelCase__  :       int				=  CLIPTextModel(__lowerCamelCase				)

						lowerCAmelCase__  :       str				=  XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''				)
						lowerCAmelCase__  :       Union[str, Any]				=  77

						lowerCAmelCase__  :       str				=  {
						    '''unet''': unet,
						    '''scheduler''': scheduler,
						    '''vae''': vae,
						    '''text_encoder''': text_encoder,
						    '''tokenizer''': tokenizer,
						    '''safety_checker''': None,
						    '''feature_extractor''': None,
						}
						return components



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase=0				)  ->    Optional[Any]:

						"""simple docstring"""






						if str(__lowerCamelCase				).startswith('''mps'''				):
										lowerCAmelCase__  :       Tuple				=  torch.manual_seed(__lowerCamelCase				)
						else:
										lowerCAmelCase__  :       Any				=  torch.Generator(device=__lowerCamelCase				).manual_seed(__lowerCamelCase				)
						lowerCAmelCase__  :       Optional[Any]				=  {
						    '''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 	lowerCAmelCase__   (self				)  ->    str:

						"""simple docstring"""






						super().test_attention_slicing_forward_pass(expected_max_diff=3e-3				)



		def 	lowerCAmelCase__   (self				)  ->    Tuple:

						"""simple docstring"""






						super().test_inference_batch_single_identical(expected_max_diff=3e-3				)



		def 	lowerCAmelCase__   (self				)  ->    Tuple:

						"""simple docstring"""






						lowerCAmelCase__  :       List[str]				=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator

						lowerCAmelCase__  :       int				=  self.get_dummy_components()
						torch.manual_seed(0				)
						lowerCAmelCase__  :       str				=  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=50_02	,)
						# TODO: remove after fixing the non-deterministic text encoder
						lowerCAmelCase__  :       Any				=  RobertaSeriesModelWithTransformation(__lowerCamelCase				)
						lowerCAmelCase__  :       List[str]				=  text_encoder

						lowerCAmelCase__  :       List[Any]				=  AltDiffusionPipeline(**__lowerCamelCase				)
						lowerCAmelCase__  :       Optional[Any]				=  alt_pipe.to(__lowerCamelCase				)
						alt_pipe.set_progress_bar_config(disable=__lowerCamelCase				)

						lowerCAmelCase__  :       List[str]				=  self.get_dummy_inputs(__lowerCamelCase				)
						lowerCAmelCase__  :       str				=  '''A photo of an astronaut'''
						lowerCAmelCase__  :       Any				=  alt_pipe(**__lowerCamelCase				)
						lowerCAmelCase__  :       int				=  output.images
						lowerCAmelCase__  :       List[Any]				=  image[0, -3:, -3:, -1]

						assert image.shape == (1, 64, 64, 3)
						lowerCAmelCase__  :       Optional[Any]				=  np.array(
						    [0.574_8162, 0.6044_7145, 0.4882_1217, 0.5010_0636, 0.543_1185, 0.4576_3683, 0.4965_7696, 0.4813_2733, 0.4757_3093]				)

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



		def 	lowerCAmelCase__   (self				)  ->    Optional[int]:

						"""simple docstring"""






						lowerCAmelCase__  :       List[Any]				=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator

						lowerCAmelCase__  :       str				=  self.get_dummy_components()
						lowerCAmelCase__  :       Optional[Any]				=  PNDMScheduler(skip_prk_steps=__lowerCamelCase				)
						torch.manual_seed(0				)
						lowerCAmelCase__  :       int				=  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=50_02	,)
						# TODO: remove after fixing the non-deterministic text encoder
						lowerCAmelCase__  :       Tuple				=  RobertaSeriesModelWithTransformation(__lowerCamelCase				)
						lowerCAmelCase__  :       List[str]				=  text_encoder
						lowerCAmelCase__  :       List[Any]				=  AltDiffusionPipeline(**__lowerCamelCase				)
						lowerCAmelCase__  :       str				=  alt_pipe.to(__lowerCamelCase				)
						alt_pipe.set_progress_bar_config(disable=__lowerCamelCase				)

						lowerCAmelCase__  :       Union[str, Any]				=  self.get_dummy_inputs(__lowerCamelCase				)
						lowerCAmelCase__  :       Optional[int]				=  alt_pipe(**__lowerCamelCase				)
						lowerCAmelCase__  :       Optional[int]				=  output.images
						lowerCAmelCase__  :       str				=  image[0, -3:, -3:, -1]

						assert image.shape == (1, 64, 64, 3)
						lowerCAmelCase__  :       Union[str, Any]				=  np.array(
						    [0.5160_5093, 0.570_7241, 0.4736_5507, 0.5057_8886, 0.563_3877, 0.464_2503, 0.518_2081, 0.4876_3484, 0.4908_4237]				)

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



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




		'''simple docstring'''



		def 	lowerCAmelCase__   (self				)  ->    List[Any]:

						"""simple docstring"""






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



		def 	lowerCAmelCase__   (self				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__  :       Union[str, Any]				=  AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''	,safety_checker=__lowerCamelCase				)
						lowerCAmelCase__  :       List[Any]				=  alt_pipe.to(__lowerCamelCase				)
						alt_pipe.set_progress_bar_config(disable=__lowerCamelCase				)

						lowerCAmelCase__  :       Any				=  '''A painting of a squirrel eating a burger'''
						lowerCAmelCase__  :       List[Any]				=  torch.manual_seed(0				)
						lowerCAmelCase__  :       str				=  alt_pipe([prompt]	,generator=__lowerCamelCase	,guidance_scale=6.0	,num_inference_steps=20	,output_type='''np'''				)

						lowerCAmelCase__  :       str				=  output.images

						lowerCAmelCase__  :       int				=  image[0, -3:, -3:, -1]

						assert image.shape == (1, 5_12, 5_12, 3)
						lowerCAmelCase__  :       Dict				=  np.array([0.1010, 0.0800, 0.0794, 0.0885, 0.0843, 0.0762, 0.0769, 0.0729, 0.0586]				)

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



		def 	lowerCAmelCase__   (self				)  ->    Any:

						"""simple docstring"""






						lowerCAmelCase__  :       List[Any]				=  DDIMScheduler.from_pretrained('''BAAI/AltDiffusion'''	,subfolder='''scheduler'''				)

						lowerCAmelCase__  :       List[str]				=  AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''	,scheduler=__lowerCamelCase	,safety_checker=__lowerCamelCase				)
						lowerCAmelCase__  :       Union[str, Any]				=  alt_pipe.to(__lowerCamelCase				)
						alt_pipe.set_progress_bar_config(disable=__lowerCamelCase				)

						lowerCAmelCase__  :       List[Any]				=  '''A painting of a squirrel eating a burger'''
						lowerCAmelCase__  :       Optional[Any]				=  torch.manual_seed(0				)

						lowerCAmelCase__  :       List[str]				=  alt_pipe([prompt]	,generator=__lowerCamelCase	,num_inference_steps=2	,output_type='''numpy'''				)
						lowerCAmelCase__  :       List[str]				=  output.images

						lowerCAmelCase__  :       Optional[int]				=  image[0, -3:, -3:, -1]

						assert image.shape == (1, 5_12, 5_12, 3)
						lowerCAmelCase__  :       List[Any]				=  np.array([0.4019, 0.4052, 0.3810, 0.4119, 0.3916, 0.3982, 0.4651, 0.4195, 0.5323]				)

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

