



"""simple docstring"""
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     lowerCamelCase  (      _lowerCAmelCase   ,						_lowerCAmelCase   ,						unittest.TestCase    ):






					'''simple docstring'''





					_A		:      Dict						=					StableDiffusionParadigmsPipeline
					_A		:      Dict						=					TEXT_TO_IMAGE_PARAMS
					_A		:      List[Any]						=					TEXT_TO_IMAGE_BATCH_PARAMS
					_A		:      List[Any]						=					TEXT_TO_IMAGE_IMAGE_PARAMS
					_A		:      Dict						=					TEXT_TO_IMAGE_IMAGE_PARAMS






					def        lowerCAmelCase_    (	self:							Optional[Any]						)       ->     Optional[Any]:
											torch.manual_seed(0						)
											snake_case_			: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						,					attention_head_dim=(2, 4)						,					use_linear_projection=snake_case						,					)
											snake_case_			:Optional[Any]						=      DDIMParallelScheduler(
											    beta_start=0.0_0_0_8_5						,					beta_end=0.0_1_2						,					beta_schedule="""scaled_linear"""						,					clip_sample=snake_case						,					set_alpha_to_one=snake_case						,					)
											torch.manual_seed(0						)
											snake_case_			:List[str]						=      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=128						,					)
											torch.manual_seed(0						)
											snake_case_			:Union[str, Any]						=      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=1_000						,					hidden_act="""gelu"""						,					projection_dim=512						,					)
											snake_case_			:Optional[Any]						=      CLIPTextModel(snake_case						)
											snake_case_			:Optional[Any]						=      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

											snake_case_			:Optional[int]						=      {
											    """unet""": unet,
											    """scheduler""": scheduler,
											    """vae""": vae,
											    """text_encoder""": text_encoder,
											    """tokenizer""": tokenizer,
											    """safety_checker""": None,
											    """feature_extractor""": None,
											}
											return components






					def        lowerCAmelCase_    (	self:							int						,					snake_case:							Union[str, Any]						,					snake_case:							Union[str, Any]=0						)       ->     List[Any]:
											if str(snake_case						).startswith("""mps"""						):
																	snake_case_			:Tuple						=      torch.manual_seed(snake_case						)
											else:
																	snake_case_			:Any						=      torch.Generator(device=snake_case						).manual_seed(snake_case						)
											snake_case_			:Union[str, Any]						=      {
											    """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        lowerCAmelCase_    (	self:							Optional[Any]						)       ->     Optional[int]:
											snake_case_			:Tuple						=      """cpu"""  # ensure determinism for the device-dependent torch.Generator
											snake_case_			:Any						=      self.get_dummy_components()
											snake_case_			:Union[str, Any]						=      StableDiffusionParadigmsPipeline(**snake_case						)
											snake_case_			:Any						=      sd_pipe.to(snake_case						)
											sd_pipe.set_progress_bar_config(disable=snake_case						)

											snake_case_			:Optional[int]						=      self.get_dummy_inputs(snake_case						)
											snake_case_			:Dict						=      sd_pipe(**snake_case						).images
											snake_case_			:Dict						=      image[0, -3:, -3:, -1]
											assert image.shape == (1, 64, 64, 3)

											snake_case_			:Optional[int]						=      np.array([0.4_7_7_3, 0.5_4_1_7, 0.4_7_2_3, 0.4_9_2_5, 0.5_6_3_1, 0.4_7_5_2, 0.5_2_4_0, 0.4_9_3_5, 0.5_0_2_3]						)

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






					def        lowerCAmelCase_    (	self:							int						)       ->     Optional[int]:
											snake_case_			:Tuple						=      """cpu"""  # ensure determinism for the device-dependent torch.Generator
											snake_case_			:Tuple						=      self.get_dummy_components()
											torch.manual_seed(0						)
											snake_case_			:Tuple						=      DDPMParallelScheduler()
											torch.manual_seed(0						)
											snake_case_			:Optional[int]						=      StableDiffusionParadigmsPipeline(**snake_case						)
											snake_case_			:List[Any]						=      sd_pipe.to(snake_case						)
											sd_pipe.set_progress_bar_config(disable=snake_case						)

											snake_case_			:List[str]						=      self.get_dummy_inputs(snake_case						)
											snake_case_			:Tuple						=      sd_pipe(**snake_case						).images
											snake_case_			:List[str]						=      image[0, -3:, -3:, -1]
											assert image.shape == (1, 64, 64, 3)

											snake_case_			:str						=      np.array([0.3_5_7_3, 0.4_4_2_0, 0.4_9_6_0, 0.4_7_9_9, 0.3_7_9_6, 0.3_8_7_9, 0.4_8_1_9, 0.4_3_6_5, 0.4_4_6_8]						)

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






					def        lowerCAmelCase_    (	self:							Any						)       ->     Optional[int]:
											super().test_inference_batch_consistent(batch_sizes=[1, 2]						)






					def        lowerCAmelCase_    (	self:							Optional[int]						)       ->     Optional[int]:
											super().test_inference_batch_single_identical(batch_size=2						,					expected_max_diff=3E-3						)






					def        lowerCAmelCase_    (	self:							int						)       ->     Optional[Any]:
											snake_case_			:str						=      """cpu"""  # ensure determinism for the device-dependent torch.Generator
											snake_case_			:List[str]						=      self.get_dummy_components()
											snake_case_			:Dict						=      StableDiffusionParadigmsPipeline(**snake_case						)
											snake_case_			:Optional[int]						=      sd_pipe.to(snake_case						)
											sd_pipe.set_progress_bar_config(disable=snake_case						)

											snake_case_			:Optional[int]						=      self.get_dummy_inputs(snake_case						)
											snake_case_			:List[str]						=      """french fries"""
											snake_case_			:Any						=      sd_pipe(**snake_case						,					negative_prompt=snake_case						)
											snake_case_			:Union[str, Any]						=      output.images
											snake_case_			:str						=      image[0, -3:, -3:, -1]

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

											snake_case_			:Optional[int]						=      np.array([0.4_7_7_1, 0.5_4_2_0, 0.4_6_8_3, 0.4_9_1_8, 0.5_6_3_6, 0.4_7_2_5, 0.5_2_3_0, 0.4_9_2_3, 0.5_0_1_5]						)

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

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






					'''simple docstring'''





					def        lowerCAmelCase_    (	self:							Union[str, Any]						)       ->     Any:
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()






					def        lowerCAmelCase_    (	self:							Any						,					snake_case:							Optional[int]=0						)       ->     Tuple:
											snake_case_			:Dict						=      torch.Generator(device=snake_case						).manual_seed(snake_case						)
											snake_case_			:Tuple						=      {
											    """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        lowerCAmelCase_    (	self:							str						)       ->     Any:
											snake_case_			:Any						=      """stabilityai/stable-diffusion-2-base"""
											snake_case_			:Optional[Any]						=      DDIMParallelScheduler.from_pretrained(snake_case						,					subfolder="""scheduler"""						)
											snake_case_			:Tuple						=      StableDiffusionParadigmsPipeline.from_pretrained(snake_case						,					scheduler=snake_case						,					safety_checker=snake_case						)
											pipe.to(snake_case						)
											pipe.set_progress_bar_config(disable=snake_case						)
											pipe.enable_attention_slicing()

											snake_case_			:int						=      self.get_inputs()
											snake_case_			:List[str]						=      pipe(**snake_case						).images
											snake_case_			:Optional[int]						=      image[0, -3:, -3:, -1].flatten()

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

											snake_case_			:Tuple						=      np.array([0.9_6_2_2, 0.9_6_0_2, 0.9_7_4_8, 0.9_5_9_1, 0.9_6_3_0, 0.9_6_9_1, 0.9_6_6_1, 0.9_6_3_1, 0.9_7_4_1]						)

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





"""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        A_  (      _lowercase, _lowercase = 512     ):







						'''simple docstring'''


						snake_case_			:int						=      x.shape[1] if len(x.shape     ) > 1 else 1
						snake_case_			:List[Any]						=      1 + np.arange(0, feature_num * offset, _lowercase, dtype=np.intaa     )
						snake_case_			:int						=      x + feature_offset
						return x
def        A_  (      _lowercase, _lowercase=True     ):







						'''simple docstring'''


						requires_backends(_lowercase, ["""cython"""]     )

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

						snake_case_			:Any						=      np.asarray(item["""edge_index"""], dtype=np.intaa     )

						snake_case_			:str						=      convert_to_single_emb(_lowercase     ) + 1
						snake_case_			:str						=      item["""num_nodes"""]

						if len(edge_attr.shape     ) == 1:
												snake_case_			:Any						=      edge_attr[:, None]
						snake_case_			:Dict						=      np.zeros([num_nodes, num_nodes, edge_attr.shape[-1]], dtype=np.intaa     )
						snake_case_			:str						=      convert_to_single_emb(_lowercase     ) + 1

						# node adj matrix [num_nodes, num_nodes] bool
						snake_case_			:List[str]						=      np.zeros([num_nodes, num_nodes], dtype=_lowercase     )
						snake_case_			:int						=      True

						snake_case_,	snake_case_			:Tuple						=      algos_graphormer.floyd_warshall(_lowercase     )
						snake_case_			:List[str]						=      np.amax(_lowercase     )

						snake_case_			:Dict						=      algos_graphormer.gen_edge_input(_lowercase, _lowercase, _lowercase     )
						snake_case_			:Tuple						=      np.zeros([num_nodes + 1, num_nodes + 1], dtype=np.single     )  # with graph token

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

						return item

class     lowerCamelCase  :






					'''simple docstring'''





					def __init__(	self:							Dict						,					snake_case:							Optional[int]=20						,					snake_case:							List[Any]=False						)       ->     Union[str, Any]:
											if not is_cython_available():
																	raise ImportError("""Graphormer preprocessing needs Cython (pyximport)"""						)

											snake_case_			:List[str]						=      spatial_pos_max
											snake_case_			:int						=      on_the_fly_processing






					def __call__(	self:							str						,					snake_case:							List[dict]						)       ->     Dict[str, Any]:
											if self.on_the_fly_processing:
																	snake_case_			:int						=      [preprocess_item(snake_case						) for i in features]

											if not isinstance(features[0]						,					snake_case						):
																	snake_case_			:Tuple						=      [vars(snake_case						) for f in features]
											snake_case_			:Any						=      {}

											snake_case_			:Dict						=      max(len(i["""input_nodes"""]						) for i in features						)
											snake_case_			:List[Any]						=      len(features[0]["""input_nodes"""][0]						)
											snake_case_			:Dict						=      len(features[0]["""attn_edge_type"""][0][0]						)
											snake_case_			:List[Any]						=      max(len(i["""input_edges"""][0][0]						) for i in features						)
											snake_case_			:Any						=      len(features[0]["""input_edges"""][0][0][0]						)
											snake_case_			:List[str]						=      len(snake_case						)

											snake_case_			:Union[str, Any]						=      torch.zeros(snake_case						,					max_node_num + 1						,					max_node_num + 1						,					dtype=torch.float						)
											snake_case_			:int						=      torch.zeros(snake_case						,					snake_case						,					snake_case						,					snake_case						,					dtype=torch.long						)
											snake_case_			:Optional[Any]						=      torch.zeros(snake_case						,					snake_case						,					snake_case						,					dtype=torch.long						)
											snake_case_			:Tuple						=      torch.zeros(snake_case						,					snake_case						,					dtype=torch.long						)
											snake_case_			:Dict						=      torch.zeros(snake_case						,					snake_case						,					snake_case						,					dtype=torch.long						)
											snake_case_			:Union[str, Any]						=      torch.zeros(
											    snake_case						,					snake_case						,					snake_case						,					snake_case						,					snake_case						,					dtype=torch.long						)

											for ix, f in enumerate(snake_case						):
																	for k in ["attn_bias", "attn_edge_type", "spatial_pos", "in_degree", "input_nodes", "input_edges"]:
																							snake_case_			:List[str]						=      torch.tensor(f[k]						)

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

																	snake_case_			:Optional[Any]						=      f["""attn_bias"""]
																	snake_case_			:int						=      f[
																	    """attn_edge_type"""
																	]
																	snake_case_			:Dict						=      f["""spatial_pos"""]
																	snake_case_			:Tuple						=      f["""in_degree"""]
																	snake_case_			:List[Any]						=      f["""input_nodes"""]
																	snake_case_			:List[Any]						=      f["""input_edges"""]

											snake_case_			:int						=      batch["""in_degree"""]

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

											return batch
