





import argparse
import json
from collections import OrderedDict
from pathlib import Path

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

from transformers import (
    ConditionalDetrConfig,
    ConditionalDetrForObjectDetection,
    ConditionalDetrForSegmentation,
    ConditionalDetrImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
_SCREAMING_SNAKE_CASE 		=						logging.get_logger(__name__)

# here we list all keys to be renamed (original name on the left, our name on the right)
_SCREAMING_SNAKE_CASE 		=						[]
for i in range(6):
						# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
						rename_keys.append(
						    (F'''transformer.encoder.layers.{i}.self_attn.out_proj.weight''', F'''encoder.layers.{i}.self_attn.out_proj.weight''')
						)
						rename_keys.append(
						    (F'''transformer.encoder.layers.{i}.self_attn.out_proj.bias''', F'''encoder.layers.{i}.self_attn.out_proj.bias''')
						)
						rename_keys.append((F'''transformer.encoder.layers.{i}.linear1.weight''', F'''encoder.layers.{i}.fc1.weight'''))
						rename_keys.append((F'''transformer.encoder.layers.{i}.linear1.bias''', F'''encoder.layers.{i}.fc1.bias'''))
						rename_keys.append((F'''transformer.encoder.layers.{i}.linear2.weight''', F'''encoder.layers.{i}.fc2.weight'''))
						rename_keys.append((F'''transformer.encoder.layers.{i}.linear2.bias''', F'''encoder.layers.{i}.fc2.bias'''))
						rename_keys.append(
						    (F'''transformer.encoder.layers.{i}.norm1.weight''', F'''encoder.layers.{i}.self_attn_layer_norm.weight''')
						)
						rename_keys.append((F'''transformer.encoder.layers.{i}.norm1.bias''', F'''encoder.layers.{i}.self_attn_layer_norm.bias'''))
						rename_keys.append((F'''transformer.encoder.layers.{i}.norm2.weight''', F'''encoder.layers.{i}.final_layer_norm.weight'''))
						rename_keys.append((F'''transformer.encoder.layers.{i}.norm2.bias''', F'''encoder.layers.{i}.final_layer_norm.bias'''))
						# decoder layers: 2 times output projection, 2 feedforward neural networks and 3 layernorms
						rename_keys.append(
						    (F'''transformer.decoder.layers.{i}.self_attn.out_proj.weight''', F'''decoder.layers.{i}.self_attn.out_proj.weight''')
						)
						rename_keys.append(
						    (F'''transformer.decoder.layers.{i}.self_attn.out_proj.bias''', F'''decoder.layers.{i}.self_attn.out_proj.bias''')
						)
						rename_keys.append(
						    (
						        F'''transformer.decoder.layers.{i}.cross_attn.out_proj.weight''',
						        F'''decoder.layers.{i}.encoder_attn.out_proj.weight''',
						    )
						)
						rename_keys.append(
						    (
						        F'''transformer.decoder.layers.{i}.cross_attn.out_proj.bias''',
						        F'''decoder.layers.{i}.encoder_attn.out_proj.bias''',
						    )
						)
						rename_keys.append((F'''transformer.decoder.layers.{i}.linear1.weight''', F'''decoder.layers.{i}.fc1.weight'''))
						rename_keys.append((F'''transformer.decoder.layers.{i}.linear1.bias''', F'''decoder.layers.{i}.fc1.bias'''))
						rename_keys.append((F'''transformer.decoder.layers.{i}.linear2.weight''', F'''decoder.layers.{i}.fc2.weight'''))
						rename_keys.append((F'''transformer.decoder.layers.{i}.linear2.bias''', F'''decoder.layers.{i}.fc2.bias'''))
						rename_keys.append(
						    (F'''transformer.decoder.layers.{i}.norm1.weight''', F'''decoder.layers.{i}.self_attn_layer_norm.weight''')
						)
						rename_keys.append((F'''transformer.decoder.layers.{i}.norm1.bias''', F'''decoder.layers.{i}.self_attn_layer_norm.bias'''))
						rename_keys.append(
						    (F'''transformer.decoder.layers.{i}.norm2.weight''', F'''decoder.layers.{i}.encoder_attn_layer_norm.weight''')
						)
						rename_keys.append(
						    (F'''transformer.decoder.layers.{i}.norm2.bias''', F'''decoder.layers.{i}.encoder_attn_layer_norm.bias''')
						)
						rename_keys.append((F'''transformer.decoder.layers.{i}.norm3.weight''', F'''decoder.layers.{i}.final_layer_norm.weight'''))
						rename_keys.append((F'''transformer.decoder.layers.{i}.norm3.bias''', F'''decoder.layers.{i}.final_layer_norm.bias'''))

						# q, k, v projections in self/cross-attention in decoder for conditional DETR
						rename_keys.append(
						    (F'''transformer.decoder.layers.{i}.sa_qcontent_proj.weight''', F'''decoder.layers.{i}.sa_qcontent_proj.weight''')
						)
						rename_keys.append(
						    (F'''transformer.decoder.layers.{i}.sa_kcontent_proj.weight''', F'''decoder.layers.{i}.sa_kcontent_proj.weight''')
						)
						rename_keys.append(
						    (F'''transformer.decoder.layers.{i}.sa_qpos_proj.weight''', F'''decoder.layers.{i}.sa_qpos_proj.weight''')
						)
						rename_keys.append(
						    (F'''transformer.decoder.layers.{i}.sa_kpos_proj.weight''', F'''decoder.layers.{i}.sa_kpos_proj.weight''')
						)
						rename_keys.append((F'''transformer.decoder.layers.{i}.sa_v_proj.weight''', F'''decoder.layers.{i}.sa_v_proj.weight'''))
						rename_keys.append(
						    (F'''transformer.decoder.layers.{i}.ca_qcontent_proj.weight''', F'''decoder.layers.{i}.ca_qcontent_proj.weight''')
						)
						# rename_keys.append((f"transformer.decoder.layers.{i}.ca_qpos_proj.weight", f"decoder.layers.{i}.ca_qpos_proj.weight"))
						rename_keys.append(
						    (F'''transformer.decoder.layers.{i}.ca_kcontent_proj.weight''', F'''decoder.layers.{i}.ca_kcontent_proj.weight''')
						)
						rename_keys.append(
						    (F'''transformer.decoder.layers.{i}.ca_kpos_proj.weight''', F'''decoder.layers.{i}.ca_kpos_proj.weight''')
						)
						rename_keys.append((F'''transformer.decoder.layers.{i}.ca_v_proj.weight''', F'''decoder.layers.{i}.ca_v_proj.weight'''))
						rename_keys.append(
						    (F'''transformer.decoder.layers.{i}.ca_qpos_sine_proj.weight''', F'''decoder.layers.{i}.ca_qpos_sine_proj.weight''')
						)

						rename_keys.append(
						    (F'''transformer.decoder.layers.{i}.sa_qcontent_proj.bias''', F'''decoder.layers.{i}.sa_qcontent_proj.bias''')
						)
						rename_keys.append(
						    (F'''transformer.decoder.layers.{i}.sa_kcontent_proj.bias''', F'''decoder.layers.{i}.sa_kcontent_proj.bias''')
						)
						rename_keys.append((F'''transformer.decoder.layers.{i}.sa_qpos_proj.bias''', F'''decoder.layers.{i}.sa_qpos_proj.bias'''))
						rename_keys.append((F'''transformer.decoder.layers.{i}.sa_kpos_proj.bias''', F'''decoder.layers.{i}.sa_kpos_proj.bias'''))
						rename_keys.append((F'''transformer.decoder.layers.{i}.sa_v_proj.bias''', F'''decoder.layers.{i}.sa_v_proj.bias'''))
						rename_keys.append(
						    (F'''transformer.decoder.layers.{i}.ca_qcontent_proj.bias''', F'''decoder.layers.{i}.ca_qcontent_proj.bias''')
						)
						# rename_keys.append((f"transformer.decoder.layers.{i}.ca_qpos_proj.bias", f"decoder.layers.{i}.ca_qpos_proj.bias"))
						rename_keys.append(
						    (F'''transformer.decoder.layers.{i}.ca_kcontent_proj.bias''', F'''decoder.layers.{i}.ca_kcontent_proj.bias''')
						)
						rename_keys.append((F'''transformer.decoder.layers.{i}.ca_kpos_proj.bias''', F'''decoder.layers.{i}.ca_kpos_proj.bias'''))
						rename_keys.append((F'''transformer.decoder.layers.{i}.ca_v_proj.bias''', F'''decoder.layers.{i}.ca_v_proj.bias'''))
						rename_keys.append(
						    (F'''transformer.decoder.layers.{i}.ca_qpos_sine_proj.bias''', F'''decoder.layers.{i}.ca_qpos_sine_proj.bias''')
						)

# convolutional projection + query embeddings + layernorm of decoder + class and bounding box heads
# for conditional DETR, also convert reference point head and query scale MLP
rename_keys.extend(
    [
        ("""input_proj.weight""", """input_projection.weight"""),
        ("""input_proj.bias""", """input_projection.bias"""),
        ("""query_embed.weight""", """query_position_embeddings.weight"""),
        ("""transformer.decoder.norm.weight""", """decoder.layernorm.weight"""),
        ("""transformer.decoder.norm.bias""", """decoder.layernorm.bias"""),
        ("""class_embed.weight""", """class_labels_classifier.weight"""),
        ("""class_embed.bias""", """class_labels_classifier.bias"""),
        ("""bbox_embed.layers.0.weight""", """bbox_predictor.layers.0.weight"""),
        ("""bbox_embed.layers.0.bias""", """bbox_predictor.layers.0.bias"""),
        ("""bbox_embed.layers.1.weight""", """bbox_predictor.layers.1.weight"""),
        ("""bbox_embed.layers.1.bias""", """bbox_predictor.layers.1.bias"""),
        ("""bbox_embed.layers.2.weight""", """bbox_predictor.layers.2.weight"""),
        ("""bbox_embed.layers.2.bias""", """bbox_predictor.layers.2.bias"""),
        ("""transformer.decoder.ref_point_head.layers.0.weight""", """decoder.ref_point_head.layers.0.weight"""),
        ("""transformer.decoder.ref_point_head.layers.0.bias""", """decoder.ref_point_head.layers.0.bias"""),
        ("""transformer.decoder.ref_point_head.layers.1.weight""", """decoder.ref_point_head.layers.1.weight"""),
        ("""transformer.decoder.ref_point_head.layers.1.bias""", """decoder.ref_point_head.layers.1.bias"""),
        ("""transformer.decoder.query_scale.layers.0.weight""", """decoder.query_scale.layers.0.weight"""),
        ("""transformer.decoder.query_scale.layers.0.bias""", """decoder.query_scale.layers.0.bias"""),
        ("""transformer.decoder.query_scale.layers.1.weight""", """decoder.query_scale.layers.1.weight"""),
        ("""transformer.decoder.query_scale.layers.1.bias""", """decoder.query_scale.layers.1.bias"""),
        ("""transformer.decoder.layers.0.ca_qpos_proj.weight""", """decoder.layers.0.ca_qpos_proj.weight"""),
        ("""transformer.decoder.layers.0.ca_qpos_proj.bias""", """decoder.layers.0.ca_qpos_proj.bias"""),
    ]
)


def 							SCREAMING_SNAKE_CASE__  (      __a				,	__a				,	__a				):
	snake_case_						:     str     =		state_dict.pop(__a				)
	snake_case_						:     Optional[Any]     =		val


def 							SCREAMING_SNAKE_CASE__  (      __a				):
	snake_case_						:     List[str]     =		OrderedDict()
	for key, value in state_dict.items():
		if "backbone.0.body" in key:
			snake_case_						:     str     =		key.replace('backbone.0.body'				,	'backbone.conv_encoder.model'				)
			snake_case_						:     List[Any]     =		value
		else:
			snake_case_						:     Optional[int]     =		value

	return new_state_dict


def 							SCREAMING_SNAKE_CASE__  (      __a				,	__a=False				):
	snake_case_						:     Optional[Any]     =		''
	if is_panoptic:
		snake_case_						:     Dict     =		'conditional_detr.'

	# first: transformer encoder
	for i in range(6				):
		# read in weights + bias of input projection layer (in PyTorch's MultiHeadAttention, this is a single matrix + bias)
		snake_case_						:     int     =		state_dict.pop(f"""{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight"""				)
		snake_case_						:     str     =		state_dict.pop(f"""{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_bias"""				)
		# next, add query, keys and values (in that order) to the state dict
		snake_case_						:     List[Any]     =		in_proj_weight[:2_56, :]
		snake_case_						:     List[str]     =		in_proj_bias[:2_56]
		snake_case_						:     Dict     =		in_proj_weight[2_56:5_12, :]
		snake_case_						:     Union[str, Any]     =		in_proj_bias[2_56:5_12]
		snake_case_						:     str     =		in_proj_weight[-2_56:, :]
		snake_case_						:     List[str]     =		in_proj_bias[-2_56:]


def 							SCREAMING_SNAKE_CASE__  (      ):
	snake_case_						:     Optional[int]     =		'http://images.cocodataset.org/val2017/000000039769.jpg'
	snake_case_						:     List[Any]     =		Image.open(requests.get(__a				,	stream=__a				).raw				)

	return im


@torch.no_grad()
def 							SCREAMING_SNAKE_CASE__  (      __a				,	__a				):
	snake_case_						:     List[str]     =		ConditionalDetrConfig()
	# set backbone and dilation attributes
	if "resnet101" in model_name:
		snake_case_						:     Tuple     =		'resnet101'
	if "dc5" in model_name:
		snake_case_						:     int     =		True
	snake_case_						:     Union[str, Any]     =		'panoptic' in model_name
	if is_panoptic:
		snake_case_						:     Optional[int]     =		2_50
	else:
		snake_case_						:     Union[str, Any]     =		91
		snake_case_						:     str     =		'huggingface/label-files'
		snake_case_						:     Optional[Any]     =		'coco-detection-id2label.json'
		snake_case_						:     Optional[Any]     =		json.load(open(hf_hub_download(__a				,	__a				,	repo_type='dataset'				)				,	'r'				)				)
		snake_case_						:     Dict     =		{int(__a				): v for k, v in idalabel.items()}
		snake_case_						:     List[str]     =		idalabel
		snake_case_						:     Union[str, Any]     =		{v: k for k, v in idalabel.items()}

	# load image processor
	snake_case_						:     Any     =		'coco_panoptic' if is_panoptic else 'coco_detection'
	snake_case_						:     List[Any]     =		ConditionalDetrImageProcessor(format=__a				)

	# prepare image
	snake_case_						:     Any     =		prepare_img()
	snake_case_						:     List[Any]     =		image_processor(images=__a				,	return_tensors='pt'				)
	snake_case_						:     Tuple     =		encoding['pixel_values']

	logger.info(f"""Converting model {model_name}..."""				)

	# load original model from torch hub
	snake_case_						:     Tuple     =		torch.hub.load('DeppMeng/ConditionalDETR'				,	__a				,	pretrained=__a				).eval()
	snake_case_						:     str     =		conditional_detr.state_dict()
	# rename keys
	for src, dest in rename_keys:
		if is_panoptic:
			snake_case_						:     Optional[int]     =		'conditional_detr.' + src
		rename_key(__a				,	__a				,	__a				)
	snake_case_						:     List[str]     =		rename_backbone_keys(__a				)
	# query, key and value matrices need special treatment
	read_in_q_k_v(__a				,	is_panoptic=__a				)
	# important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
	snake_case_						:     List[str]     =		'conditional_detr.model.' if is_panoptic else 'model.'
	for key in state_dict.copy().keys():
		if is_panoptic:
			if (
			    key.startswith('conditional_detr'				)
			    and not key.startswith('class_labels_classifier'				)
			    and not key.startswith('bbox_predictor'				)
			):
				snake_case_						:     Dict     =		state_dict.pop(__a				)
				snake_case_						:     List[Any]     =		val
			elif "class_labels_classifier" in key or "bbox_predictor" in key:
				snake_case_						:     Any     =		state_dict.pop(__a				)
				snake_case_						:     str     =		val
			elif key.startswith('bbox_attention'				) or key.startswith('mask_head'				):
				continue
			else:
				snake_case_						:     str     =		state_dict.pop(__a				)
				snake_case_						:     Dict     =		val
		else:
			if not key.startswith('class_labels_classifier'				) and not key.startswith('bbox_predictor'				):
				snake_case_						:     List[Any]     =		state_dict.pop(__a				)
				snake_case_						:     Union[str, Any]     =		val
    # finally, create HuggingFace model and load state dict
	snake_case_						:     Optional[Any]     =		ConditionalDetrForSegmentation(__a				) if is_panoptic else ConditionalDetrForObjectDetection(__a				)
	model.load_state_dict(__a				)
	model.eval()
	model.push_to_hub(repo_id=__a				,	organization='DepuMeng'				,	commit_message='Add model'				)
	# verify our conversion
	snake_case_						:     Optional[Any]     =		conditional_detr(__a				)
	snake_case_						:     Optional[int]     =		model(__a				)
	assert torch.allclose(outputs.logits				,	original_outputs['pred_logits']				,	atol=1E-4				)
	assert torch.allclose(outputs.pred_boxes				,	original_outputs['pred_boxes']				,	atol=1E-4				)
	if is_panoptic:
		assert torch.allclose(outputs.pred_masks				,	original_outputs['pred_masks']				,	atol=1E-4				)

	# Save model and image processor
	logger.info(f"""Saving PyTorch model and image processor to {pytorch_dump_folder_path}..."""				)
	Path(__a				).mkdir(exist_ok=__a				)
	model.save_pretrained(__a				)
	image_processor.save_pretrained(__a				)


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

						parser.add_argument(
						    """--model_name""",
						    default="""conditional_detr_resnet50""",
						    type=str,
						    help="""Name of the CONDITIONAL_DETR model you'd like to convert.""",
						)
						parser.add_argument(
						    """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the folder to output PyTorch model."""
						)
						_SCREAMING_SNAKE_CASE 		=						parser.parse_args()
						convert_conditional_detr_checkpoint(args.model_name, args.pytorch_dump_folder_path)








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 SCREAMING_SNAKE_CASE_  (   snake_case_    ,			snake_case_    ,			snake_case_    ,			unittest.TestCase   ):
						__magic_name__:     int	 =    AltDiffusionPipeline
						__magic_name__:     Any	 =    TEXT_TO_IMAGE_PARAMS
						__magic_name__:     Any	 =    TEXT_TO_IMAGE_BATCH_PARAMS
						__magic_name__:     Any	 =    TEXT_TO_IMAGE_IMAGE_PARAMS
						__magic_name__:     Union[str, Any]	 =    TEXT_TO_IMAGE_IMAGE_PARAMS






						def UpperCAmelCase_			( self : List[Any]		)							->      int:







							"""simple docstring"""


							torch.manual_seed(0		)
							snake_case_						:     List[str]     =		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						,				)
							snake_case_						:     Optional[Any]     =		DDIMScheduler(
							    beta_start=0.0_0_0_8_5						,				beta_end=0.0_1_2						,				beta_schedule='scaled_linear'						,				clip_sample=_A						,				set_alpha_to_one=_A						,				)
							torch.manual_seed(0		)
							snake_case_						:     Union[str, 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		)
							snake_case_						:     Optional[Any]     =		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=5002						,				)
							snake_case_						:     Any     =		CLIPTextModel(_A		)

							snake_case_						:     Any     =		XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'		)
							snake_case_						:     Dict     =		77

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






						def UpperCAmelCase_			( self : int						,				_A : Optional[int]						,				_A : int=0		)							->      Dict:







							"""simple docstring"""


							if str(_A		).startswith('mps'		):
								snake_case_						:     Union[str, Any]     =		torch.manual_seed(_A		)
							else:
								snake_case_						:     Union[str, Any]     =		torch.Generator(device=_A		).manual_seed(_A		)
							snake_case_						:     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 UpperCAmelCase_			( self : Optional[Any]		)							->      Any:







							"""simple docstring"""


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






						def UpperCAmelCase_			( self : List[Any]		)							->      Dict:







							"""simple docstring"""


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






						def UpperCAmelCase_			( self : Dict		)							->      Any:







							"""simple docstring"""


							snake_case_						:     Tuple     =		'cpu'  # ensure determinism for the device-dependent torch.Generator

							snake_case_						:     Any     =		self.get_dummy_components()
							torch.manual_seed(0		)
							snake_case_						:     Any     =		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						,				)
							# TODO: remove after fixing the non-deterministic text encoder
							snake_case_						:     Optional[Any]     =		RobertaSeriesModelWithTransformation(_A		)
							snake_case_						:     Optional[Any]     =		text_encoder

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

							snake_case_						:     Optional[Any]     =		self.get_dummy_inputs(_A		)
							snake_case_						:     int     =		'A photo of an astronaut'
							snake_case_						:     Tuple     =		alt_pipe(**_A		)
							snake_case_						:     Any     =		output.images
							snake_case_						:     Optional[int]     =		image[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)
							snake_case_						:     Any     =		np.array(
							    [0.5_7_4_8_1_6_2, 0.6_0_4_4_7_1_4_5, 0.4_8_8_2_1_2_1_7, 0.5_0_1_0_0_6_3_6, 0.5_4_3_1_1_8_5, 0.4_5_7_6_3_6_8_3, 0.4_9_6_5_7_6_9_6, 0.4_8_1_3_2_7_3_3, 0.4_7_5_7_3_0_9_3]		)

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






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







							"""simple docstring"""


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

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

							snake_case_						:     Dict     =		self.get_dummy_inputs(_A		)
							snake_case_						:     Tuple     =		alt_pipe(**_A		)
							snake_case_						:     int     =		output.images
							snake_case_						:     List[Any]     =		image[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)
							snake_case_						:     Optional[int]     =		np.array(
							    [0.5_1_6_0_5_0_9_3, 0.5_7_0_7_2_4_1, 0.4_7_3_6_5_5_0_7, 0.5_0_5_7_8_8_8_6, 0.5_6_3_3_8_7_7, 0.4_6_4_2_5_0_3, 0.5_1_8_2_0_8_1, 0.4_8_7_6_3_4_8_4, 0.4_9_0_8_4_2_3_7]		)

							assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2
@slow
@require_torch_gpu
class SCREAMING_SNAKE_CASE_  (   unittest.TestCase   ):






						def UpperCAmelCase_			( self : List[Any]		)							->      Tuple:







							"""simple docstring"""


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






						def UpperCAmelCase_			( self : int		)							->      List[str]:







							"""simple docstring"""


							snake_case_						:     Optional[int]     =		AltDiffusionPipeline.from_pretrained('BAAI/AltDiffusion'						,				safety_checker=_A		)
							snake_case_						:     Optional[int]     =		alt_pipe.to(_A		)
							alt_pipe.set_progress_bar_config(disable=_A		)

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

							snake_case_						:     Any     =		output.images

							snake_case_						:     List[Any]     =		image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							snake_case_						:     Union[str, Any]     =		np.array([0.1_0_1_0, 0.0_8_0_0, 0.0_7_9_4, 0.0_8_8_5, 0.0_8_4_3, 0.0_7_6_2, 0.0_7_6_9, 0.0_7_2_9, 0.0_5_8_6]		)

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






						def UpperCAmelCase_			( self : Optional[int]		)							->      Union[str, Any]:







							"""simple docstring"""


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

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

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

							snake_case_						:     List[Any]     =		alt_pipe([prompt]						,				generator=_A						,				num_inference_steps=2						,				output_type='numpy'		)
							snake_case_						:     Any     =		output.images

							snake_case_						:     Dict     =		image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							snake_case_						:     List[Any]     =		np.array([0.4_0_1_9, 0.4_0_5_2, 0.3_8_1_0, 0.4_1_1_9, 0.3_9_1_6, 0.3_9_8_2, 0.4_6_5_1, 0.4_1_9_5, 0.5_3_2_3]		)

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

