






"""simple docstring"""



import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLPipeline,
    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 							__lowercase(							_UpperCamelCase    ,				_UpperCamelCase    ,				unittest.TestCase	):
						lowerCAmelCase_		:				Optional[Any]=							StableDiffusionXLPipeline
						lowerCAmelCase_		:				str=							TEXT_TO_IMAGE_PARAMS
						lowerCAmelCase_		:				Optional[int]=							TEXT_TO_IMAGE_BATCH_PARAMS
						lowerCAmelCase_		:				Dict=							TEXT_TO_IMAGE_IMAGE_PARAMS
						lowerCAmelCase_		:				int=							TEXT_TO_IMAGE_IMAGE_PARAMS






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




											"""simple docstring"""





											torch.manual_seed(0)
											lowerCamelCase__					=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')     ,    attention_head_dim=(2, 4)     ,    use_linear_projection=_SCREAMING_SNAKE_CASE     ,    addition_embed_type='text_time'     ,    addition_time_embed_dim=8     ,    transformer_layers_per_block=(1, 2)     ,    projection_class_embeddings_input_dim=80     ,    cross_attention_dim=64     ,    )
											lowerCamelCase__					=EulerDiscreteScheduler(
											    beta_start=0.00085     ,    beta_end=0.012     ,    steps_offset=1     ,    beta_schedule='scaled_linear'     ,    timestep_spacing='leading'     ,    )
											torch.manual_seed(0)
											lowerCamelCase__					=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)
											lowerCamelCase__					=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=32     ,    )
											lowerCamelCase__					=CLIPTextModel(_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					=CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'     ,    local_files_only=_SCREAMING_SNAKE_CASE)

											lowerCamelCase__					=CLIPTextModelWithProjection(_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					=CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'     ,    local_files_only=_SCREAMING_SNAKE_CASE)

											lowerCamelCase__					={
											    'unet': unet,
											    'scheduler': scheduler,
											    'vae': vae,
											    'text_encoder': text_encoder,
											    'tokenizer': tokenizer,
											    'text_encoder_2': text_encoder_a,
											    'tokenizer_2': tokenizer_a,
											    # "safety_checker": None,
											    # "feature_extractor": None,
											}
											return components






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE=0)      ->	Optional[int]:




											"""simple docstring"""





											if str(_SCREAMING_SNAKE_CASE).startswith('mps'):
																lowerCamelCase__					=torch.manual_seed(_SCREAMING_SNAKE_CASE)
											else:
																lowerCamelCase__					=torch.Generator(device=_SCREAMING_SNAKE_CASE).manual_seed(_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					={
											    'prompt': 'A painting of a squirrel eating a burger',
											    'generator': generator,
											    'num_inference_steps': 2,
											    'guidance_scale': 5.0,
											    'output_type': 'numpy',
											}
											return inputs






						def 		lowerCAmelCase (       self)      ->	Optional[Any]:




											"""simple docstring"""





											lowerCamelCase__					='cpu'  # ensure determinism for the device-dependent torch.Generator
											lowerCamelCase__					=self.get_dummy_components()
											lowerCamelCase__					=StableDiffusionXLPipeline(**_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					=sd_pipe.to(_SCREAMING_SNAKE_CASE)
											sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)

											lowerCamelCase__					=self.get_dummy_inputs(_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					=sd_pipe(**_SCREAMING_SNAKE_CASE).images
											lowerCamelCase__					=image[0, -3:, -3:, -1]

											assert image.shape == (1, 64, 64, 3)
											lowerCamelCase__					=np.array([0.5873, 0.6128, 0.4797, 0.5122, 0.5674, 0.4639, 0.5227, 0.5149, 0.4747])

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






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




											"""simple docstring"""





											lowerCamelCase__					=self.get_dummy_components()
											lowerCamelCase__					=StableDiffusionXLPipeline(**_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					=sd_pipe.to(_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					=sd_pipe.to(_SCREAMING_SNAKE_CASE)
											sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)

											# forward without prompt embeds
											lowerCamelCase__					=self.get_dummy_inputs(_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					=3 * ['this is a negative prompt']
											lowerCamelCase__					=negative_prompt
											lowerCamelCase__					=3 * [inputs['prompt']]

											lowerCamelCase__					=sd_pipe(**_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					=output.images[0, -3:, -3:, -1]

											# forward with prompt embeds
											lowerCamelCase__					=self.get_dummy_inputs(_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					=3 * ['this is a negative prompt']
											lowerCamelCase__					=3 * [inputs.pop('prompt')]

											(
											    (
											    lowerCamelCase__
											)					,  (
											    lowerCamelCase__
											)					,  (
											    lowerCamelCase__
											)					,  (
											    lowerCamelCase__
											)					,  
											)					=sd_pipe.encode_prompt(_SCREAMING_SNAKE_CASE     ,    negative_prompt=_SCREAMING_SNAKE_CASE)

											lowerCamelCase__					=sd_pipe(
											    **_SCREAMING_SNAKE_CASE     ,    prompt_embeds=_SCREAMING_SNAKE_CASE     ,    negative_prompt_embeds=_SCREAMING_SNAKE_CASE     ,    pooled_prompt_embeds=_SCREAMING_SNAKE_CASE     ,    negative_pooled_prompt_embeds=_SCREAMING_SNAKE_CASE     ,    )
											lowerCamelCase__					=output.images[0, -3:, -3:, -1]

											# make sure that it's equal
											assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()).max() < 1E-4






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




											"""simple docstring"""





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






						def 		lowerCAmelCase (       self)      ->	Any:




											"""simple docstring"""





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




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






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




											"""simple docstring"""





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






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE="cpu"     ,    _SCREAMING_SNAKE_CASE=torch.floataa     ,    _SCREAMING_SNAKE_CASE=0)      ->	str:




											"""simple docstring"""





											lowerCamelCase__					=torch.Generator(device=_SCREAMING_SNAKE_CASE).manual_seed(_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					=np.random.RandomState(_SCREAMING_SNAKE_CASE).standard_normal((1, 4, 64, 64))
											lowerCamelCase__					=torch.from_numpy(_SCREAMING_SNAKE_CASE).to(device=_SCREAMING_SNAKE_CASE     ,    dtype=_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					={
											    'prompt': 'a photograph of an astronaut riding a horse',
											    'latents': latents,
											    'generator': generator,
											    'num_inference_steps': 3,
											    'guidance_scale': 7.5,
											    'output_type': 'numpy',
											}
											return inputs






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




											"""simple docstring"""





											lowerCamelCase__					=DiffusionPipeline.from_pretrained('stabilityai/stable-diffusion-2-base')
											pipe.to(_SCREAMING_SNAKE_CASE)
											pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)

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

											assert image.shape == (1, 5_12, 5_12, 3)
											lowerCamelCase__					=np.array([0.49493, 0.47896, 0.40798, 0.54214, 0.53212, 0.48202, 0.47656, 0.46329, 0.48506])
											assert np.abs(image_slice - expected_slice).max() < 7E-3







"""simple docstring"""



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     lowerCamelCase_					(	__A	:		List[str]      ,						__A	:		int = 512      ):
					lowerCamelCase__					=x.shape[1] if len(x.shape      ) > 1 else 1
					lowerCamelCase__					=1 + np.arange(0      ,						feature_num * offset      ,						__A      ,						dtype=np.intaa      )
					lowerCamelCase__					=x + feature_offset
					return x




def     lowerCamelCase_					(	__A	:		str      ,						__A	:		Optional[Any]=True      ):
					requires_backends(__A      ,						['cython']      )

					if keep_features and "edge_attr" in item.keys():  # edge_attr
										lowerCamelCase__					=np.asarray(item['edge_attr']      ,						dtype=np.intaa      )
					else:
										lowerCamelCase__					=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
										lowerCamelCase__					=np.asarray(item['node_feat']      ,						dtype=np.intaa      )
					else:
										lowerCamelCase__					=np.ones((item['num_nodes'], 1)      ,						dtype=np.intaa      )  # same embedding for all

					lowerCamelCase__					=np.asarray(item['edge_index']      ,						dtype=np.intaa      )

					lowerCamelCase__					=convert_to_single_emb(__A      ) + 1
					lowerCamelCase__					=item['num_nodes']

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

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

					lowerCamelCase__					,  lowerCamelCase__					=algos_graphormer.floyd_warshall(__A      )
					lowerCamelCase__					=np.amax(__A      )

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

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

					return item




class 							__lowercase:






						def __init__(       self     ,    _SCREAMING_SNAKE_CASE=20     ,    _SCREAMING_SNAKE_CASE=False)      ->	Any:




											"""simple docstring"""





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

											lowerCamelCase__					=spatial_pos_max
											lowerCamelCase__					=on_the_fly_processing






						def __call__(       self     ,    _SCREAMING_SNAKE_CASE)      ->	Dict[str, Any]:




											"""simple docstring"""





											if self.on_the_fly_processing:
																lowerCamelCase__					=[preprocess_item(_SCREAMING_SNAKE_CASE) for i in features]

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

											lowerCamelCase__					=max(len(i['input_nodes']) for i in features)
											lowerCamelCase__					=len(features[0]['input_nodes'][0])
											lowerCamelCase__					=len(features[0]['attn_edge_type'][0][0])
											lowerCamelCase__					=max(len(i['input_edges'][0][0]) for i in features)
											lowerCamelCase__					=len(features[0]['input_edges'][0][0][0])
											lowerCamelCase__					=len(_SCREAMING_SNAKE_CASE)

											lowerCamelCase__					=torch.zeros(_SCREAMING_SNAKE_CASE     ,    max_node_num + 1     ,    max_node_num + 1     ,    dtype=torch.float)
											lowerCamelCase__					=torch.zeros(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    dtype=torch.long)
											lowerCamelCase__					=torch.zeros(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    dtype=torch.long)
											lowerCamelCase__					=torch.zeros(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    dtype=torch.long)
											lowerCamelCase__					=torch.zeros(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    dtype=torch.long)
											lowerCamelCase__					=torch.zeros(
											    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    dtype=torch.long)

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

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

																lowerCamelCase__					=f['attn_bias']
																lowerCamelCase__					=f[
																    'attn_edge_type'
																]
																lowerCamelCase__					=f['spatial_pos']
																lowerCamelCase__					=f['in_degree']
																lowerCamelCase__					=f['input_nodes']
																lowerCamelCase__					=f['input_edges']

											lowerCamelCase__					=batch['in_degree']

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

											return batch