





"""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 lowerCamelCase__  (    __UpperCamelCase  ,				__UpperCamelCase  ,				unittest.TestCase				):
   __UpperCAmelCase        =					StableDiffusionXLPipeline
   __UpperCAmelCase        =					TEXT_TO_IMAGE_PARAMS
   __UpperCAmelCase        =					TEXT_TO_IMAGE_BATCH_PARAMS
   __UpperCAmelCase        =					TEXT_TO_IMAGE_IMAGE_PARAMS
   __UpperCAmelCase        =					TEXT_TO_IMAGE_IMAGE_PARAMS



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







      """simple docstring"""
      torch.manual_seed(0	)
      lowercase    :			int  =       UNetaDConditionModel(
          block_out_channels=(3_2, 6_4)      ,   layers_per_block=2      ,   sample_size=3_2      ,   in_channels=4      ,   out_channels=4      ,   down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")      ,   up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")      ,   attention_head_dim=(2, 4)      ,   use_linear_projection=snake_case      ,   addition_embed_type="""text_time"""      ,   addition_time_embed_dim=8      ,   transformer_layers_per_block=(1, 2)      ,   projection_class_embeddings_input_dim=8_0      ,   cross_attention_dim=6_4      ,   )
      lowercase    :			Union[str, Any]  =       EulerDiscreteScheduler(
          beta_start=0.0_00_85      ,   beta_end=0.0_12      ,   steps_offset=1      ,   beta_schedule="""scaled_linear"""      ,   timestep_spacing="""leading"""      ,   )
      torch.manual_seed(0	)
      lowercase    :			Union[str, Any]  =       AutoencoderKL(
          block_out_channels=[3_2, 6_4]      ,   in_channels=3      ,   out_channels=3      ,   down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]      ,   up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]      ,   latent_channels=4      ,   sample_size=1_2_8      ,   )
      torch.manual_seed(0	)
      lowercase    :			Union[str, Any]  =       CLIPTextConfig(
          bos_token_id=0      ,   eos_token_id=2      ,   hidden_size=3_2      ,   intermediate_size=3_7      ,   layer_norm_eps=1E-05      ,   num_attention_heads=4      ,   num_hidden_layers=5      ,   pad_token_id=1      ,   vocab_size=1_0_0_0      ,   hidden_act="""gelu"""      ,   projection_dim=3_2      ,   )
      lowercase    :			Optional[Any]  =       CLIPTextModel(snake_case	)
      lowercase    :			Any  =       CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""      ,   local_files_only=snake_case	)

      lowercase    :			Optional[Any]  =       CLIPTextModelWithProjection(snake_case	)
      lowercase    :			Dict  =       CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""      ,   local_files_only=snake_case	)

      lowercase    :			Optional[Any]  =       {
          """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    _UpperCAmelCase	(       self      ,   snake_case      ,   snake_case=0	) ->    List[Any]:







      """simple docstring"""
      if str(snake_case	).startswith("""mps"""	):
         lowercase    :			Any  =       torch.manual_seed(snake_case	)
      else:
         lowercase    :			Dict  =       torch.Generator(device=snake_case	).manual_seed(snake_case	)
      lowercase    :			List[str]  =       {
          """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    _UpperCAmelCase	(       self	) ->    Tuple:







      """simple docstring"""
      lowercase    :			int  =       """cpu"""  # ensure determinism for the device-dependent torch.Generator
      lowercase    :			List[Any]  =       self.get_dummy_components()
      lowercase    :			List[str]  =       StableDiffusionXLPipeline(**snake_case	)
      lowercase    :			List[Any]  =       sd_pipe.to(snake_case	)
      sd_pipe.set_progress_bar_config(disable=snake_case	)

      lowercase    :			Tuple  =       self.get_dummy_inputs(snake_case	)
      lowercase    :			Optional[Any]  =       sd_pipe(**snake_case	).images
      lowercase    :			Optional[int]  =       image[0, -3:, -3:, -1]

      assert image.shape == (1, 6_4, 6_4, 3)
      lowercase    :			Any  =       np.array([0.58_73, 0.61_28, 0.47_97, 0.51_22, 0.56_74, 0.46_39, 0.52_27, 0.51_49, 0.47_47]	)

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



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







      """simple docstring"""
      lowercase    :			List[Any]  =       self.get_dummy_components()
      lowercase    :			int  =       StableDiffusionXLPipeline(**snake_case	)
      lowercase    :			Union[str, Any]  =       sd_pipe.to(snake_case	)
      lowercase    :			Any  =       sd_pipe.to(snake_case	)
      sd_pipe.set_progress_bar_config(disable=snake_case	)

      # forward without prompt embeds
      lowercase    :			Dict  =       self.get_dummy_inputs(snake_case	)
      lowercase    :			Dict  =       3 * ["""this is a negative prompt"""]
      lowercase    :			Any  =       negative_prompt
      lowercase    :			Any  =       3 * [inputs["""prompt"""]]

      lowercase    :			List[Any]  =       sd_pipe(**snake_case	)
      lowercase    :			Union[str, Any]  =       output.images[0, -3:, -3:, -1]

      # forward with prompt embeds
      lowercase    :			Optional[Any]  =       self.get_dummy_inputs(snake_case	)
      lowercase    :			Tuple  =       3 * ["""this is a negative prompt"""]
      lowercase    :			Optional[int]  =       3 * [inputs.pop("""prompt"""	)]

      (
          (
          lowercase
      )  ,					(
          lowercase
      )  ,					(
          lowercase
      )  ,					(
          lowercase
      )  ,					
      )    :			List[Any]  =       sd_pipe.encode_prompt(snake_case      ,   negative_prompt=snake_case	)

      lowercase    :			List[str]  =       sd_pipe(
          **snake_case      ,   prompt_embeds=snake_case      ,   negative_prompt_embeds=snake_case      ,   pooled_prompt_embeds=snake_case      ,   negative_pooled_prompt_embeds=snake_case      ,   )
      lowercase    :			Union[str, Any]  =       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    _UpperCAmelCase	(       self	) ->    Union[str, Any]:







      """simple docstring"""
      super().test_attention_slicing_forward_pass(expected_max_diff=3E-3	)





   def    _UpperCAmelCase	(       self	) ->    int:







      """simple docstring"""
      super().test_inference_batch_single_identical(expected_max_diff=3E-3	)




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



   def    _UpperCAmelCase	(       self	) ->    int:







      """simple docstring"""
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()



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







      """simple docstring"""
      lowercase    :			Any  =       torch.Generator(device=snake_case	).manual_seed(snake_case	)
      lowercase    :			Tuple  =       np.random.RandomState(snake_case	).standard_normal((1, 4, 6_4, 6_4)	)
      lowercase    :			List[Any]  =       torch.from_numpy(snake_case	).to(device=snake_case      ,   dtype=snake_case	)
      lowercase    :			Union[str, Any]  =       {
          """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    _UpperCAmelCase	(       self	) ->    Optional[int]:







      """simple docstring"""
      lowercase    :			Dict  =       DiffusionPipeline.from_pretrained("""stabilityai/stable-diffusion-2-base"""	)
      pipe.to(snake_case	)
      pipe.set_progress_bar_config(disable=snake_case	)

      lowercase    :			Tuple  =       self.get_inputs(snake_case	)
      lowercase    :			Any  =       pipe(**snake_case	).images
      lowercase    :			int  =       image[0, -3:, -3:, -1].flatten()

      assert image.shape == (1, 5_1_2, 5_1_2, 3)
      lowercase    :			Dict  =       np.array([0.4_94_93, 0.4_78_96, 0.4_07_98, 0.5_42_14, 0.5_32_12, 0.4_82_02, 0.4_76_56, 0.4_63_29, 0.4_85_06]	)
      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 			__snake_case						(      __A ,__A = 512    )		->				Optional[Any]:
   lowercase    :			Optional[Any]  =       x.shape[1] if len(x.shape    ) > 1 else 1
   lowercase    :			Tuple  =       1 + np.arange(0 ,feature_num * offset ,__A ,dtype=np.intaa    )
   lowercase    :			Any  =       x + feature_offset
   return x




def 			__snake_case						(      __A ,__A=True    )		->				List[Any]:
   requires_backends(__A ,["""cython"""]    )

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

   lowercase    :			str  =       np.asarray(item["""edge_index"""] ,dtype=np.intaa    )

   lowercase    :			Tuple  =       convert_to_single_emb(__A    ) + 1
   lowercase    :			Dict  =       item["""num_nodes"""]

   if len(edge_attr.shape    ) == 1:
      lowercase    :			Dict  =       edge_attr[:, None]
   lowercase    :			Tuple  =       np.zeros([num_nodes, num_nodes, edge_attr.shape[-1]] ,dtype=np.intaa    )
   lowercase    :			List[Any]  =       convert_to_single_emb(__A    ) + 1

   # node adj matrix [num_nodes, num_nodes] bool
   lowercase    :			List[Any]  =       np.zeros([num_nodes, num_nodes] ,dtype=__A    )
   lowercase    :			str  =       True

   lowercase  ,					lowercase    :			List[str]  =       algos_graphormer.floyd_warshall(__A    )
   lowercase    :			str  =       np.amax(__A    )

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

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

   return item







class lowerCamelCase__  :



   def __init__(       self      ,   snake_case=2_0      ,   snake_case=False	) ->    List[Any]:







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

      lowercase    :			Any  =       spatial_pos_max
      lowercase    :			Dict  =       on_the_fly_processing





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







      """simple docstring"""
      if self.on_the_fly_processing:
         lowercase    :			Dict  =       [preprocess_item(snake_case	) for i in features]

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

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

      lowercase    :			List[str]  =       torch.zeros(snake_case      ,   max_node_num + 1      ,   max_node_num + 1      ,   dtype=torch.float	)
      lowercase    :			List[Any]  =       torch.zeros(snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   dtype=torch.long	)
      lowercase    :			int  =       torch.zeros(snake_case      ,   snake_case      ,   snake_case      ,   dtype=torch.long	)
      lowercase    :			List[str]  =       torch.zeros(snake_case      ,   snake_case      ,   dtype=torch.long	)
      lowercase    :			Tuple  =       torch.zeros(snake_case      ,   snake_case      ,   snake_case      ,   dtype=torch.long	)
      lowercase    :			Dict  =       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"]:
            lowercase    :			Optional[int]  =       torch.tensor(f[k]	)

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

         lowercase    :			str  =       f["""attn_bias"""]
         lowercase    :			Union[str, Any]  =       f[
             """attn_edge_type"""
         ]
         lowercase    :			Union[str, Any]  =       f["""spatial_pos"""]
         lowercase    :			int  =       f["""in_degree"""]
         lowercase    :			int  =       f["""input_nodes"""]
         lowercase    :			Any  =       f["""input_edges"""]

      lowercase    :			Dict  =       batch["""in_degree"""]

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

      return batch


