




"""simple docstring"""





import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import cached_download, hf_hub_url
from PIL import Image

from transformers import DPTConfig, DPTForDepthEstimation, DPTForSemanticSegmentation, DPTImageProcessor
from transformers.utils import logging


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

def 					_UpperCamelCase    (			UpperCamelCase							)			->			Optional[int]:


			"""simple docstring"""
			__UpperCAmelCase	: Union[str, Any]		=					DPTConfig()

			if "large" in checkpoint_url:
						__UpperCAmelCase	: Any		=					1024
						__UpperCAmelCase	: List[str]		=					4096
						__UpperCAmelCase	: List[Any]		=					24
						__UpperCAmelCase	: Any		=					16
						__UpperCAmelCase	: Optional[Any]		=					[5, 11, 17, 23]
						__UpperCAmelCase	: List[Any]		=					[256, 512, 1024, 1024]
						__UpperCAmelCase	: Tuple		=					(1, 384, 384)

			if "ade" in checkpoint_url:
						__UpperCAmelCase	: int		=					True

						__UpperCAmelCase	: Dict		=					150
						__UpperCAmelCase	: Optional[Any]		=					"huggingface/label-files"
						__UpperCAmelCase	: Dict		=					"ade20k-id2label.json"
						__UpperCAmelCase	: Union[str, Any]		=					json.load(open(cached_download(hf_hub_url(UpperCamelCase				,       UpperCamelCase				,       repo_type="dataset"							)							)				,       "r"							)							)
						__UpperCAmelCase	: Tuple		=					{int(UpperCamelCase							): v for k, v in idalabel.items()}
						__UpperCAmelCase	: List[Any]		=					idalabel
						__UpperCAmelCase	: Optional[Any]		=					{v: k for k, v in idalabel.items()}
						__UpperCAmelCase	: str		=					[1, 150, 480, 480]

			return config, expected_shape

def 					_UpperCamelCase    (			UpperCamelCase							)			->			Union[str, Any]:


			"""simple docstring"""
			__UpperCAmelCase	: Optional[Any]		=					["pretrained.model.head.weight", "pretrained.model.head.bias"]
			for k in ignore_keys:
						state_dict.pop(UpperCamelCase				,       UpperCamelCase							)

def 					_UpperCamelCase    (			UpperCamelCase							)			->			Any:


			"""simple docstring"""
			if (
			    "pretrained.model" in name
			    and "cls_token" not in name
			    and "pos_embed" not in name
			    and "patch_embed" not in name
			):
						__UpperCAmelCase	: Optional[Any]		=					name.replace("pretrained.model"				,       "dpt.encoder"							)
			if "pretrained.model" in name:
						__UpperCAmelCase	: str		=					name.replace("pretrained.model"				,       "dpt.embeddings"							)
			if "patch_embed" in name:
						__UpperCAmelCase	: List[Any]		=					name.replace("patch_embed"				,       "patch_embeddings"							)
			if "pos_embed" in name:
						__UpperCAmelCase	: Dict		=					name.replace("pos_embed"				,       "position_embeddings"							)
			if "attn.proj" in name:
						__UpperCAmelCase	: int		=					name.replace("attn.proj"				,       "attention.output.dense"							)
			if "proj" in name and "project" not in name:
						__UpperCAmelCase	: List[Any]		=					name.replace("proj"				,       "projection"							)
			if "blocks" in name:
						__UpperCAmelCase	: Union[str, Any]		=					name.replace("blocks"				,       "layer"							)
			if "mlp.fc1" in name:
						__UpperCAmelCase	: Dict		=					name.replace("mlp.fc1"				,       "intermediate.dense"							)
			if "mlp.fc2" in name:
						__UpperCAmelCase	: str		=					name.replace("mlp.fc2"				,       "output.dense"							)
			if "norm1" in name:
						__UpperCAmelCase	: List[Any]		=					name.replace("norm1"				,       "layernorm_before"							)
			if "norm2" in name:
						__UpperCAmelCase	: List[str]		=					name.replace("norm2"				,       "layernorm_after"							)
			if "scratch.output_conv" in name:
						__UpperCAmelCase	: str		=					name.replace("scratch.output_conv"				,       "head"							)
			if "scratch" in name:
						__UpperCAmelCase	: Dict		=					name.replace("scratch"				,       "neck"							)
			if "layer1_rn" in name:
						__UpperCAmelCase	: Optional[Any]		=					name.replace("layer1_rn"				,       "convs.0"							)
			if "layer2_rn" in name:
						__UpperCAmelCase	: Any		=					name.replace("layer2_rn"				,       "convs.1"							)
			if "layer3_rn" in name:
						__UpperCAmelCase	: Union[str, Any]		=					name.replace("layer3_rn"				,       "convs.2"							)
			if "layer4_rn" in name:
						__UpperCAmelCase	: List[Any]		=					name.replace("layer4_rn"				,       "convs.3"							)
			if "refinenet" in name:
						__UpperCAmelCase	: int		=					int(name[len("neck.refinenet"							) : len("neck.refinenet"							) + 1]							)
						# tricky here: we need to map 4 to 0, 3 to 1, 2 to 2 and 1 to 3
						__UpperCAmelCase	: Union[str, Any]		=					name.replace(f"refinenet{layer_idx}"				,       f"fusion_stage.layers.{abs(layer_idx-4							)}"							)
			if "out_conv" in name:
						__UpperCAmelCase	: List[Any]		=					name.replace("out_conv"				,       "projection"							)
			if "resConfUnit1" in name:
						__UpperCAmelCase	: Any		=					name.replace("resConfUnit1"				,       "residual_layer1"							)
			if "resConfUnit2" in name:
						__UpperCAmelCase	: Dict		=					name.replace("resConfUnit2"				,       "residual_layer2"							)
			if "conv1" in name:
						__UpperCAmelCase	: Union[str, Any]		=					name.replace("conv1"				,       "convolution1"							)
			if "conv2" in name:
						__UpperCAmelCase	: str		=					name.replace("conv2"				,       "convolution2"							)
			# readout blocks
			if "pretrained.act_postprocess1.0.project.0" in name:
						__UpperCAmelCase	: Dict		=					name.replace("pretrained.act_postprocess1.0.project.0"				,       "neck.reassemble_stage.readout_projects.0.0"							)
			if "pretrained.act_postprocess2.0.project.0" in name:
						__UpperCAmelCase	: Dict		=					name.replace("pretrained.act_postprocess2.0.project.0"				,       "neck.reassemble_stage.readout_projects.1.0"							)
			if "pretrained.act_postprocess3.0.project.0" in name:
						__UpperCAmelCase	: str		=					name.replace("pretrained.act_postprocess3.0.project.0"				,       "neck.reassemble_stage.readout_projects.2.0"							)
			if "pretrained.act_postprocess4.0.project.0" in name:
						__UpperCAmelCase	: Dict		=					name.replace("pretrained.act_postprocess4.0.project.0"				,       "neck.reassemble_stage.readout_projects.3.0"							)
			# resize blocks
			if "pretrained.act_postprocess1.3" in name:
						__UpperCAmelCase	: str		=					name.replace("pretrained.act_postprocess1.3"				,       "neck.reassemble_stage.layers.0.projection"							)
			if "pretrained.act_postprocess1.4" in name:
						__UpperCAmelCase	: List[str]		=					name.replace("pretrained.act_postprocess1.4"				,       "neck.reassemble_stage.layers.0.resize"							)
			if "pretrained.act_postprocess2.3" in name:
						__UpperCAmelCase	: Tuple		=					name.replace("pretrained.act_postprocess2.3"				,       "neck.reassemble_stage.layers.1.projection"							)
			if "pretrained.act_postprocess2.4" in name:
						__UpperCAmelCase	: Optional[Any]		=					name.replace("pretrained.act_postprocess2.4"				,       "neck.reassemble_stage.layers.1.resize"							)
			if "pretrained.act_postprocess3.3" in name:
						__UpperCAmelCase	: Any		=					name.replace("pretrained.act_postprocess3.3"				,       "neck.reassemble_stage.layers.2.projection"							)
			if "pretrained.act_postprocess4.3" in name:
						__UpperCAmelCase	: str		=					name.replace("pretrained.act_postprocess4.3"				,       "neck.reassemble_stage.layers.3.projection"							)
			if "pretrained.act_postprocess4.4" in name:
						__UpperCAmelCase	: str		=					name.replace("pretrained.act_postprocess4.4"				,       "neck.reassemble_stage.layers.3.resize"							)
			if "pretrained" in name:
						__UpperCAmelCase	: Union[str, Any]		=					name.replace("pretrained"				,       "dpt"							)
			if "bn" in name:
						__UpperCAmelCase	: int		=					name.replace("bn"				,       "batch_norm"							)
			if "head" in name:
						__UpperCAmelCase	: Optional[int]		=					name.replace("head"				,       "head.head"							)
			if "encoder.norm" in name:
						__UpperCAmelCase	: Union[str, Any]		=					name.replace("encoder.norm"				,       "layernorm"							)
			if "auxlayer" in name:
						__UpperCAmelCase	: Any		=					name.replace("auxlayer"				,       "auxiliary_head.head"							)

			return name

def 					_UpperCamelCase    (			UpperCamelCase				,       UpperCamelCase							)			->			Dict:


			"""simple docstring"""
			for i in range(config.num_hidden_layers							):
						# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
						__UpperCAmelCase	: List[Any]		=					state_dict.pop(f"dpt.encoder.layer.{i}.attn.qkv.weight"							)
						__UpperCAmelCase	: List[str]		=					state_dict.pop(f"dpt.encoder.layer.{i}.attn.qkv.bias"							)
						# next, add query, keys and values (in that order) to the state dict
						__UpperCAmelCase	: List[str]		=					in_proj_weight[: config.hidden_size, :]
						__UpperCAmelCase	: List[str]		=					in_proj_bias[: config.hidden_size]
						__UpperCAmelCase	: List[str]		=					in_proj_weight[
						    config.hidden_size : config.hidden_size * 2, :
						]
						__UpperCAmelCase	: List[str]		=					in_proj_bias[
						    config.hidden_size : config.hidden_size * 2
						]
						__UpperCAmelCase	: List[Any]		=					in_proj_weight[
						    -config.hidden_size :, :
						]
						__UpperCAmelCase	: Dict		=					in_proj_bias[-config.hidden_size :]

def 					_UpperCamelCase    (			)			->			Union[str, Any]:


			"""simple docstring"""
			__UpperCAmelCase	: Dict		=					"http://images.cocodataset.org/val2017/000000039769.jpg"
			__UpperCAmelCase	: Optional[int]		=					Image.open(requests.get(UpperCamelCase				,       stream=UpperCamelCase							).raw							)
			return im

@torch.no_grad()
def 					_UpperCamelCase    (			UpperCamelCase				,       UpperCamelCase				,       UpperCamelCase				,       UpperCamelCase							)			->			Union[str, Any]:


			"""simple docstring"""
			__UpperCAmelCase	,     __UpperCAmelCase	: Optional[int]		=					get_dpt_config(UpperCamelCase							)
			# load original state_dict from URL
			__UpperCAmelCase	: List[str]		=					torch.hub.load_state_dict_from_url(UpperCamelCase				,       map_location="cpu"							)
			# remove certain keys
			remove_ignore_keys_(UpperCamelCase							)
			# rename keys
			for key in state_dict.copy().keys():
						__UpperCAmelCase	: int		=					state_dict.pop(UpperCamelCase							)
						__UpperCAmelCase	: Union[str, Any]		=					val
			# read in qkv matrices
			read_in_q_k_v(UpperCamelCase				,       UpperCamelCase							)

			# load HuggingFace model
			__UpperCAmelCase	: int		=					DPTForSemanticSegmentation(UpperCamelCase							) if "ade" in checkpoint_url else DPTForDepthEstimation(UpperCamelCase							)
			model.load_state_dict(UpperCamelCase							)
			model.eval()

			# Check outputs on an image
			__UpperCAmelCase	: Tuple		=					480 if "ade" in checkpoint_url else 384
			__UpperCAmelCase	: Dict		=					DPTImageProcessor(size=UpperCamelCase							)

			__UpperCAmelCase	: Union[str, Any]		=					prepare_img()
			__UpperCAmelCase	: Union[str, Any]		=					image_processor(UpperCamelCase				,       return_tensors="pt"							)

			# forward pass
			__UpperCAmelCase	: Optional[int]		=					model(**UpperCamelCase							).logits if "ade" in checkpoint_url else model(**UpperCamelCase							).predicted_depth

			# Assert logits
			__UpperCAmelCase	: List[Any]		=					torch.tensor([[6.3199, 6.3629, 6.4148], [6.3850, 6.3615, 6.4166], [6.3519, 6.3176, 6.3575]]							)
			if "ade" in checkpoint_url:
						__UpperCAmelCase	: Any		=					torch.tensor([[4.0480, 4.2420, 4.4360], [4.3124, 4.5693, 4.8261], [4.5768, 4.8965, 5.2163]]							)
			assert outputs.shape == torch.Size(UpperCamelCase							)
			assert (
			    torch.allclose(outputs[0, 0, :3, :3]				,       UpperCamelCase				,       atol=1e-4							)
			    if "ade" in checkpoint_url
			    else torch.allclose(outputs[0, :3, :3]				,       UpperCamelCase							)
			)

			Path(UpperCamelCase							).mkdir(exist_ok=UpperCamelCase							)
			print(f"Saving model to {pytorch_dump_folder_path}"							)
			model.save_pretrained(UpperCamelCase							)
			print(f"Saving image processor to {pytorch_dump_folder_path}"							)
			image_processor.save_pretrained(UpperCamelCase							)

			if push_to_hub:
						print("Pushing model to hub..."							)
						model.push_to_hub(
						    repo_path_or_name=Path(UpperCamelCase				,       UpperCamelCase							)				,       organization="nielsr"				,       commit_message="Add model"				,       use_temp_dir=UpperCamelCase				,       )
						image_processor.push_to_hub(
						    repo_path_or_name=Path(UpperCamelCase				,       UpperCamelCase							)				,       organization="nielsr"				,       commit_message="Add image processor"				,       use_temp_dir=UpperCamelCase				,       )


if __name__ == "__main__":
					A   =       argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    """--checkpoint_url""",
					    default="""https://github.com/intel-isl/DPT/releases/download/1_0/dpt_large-midas-2f21e586.pt""",
					    type=str,
					    help="""URL of the original DPT checkpoint you'd like to convert.""",
					)
					parser.add_argument(
					    """--pytorch_dump_folder_path""",
					    default=None,
					    type=str,
					    required=True,
					    help="""Path to the output PyTorch model directory.""",
					)
					parser.add_argument(
					    """--push_to_hub""",
					    action="""store_true""",
					)
					parser.add_argument(
					    """--model_name""",
					    default="""dpt-large""",
					    type=str,
					    help="""Name of the model, in case you're pushing to the hub.""",
					)

					A   =       parser.parse_args()
					convert_dpt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub, args.model_name)








"""simple docstring"""





import gc
import unittest

import numpy as np
import torch
from torch.backends.cuda import sdp_kernel

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)
from diffusers.utils import randn_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_a, require_torch_gpu

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()



class      a__				(	__magic_name__							,    unittest.TestCase   ):
			lowercase_    =	ConsistencyModelPipeline
			lowercase_    =	UNCONDITIONAL_IMAGE_GENERATION_PARAMS
			lowercase_    =	UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

			# Override required_optional_params to remove num_images_per_prompt
			lowercase_    =	frozenset(
			    [
			        "num_inference_steps",
			        "generator",
			        "latents",
			        "output_type",
			        "return_dict",
			        "callback",
			        "callback_steps",
			    ]   )


			@property
			def       a_     (      self      :    Optional[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					UNetaDModel.from_pretrained(
						    "diffusers/consistency-models-test"      ,  subfolder="test_unet"      ,  )
						return unet


			@property
			def       a_     (      self      :    List[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Dict		=					UNetaDModel.from_pretrained(
						    "diffusers/consistency-models-test"      ,  subfolder="test_unet_class_cond"      ,  )
						return unet


			def       a_     (      self      :    Any      ,  UpperCamelCase_      :    int=False):


						"""simple docstring"""




						if class_cond:
									__UpperCAmelCase	: List[Any]		=					self.dummy_cond_unet
						else:
									__UpperCAmelCase	: Optional[int]		=					self.dummy_uncond_unet

						# Default to CM multistep sampler
						__UpperCAmelCase	: List[str]		=					CMStochasticIterativeScheduler(
						    num_train_timesteps=40      ,  sigma_min=0.002      ,  sigma_max=80.0      ,  )

						__UpperCAmelCase	: Optional[int]		=					{
						    "unet": unet,
						    "scheduler": scheduler,
						}

						return components


			def       a_     (      self      :    List[str]      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    List[Any]=0):


						"""simple docstring"""




						if str(UpperCamelCase_).startswith("mps"):
									__UpperCAmelCase	: str		=					torch.manual_seed(UpperCamelCase_)
						else:
									__UpperCAmelCase	: Optional[Any]		=					torch.Generator(device=UpperCamelCase_).manual_seed(UpperCamelCase_)

						__UpperCAmelCase	: List[Any]		=					{
						    "batch_size": 1,
						    "num_inference_steps": None,
						    "timesteps": [22, 0],
						    "generator": generator,
						    "output_type": "np",
						}

						return inputs


			def       a_     (      self      :    int):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					"cpu"  # ensure determinism for the device-dependent torch.Generator
						__UpperCAmelCase	: Union[str, Any]		=					self.get_dummy_components()
						__UpperCAmelCase	: str		=					ConsistencyModelPipeline(**UpperCamelCase_)
						__UpperCAmelCase	: Any		=					pipe.to(UpperCamelCase_)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: Any		=					self.get_dummy_inputs(UpperCamelCase_)
						__UpperCAmelCase	: str		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 32, 32, 3)

						__UpperCAmelCase	: Dict		=					image[0, -3:, -3:, -1]
						__UpperCAmelCase	: List[Any]		=					np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004])

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


			def       a_     (      self      :    Optional[int]):


						"""simple docstring"""




						__UpperCAmelCase	: List[Any]		=					"cpu"  # ensure determinism for the device-dependent torch.Generator
						__UpperCAmelCase	: Optional[int]		=					self.get_dummy_components(class_cond=UpperCamelCase_)
						__UpperCAmelCase	: Optional[Any]		=					ConsistencyModelPipeline(**UpperCamelCase_)
						__UpperCAmelCase	: str		=					pipe.to(UpperCamelCase_)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: Tuple		=					self.get_dummy_inputs(UpperCamelCase_)
						__UpperCAmelCase	: Dict		=					0
						__UpperCAmelCase	: int		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 32, 32, 3)

						__UpperCAmelCase	: Any		=					image[0, -3:, -3:, -1]
						__UpperCAmelCase	: List[Any]		=					np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004])

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


			def       a_     (      self      :    Optional[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					"cpu"  # ensure determinism for the device-dependent torch.Generator
						__UpperCAmelCase	: Optional[int]		=					self.get_dummy_components()
						__UpperCAmelCase	: Optional[Any]		=					ConsistencyModelPipeline(**UpperCamelCase_)
						__UpperCAmelCase	: Optional[Any]		=					pipe.to(UpperCamelCase_)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: List[str]		=					self.get_dummy_inputs(UpperCamelCase_)
						__UpperCAmelCase	: Dict		=					1
						__UpperCAmelCase	: int		=					None
						__UpperCAmelCase	: List[Any]		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 32, 32, 3)

						__UpperCAmelCase	: Tuple		=					image[0, -3:, -3:, -1]
						__UpperCAmelCase	: Union[str, Any]		=					np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987])

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


			def       a_     (      self      :    List[str]):


						"""simple docstring"""




						__UpperCAmelCase	: Dict		=					"cpu"  # ensure determinism for the device-dependent torch.Generator
						__UpperCAmelCase	: List[str]		=					self.get_dummy_components(class_cond=UpperCamelCase_)
						__UpperCAmelCase	: Tuple		=					ConsistencyModelPipeline(**UpperCamelCase_)
						__UpperCAmelCase	: int		=					pipe.to(UpperCamelCase_)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: Tuple		=					self.get_dummy_inputs(UpperCamelCase_)
						__UpperCAmelCase	: Optional[int]		=					1
						__UpperCAmelCase	: Optional[int]		=					None
						__UpperCAmelCase	: Tuple		=					0
						__UpperCAmelCase	: Tuple		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 32, 32, 3)

						__UpperCAmelCase	: List[str]		=					image[0, -3:, -3:, -1]
						__UpperCAmelCase	: Dict		=					np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987])

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






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


			def       a_     (      self      :    Optional[int]):


						"""simple docstring"""




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


			def       a_     (      self      :    Union[str, Any]      ,  UpperCamelCase_      :    Optional[int]=0      ,  UpperCamelCase_      :    Tuple=False      ,  UpperCamelCase_      :    int="cpu"      ,  UpperCamelCase_      :    Any=torch.floataa      ,  UpperCamelCase_      :    List[str]=(1, 3, 64, 64)):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					torch.manual_seed(UpperCamelCase_)

						__UpperCAmelCase	: int		=					{
						    "num_inference_steps": None,
						    "timesteps": [22, 0],
						    "class_labels": 0,
						    "generator": generator,
						    "output_type": "np",
						}

						if get_fixed_latents:
									__UpperCAmelCase	: int		=					self.get_fixed_latents(seed=UpperCamelCase_      ,  device=UpperCamelCase_      ,  dtype=UpperCamelCase_      ,  shape=UpperCamelCase_)
									__UpperCAmelCase	: Optional[int]		=					latents

						return inputs


			def       a_     (      self      :    Union[str, Any]      ,  UpperCamelCase_      :    int=0      ,  UpperCamelCase_      :    Tuple="cpu"      ,  UpperCamelCase_      :    Tuple=torch.floataa      ,  UpperCamelCase_      :    Optional[Any]=(1, 3, 64, 64)):


						"""simple docstring"""




						if type(UpperCamelCase_) == str:
									__UpperCAmelCase	: Union[str, Any]		=					torch.device(UpperCamelCase_)
						__UpperCAmelCase	: Union[str, Any]		=					torch.Generator(device=UpperCamelCase_).manual_seed(UpperCamelCase_)
						__UpperCAmelCase	: Optional[int]		=					randn_tensor(UpperCamelCase_      ,  generator=UpperCamelCase_      ,  device=UpperCamelCase_      ,  dtype=UpperCamelCase_)
						return latents


			def       a_     (      self      :    List[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					UNetaDModel.from_pretrained("diffusers/consistency_models"      ,  subfolder="diffusers_cd_imagenet64_l2")
						__UpperCAmelCase	: Dict		=					CMStochasticIterativeScheduler(
						    num_train_timesteps=40      ,  sigma_min=0.002      ,  sigma_max=80.0      ,  )
						__UpperCAmelCase	: Dict		=					ConsistencyModelPipeline(unet=UpperCamelCase_      ,  scheduler=UpperCamelCase_)
						pipe.to(torch_device=UpperCamelCase_)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: Dict		=					self.get_inputs()
						__UpperCAmelCase	: List[str]		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 64, 64, 3)

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

						__UpperCAmelCase	: Union[str, Any]		=					np.array([0.0888, 0.0881, 0.0666, 0.0479, 0.0292, 0.0195, 0.0201, 0.0163, 0.0254])

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


			def       a_     (      self      :    Tuple):


						"""simple docstring"""




						__UpperCAmelCase	: List[str]		=					UNetaDModel.from_pretrained("diffusers/consistency_models"      ,  subfolder="diffusers_cd_imagenet64_l2")
						__UpperCAmelCase	: Optional[Any]		=					CMStochasticIterativeScheduler(
						    num_train_timesteps=40      ,  sigma_min=0.002      ,  sigma_max=80.0      ,  )
						__UpperCAmelCase	: Any		=					ConsistencyModelPipeline(unet=UpperCamelCase_      ,  scheduler=UpperCamelCase_)
						pipe.to(torch_device=UpperCamelCase_)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: int		=					self.get_inputs()
						__UpperCAmelCase	: str		=					1
						__UpperCAmelCase	: Union[str, Any]		=					None
						__UpperCAmelCase	: Tuple		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 64, 64, 3)

						__UpperCAmelCase	: str		=					image[0, -3:, -3:, -1]

						__UpperCAmelCase	: Union[str, Any]		=					np.array([0.0340, 0.0152, 0.0063, 0.0267, 0.0221, 0.0107, 0.0416, 0.0186, 0.0217])

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


			@require_torch_a
			def       a_     (      self      :    Tuple):


						"""simple docstring"""




						__UpperCAmelCase	: int		=					UNetaDModel.from_pretrained("diffusers/consistency_models"      ,  subfolder="diffusers_cd_imagenet64_l2")
						__UpperCAmelCase	: Optional[int]		=					CMStochasticIterativeScheduler(
						    num_train_timesteps=40      ,  sigma_min=0.002      ,  sigma_max=80.0      ,  )
						__UpperCAmelCase	: int		=					ConsistencyModelPipeline(unet=UpperCamelCase_      ,  scheduler=UpperCamelCase_)
						pipe.to(torch_device=UpperCamelCase_      ,  torch_dtype=torch.floataa)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: Tuple		=					self.get_inputs(get_fixed_latents=UpperCamelCase_      ,  device=UpperCamelCase_)
						# Ensure usage of flash attention in torch 2.0
						with sdp_kernel(enable_flash=UpperCamelCase_      ,  enable_math=UpperCamelCase_      ,  enable_mem_efficient=UpperCamelCase_):
									__UpperCAmelCase	: List[str]		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 64, 64, 3)

						__UpperCAmelCase	: List[str]		=					image[0, -3:, -3:, -1]

						__UpperCAmelCase	: Optional[int]		=					np.array([0.1875, 0.1428, 0.1289, 0.2151, 0.2092, 0.1477, 0.1877, 0.1641, 0.1353])

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


			@require_torch_a
			def       a_     (      self      :    str):


						"""simple docstring"""




						__UpperCAmelCase	: Tuple		=					UNetaDModel.from_pretrained("diffusers/consistency_models"      ,  subfolder="diffusers_cd_imagenet64_l2")
						__UpperCAmelCase	: Union[str, Any]		=					CMStochasticIterativeScheduler(
						    num_train_timesteps=40      ,  sigma_min=0.002      ,  sigma_max=80.0      ,  )
						__UpperCAmelCase	: List[Any]		=					ConsistencyModelPipeline(unet=UpperCamelCase_      ,  scheduler=UpperCamelCase_)
						pipe.to(torch_device=UpperCamelCase_      ,  torch_dtype=torch.floataa)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: Optional[int]		=					self.get_inputs(get_fixed_latents=UpperCamelCase_      ,  device=UpperCamelCase_)
						__UpperCAmelCase	: List[str]		=					1
						__UpperCAmelCase	: Any		=					None
						# Ensure usage of flash attention in torch 2.0
						with sdp_kernel(enable_flash=UpperCamelCase_      ,  enable_math=UpperCamelCase_      ,  enable_mem_efficient=UpperCamelCase_):
									__UpperCAmelCase	: List[str]		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 64, 64, 3)

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

						__UpperCAmelCase	: List[str]		=					np.array([0.1663, 0.1948, 0.2275, 0.1680, 0.1204, 0.1245, 0.1858, 0.1338, 0.2095])

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


