





import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
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 PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()







class 	_UpperCamelCase	(							A_    ,							A_    ,							unittest.TestCase	):







					'''simple docstring'''

					lowerCamelCase  :  Any           =							StableDiffusionParadigmsPipeline
					lowerCamelCase  :  Union[str, Any]           =							TEXT_TO_IMAGE_PARAMS
					lowerCamelCase  :  Union[str, Any]           =							TEXT_TO_IMAGE_BATCH_PARAMS
					lowerCamelCase  :  List[Any]           =							TEXT_TO_IMAGE_IMAGE_PARAMS
					lowerCamelCase  :  int           =							TEXT_TO_IMAGE_IMAGE_PARAMS





					def 							SCREAMING_SNAKE_CASE				(    self     :	str		):






						'''simple docstring'''







						torch.manual_seed(0		)
						UpperCAmelCase_		    =  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			,       attention_head_dim=(2, 4)			,       use_linear_projection=__lowercase			,       )
						UpperCAmelCase_		    =  DDIMParallelScheduler(
						    beta_start=0.00085			,       beta_end=0.012			,       beta_schedule="""scaled_linear"""			,       clip_sample=__lowercase			,       set_alpha_to_one=__lowercase			,       )
						torch.manual_seed(0		)
						UpperCAmelCase_		    =  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			,       sample_size=1_28			,       )
						torch.manual_seed(0		)
						UpperCAmelCase_		    =  CLIPTextConfig(
						    bos_token_id=0			,       eos_token_id=2			,       hidden_size=32			,       intermediate_size=37			,       layer_norm_eps=1e-05			,       num_attention_heads=4			,       num_hidden_layers=5			,       pad_token_id=1			,       vocab_size=10_00			,       hidden_act="""gelu"""			,       projection_dim=5_12			,       )
						UpperCAmelCase_		    =  CLIPTextModel(__lowercase		)
						UpperCAmelCase_		    =  CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""		)

						UpperCAmelCase_		    =  {
						    """unet""": unet,
						    """scheduler""": scheduler,
						    """vae""": vae,
						    """text_encoder""": text_encoder,
						    """tokenizer""": tokenizer,
						    """safety_checker""": None,
						    """feature_extractor""": None,
						}
						return components





					def 							SCREAMING_SNAKE_CASE				(    self     :	List[Any]			,       __lowercase     :	Tuple			,       __lowercase     :	List[Any]=0		):






						'''simple docstring'''







						if str(__lowercase		).startswith("""mps"""		):
							UpperCAmelCase_		    =  torch.manual_seed(__lowercase		)
						else:
							UpperCAmelCase_		    =  torch.Generator(device=__lowercase		).manual_seed(__lowercase		)
						UpperCAmelCase_		    =  {
						    """prompt""": """a photograph of an astronaut riding a horse""",
						    """generator""": generator,
						    """num_inference_steps""": 10,
						    """guidance_scale""": 6.0,
						    """output_type""": """numpy""",
						    """parallel""": 3,
						    """debug""": True,
						}
						return inputs





					def 							SCREAMING_SNAKE_CASE				(    self     :	Tuple		):






						'''simple docstring'''







						UpperCAmelCase_		    =  """cpu"""  # ensure determinism for the device-dependent torch.Generator
						UpperCAmelCase_		    =  self.get_dummy_components()
						UpperCAmelCase_		    =  StableDiffusionParadigmsPipeline(**__lowercase		)
						UpperCAmelCase_		    =  sd_pipe.to(__lowercase		)
						sd_pipe.set_progress_bar_config(disable=__lowercase		)

						UpperCAmelCase_		    =  self.get_dummy_inputs(__lowercase		)
						UpperCAmelCase_		    =  sd_pipe(**__lowercase		).images
						UpperCAmelCase_		    =  image[0, -3:, -3:, -1]
						assert image.shape == (1, 64, 64, 3)

						UpperCAmelCase_		    =  np.array([0.4773, 0.5417, 0.4723, 0.4925, 0.5631, 0.4752, 0.5240, 0.4935, 0.5023]		)

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





					def 							SCREAMING_SNAKE_CASE				(    self     :	List[str]		):






						'''simple docstring'''







						UpperCAmelCase_		    =  """cpu"""  # ensure determinism for the device-dependent torch.Generator
						UpperCAmelCase_		    =  self.get_dummy_components()
						torch.manual_seed(0		)
						UpperCAmelCase_		    =  DDPMParallelScheduler()
						torch.manual_seed(0		)
						UpperCAmelCase_		    =  StableDiffusionParadigmsPipeline(**__lowercase		)
						UpperCAmelCase_		    =  sd_pipe.to(__lowercase		)
						sd_pipe.set_progress_bar_config(disable=__lowercase		)

						UpperCAmelCase_		    =  self.get_dummy_inputs(__lowercase		)
						UpperCAmelCase_		    =  sd_pipe(**__lowercase		).images
						UpperCAmelCase_		    =  image[0, -3:, -3:, -1]
						assert image.shape == (1, 64, 64, 3)

						UpperCAmelCase_		    =  np.array([0.3573, 0.4420, 0.4960, 0.4799, 0.3796, 0.3879, 0.4819, 0.4365, 0.4468]		)

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





					def 							SCREAMING_SNAKE_CASE				(    self     :	List[str]		):






						'''simple docstring'''







						super().test_inference_batch_consistent(batch_sizes=[1, 2]		)





					def 							SCREAMING_SNAKE_CASE				(    self     :	Tuple		):






						'''simple docstring'''







						super().test_inference_batch_single_identical(batch_size=2			,       expected_max_diff=3e-3		)





					def 							SCREAMING_SNAKE_CASE				(    self     :	List[Any]		):






						'''simple docstring'''







						UpperCAmelCase_		    =  """cpu"""  # ensure determinism for the device-dependent torch.Generator
						UpperCAmelCase_		    =  self.get_dummy_components()
						UpperCAmelCase_		    =  StableDiffusionParadigmsPipeline(**__lowercase		)
						UpperCAmelCase_		    =  sd_pipe.to(__lowercase		)
						sd_pipe.set_progress_bar_config(disable=__lowercase		)

						UpperCAmelCase_		    =  self.get_dummy_inputs(__lowercase		)
						UpperCAmelCase_		    =  """french fries"""
						UpperCAmelCase_		    =  sd_pipe(**__lowercase			,       negative_prompt=__lowercase		)
						UpperCAmelCase_		    =  output.images
						UpperCAmelCase_		    =  image[0, -3:, -3:, -1]

						assert image.shape == (1, 64, 64, 3)

						UpperCAmelCase_		    =  np.array([0.4771, 0.5420, 0.4683, 0.4918, 0.5636, 0.4725, 0.5230, 0.4923, 0.5015]		)

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







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







					'''simple docstring'''

					def 							SCREAMING_SNAKE_CASE				(    self     :	str		):






						'''simple docstring'''







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





					def 							SCREAMING_SNAKE_CASE				(    self     :	Dict			,       __lowercase     :	str=0		):






						'''simple docstring'''







						UpperCAmelCase_		    =  torch.Generator(device=__lowercase		).manual_seed(__lowercase		)
						UpperCAmelCase_		    =  {
						    """prompt""": """a photograph of an astronaut riding a horse""",
						    """generator""": generator,
						    """num_inference_steps""": 10,
						    """guidance_scale""": 7.5,
						    """output_type""": """numpy""",
						    """parallel""": 3,
						    """debug""": True,
						}
						return inputs





					def 							SCREAMING_SNAKE_CASE				(    self     :	Optional[int]		):






						'''simple docstring'''







						UpperCAmelCase_		    =  """stabilityai/stable-diffusion-2-base"""
						UpperCAmelCase_		    =  DDIMParallelScheduler.from_pretrained(__lowercase			,       subfolder="""scheduler"""		)
						UpperCAmelCase_		    =  StableDiffusionParadigmsPipeline.from_pretrained(__lowercase			,       scheduler=__lowercase			,       safety_checker=__lowercase		)
						pipe.to(__lowercase		)
						pipe.set_progress_bar_config(disable=__lowercase		)
						pipe.enable_attention_slicing()

						UpperCAmelCase_		    =  self.get_inputs()
						UpperCAmelCase_		    =  pipe(**__lowercase		).images
						UpperCAmelCase_		    =  image[0, -3:, -3:, -1].flatten()

						assert image.shape == (1, 5_12, 5_12, 3)

						UpperCAmelCase_		    =  np.array([0.9622, 0.9602, 0.9748, 0.9591, 0.9630, 0.9691, 0.9661, 0.9631, 0.9741]		)

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









from typing import Any, Dict, List, Mapping

import numpy as np
import torch

from ...utils import is_cython_available, requires_backends


if is_cython_available():
				import pyximport

				pyximport.install(setup_args={"""include_dirs""": np.get_include()})
				from . import algos_graphormer  # noqa E402




def  A_(      A					,   A = 512 ):
	UpperCAmelCase_		    =  x.shape[1] if len(x.shape ) > 1 else 1
	UpperCAmelCase_		    =  1 + np.arange(0					,   feature_num * offset					,   A					,   dtype=np.intaa )
	UpperCAmelCase_		    =  x + feature_offset
	return x




def  A_(      A					,   A=True ):
	requires_backends(A					,   ["""cython"""] )

	if keep_features and "edge_attr" in item.keys():  # edge_attr
		UpperCAmelCase_		    =  np.asarray(item["""edge_attr"""]					,   dtype=np.intaa )
	else:
		UpperCAmelCase_		    =  np.ones((len(item["""edge_index"""][0] ), 1)					,   dtype=np.intaa )  # same embedding for all

	if keep_features and "node_feat" in item.keys():  # input_nodes
		UpperCAmelCase_		    =  np.asarray(item["""node_feat"""]					,   dtype=np.intaa )
	else:
		UpperCAmelCase_		    =  np.ones((item["""num_nodes"""], 1)					,   dtype=np.intaa )  # same embedding for all

	UpperCAmelCase_		    =  np.asarray(item["""edge_index"""]					,   dtype=np.intaa )

	UpperCAmelCase_		    =  convert_to_single_emb(A ) + 1
	UpperCAmelCase_		    =  item["""num_nodes"""]

	if len(edge_attr.shape ) == 1:
		UpperCAmelCase_		    =  edge_attr[:, None]
	UpperCAmelCase_		    =  np.zeros([num_nodes, num_nodes, edge_attr.shape[-1]]					,   dtype=np.intaa )
	UpperCAmelCase_		    =  convert_to_single_emb(A ) + 1

	# node adj matrix [num_nodes, num_nodes] bool
	UpperCAmelCase_		    =  np.zeros([num_nodes, num_nodes]					,   dtype=A )
	UpperCAmelCase_		    =  True

	UpperCAmelCase_    ,						UpperCAmelCase_		    =  algos_graphormer.floyd_warshall(A )
	UpperCAmelCase_		    =  np.amax(A )

	UpperCAmelCase_		    =  algos_graphormer.gen_edge_input(A					,   A					,   A )
	UpperCAmelCase_		    =  np.zeros([num_nodes + 1, num_nodes + 1]					,   dtype=np.single )  # with graph token

	# combine
	UpperCAmelCase_		    =  input_nodes + 1  # we shift all indices by one for padding
	UpperCAmelCase_		    =  attn_bias
	UpperCAmelCase_		    =  attn_edge_type
	UpperCAmelCase_		    =  shortest_path_result.astype(np.intaa ) + 1  # we shift all indices by one for padding
	UpperCAmelCase_		    =  np.sum(A					,   axis=1 ).reshape(-1 ) + 1  # we shift all indices by one for padding
	UpperCAmelCase_		    =  item["""in_degree"""]  # for undirected graph
	UpperCAmelCase_		    =  input_edges + 1  # we shift all indices by one for padding
	if "labels" not in item:
		UpperCAmelCase_		    =  item["""y"""]

	return item







class 	_UpperCamelCase	:







					'''simple docstring'''

					def __init__(    self     :	str			,       __lowercase     :	int=20			,       __lowercase     :	Union[str, Any]=False		):






						'''simple docstring'''







						if not is_cython_available():
							raise ImportError("""Graphormer preprocessing needs Cython (pyximport)"""		)

						UpperCAmelCase_		    =  spatial_pos_max
						UpperCAmelCase_		    =  on_the_fly_processing





					def __call__(    self     :	List[str]			,       __lowercase     :	List[dict]		):






						'''simple docstring'''







						if self.on_the_fly_processing:
							UpperCAmelCase_		    =  [preprocess_item(__lowercase		) for i in features]

						if not isinstance(features[0]			,       __lowercase		):
							UpperCAmelCase_		    =  [vars(__lowercase		) for f in features]
						UpperCAmelCase_		    =  {}

						UpperCAmelCase_		    =  max(len(i["""input_nodes"""]		) for i in features		)
						UpperCAmelCase_		    =  len(features[0]["""input_nodes"""][0]		)
						UpperCAmelCase_		    =  len(features[0]["""attn_edge_type"""][0][0]		)
						UpperCAmelCase_		    =  max(len(i["""input_edges"""][0][0]		) for i in features		)
						UpperCAmelCase_		    =  len(features[0]["""input_edges"""][0][0][0]		)
						UpperCAmelCase_		    =  len(__lowercase		)

						UpperCAmelCase_		    =  torch.zeros(__lowercase			,       max_node_num + 1			,       max_node_num + 1			,       dtype=torch.float		)
						UpperCAmelCase_		    =  torch.zeros(__lowercase			,       __lowercase			,       __lowercase			,       __lowercase			,       dtype=torch.long		)
						UpperCAmelCase_		    =  torch.zeros(__lowercase			,       __lowercase			,       __lowercase			,       dtype=torch.long		)
						UpperCAmelCase_		    =  torch.zeros(__lowercase			,       __lowercase			,       dtype=torch.long		)
						UpperCAmelCase_		    =  torch.zeros(__lowercase			,       __lowercase			,       __lowercase			,       dtype=torch.long		)
						UpperCAmelCase_		    =  torch.zeros(
						    __lowercase			,       __lowercase			,       __lowercase			,       __lowercase			,       __lowercase			,       dtype=torch.long		)

						for ix, f in enumerate(__lowercase		):
							for k in ["attn_bias", "attn_edge_type", "spatial_pos", "in_degree", "input_nodes", "input_edges"]:
								UpperCAmelCase_		    =  torch.tensor(f[k]		)

							if len(f["""attn_bias"""][1:, 1:][f["""spatial_pos"""] >= self.spatial_pos_max]		) > 0:
								UpperCAmelCase_		    =  float("""-inf"""		)

							UpperCAmelCase_		    =  f["""attn_bias"""]
							UpperCAmelCase_		    =  f[
							    """attn_edge_type"""
							]
							UpperCAmelCase_		    =  f["""spatial_pos"""]
							UpperCAmelCase_		    =  f["""in_degree"""]
							UpperCAmelCase_		    =  f["""input_nodes"""]
							UpperCAmelCase_		    =  f["""input_edges"""]

						UpperCAmelCase_		    =  batch["""in_degree"""]

						UpperCAmelCase_		    =  features[0]["""labels"""]
						if len(__lowercase		) == 1:  # one task
							if isinstance(sample[0]			,       __lowercase		):  # regression
								UpperCAmelCase_		    =  torch.from_numpy(np.concatenate([i["""labels"""] for i in features]		)		)
							else:  # binary classification
								UpperCAmelCase_		    =  torch.from_numpy(np.concatenate([i["""labels"""] for i in features]		)		)
						else:  # multi task classification, left to float to keep the NaNs
							UpperCAmelCase_		    =  torch.from_numpy(np.stack([i["""labels"""] for i in features]			,       axis=0		)		)

						return batch


