code: 
  
  
  
  
  
  '''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 (     UpperCAmelCase       ,			UpperCAmelCase = 512    )      ->Union[str, Any]:
  
  
  
  	"""simple docstring"""
  
  
  
  
  
  
  
  	A							  =      x.shape[1] if len(x.shape    ) > 1 else 1
  	A							  =      1 + np.arange(0       ,			feature_num * offset       ,			lowercase_       ,			dtype=np.intaa    )
  	A							  =      x + feature_offset
  	return x
  
  
  
  
  def    __a (     UpperCAmelCase       ,			UpperCAmelCase=True    )      ->Tuple:
  
  
  
  	"""simple docstring"""
  
  
  
  
  
  
  
  	requires_backends(lowercase_       ,			["""cython"""]    )
  
  	if keep_features and "edge_attr" in item.keys():  # edge_attr
  		A							  =      np.asarray(item["""edge_attr"""]       ,			dtype=np.intaa    )
  	else:
  		A							  =      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
  		A							  =      np.asarray(item["""node_feat"""]       ,			dtype=np.intaa    )
  	else:
  		A							  =      np.ones((item["""num_nodes"""], 1)       ,			dtype=np.intaa    )  # same embedding for all
  
  	A							  =      np.asarray(item["""edge_index"""]       ,			dtype=np.intaa    )
  
  	A							  =      convert_to_single_emb(lowercase_    ) + 1
  	A							  =      item['num_nodes']
  
  	if len(edge_attr.shape    ) == 1:
  		A							  =      edge_attr[:, None]
  	A							  =      np.zeros([num_nodes, num_nodes, edge_attr.shape[-1]]       ,			dtype=np.intaa    )
  	A							  =      convert_to_single_emb(lowercase_    ) + 1
  
  	# node adj matrix [num_nodes, num_nodes] bool
  	A							  =      np.zeros([num_nodes, num_nodes]       ,			dtype=lowercase_    )
  	A							  =      True
  
  	A							  =      algos_graphormer.floyd_warshall(lowercase_    )
  	A							  =      np.amax(lowercase_    )
  
  	A							  =      algos_graphormer.gen_edge_input(lowercase_       ,			lowercase_       ,			lowercase_    )
  	A							  =      np.zeros([num_nodes + 1, num_nodes + 1]       ,			dtype=np.single    )  # with graph token
  
  	# combine
  	A							  =      input_nodes + 1  # we shift all indices by one for padding
  	A							  =      attn_bias
  	A							  =      attn_edge_type
  	A							  =      shortest_path_result.astype(np.intaa    ) + 1  # we shift all indices by one for padding
  	A							  =      np.sum(lowercase_       ,			axis=1    ).reshape(-1    ) + 1  # we shift all indices by one for padding
  	A							  =      item['in_degree']  # for undirected graph
  	A							  =      input_edges + 1  # we shift all indices by one for padding
  	if "labels" not in item:
  		A							  =      item['y']
  
  	return item
  
  
  
  class 							__UpperCAmelCase  :
  
  
  
  
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		def __init__(self						:      str      ,      _lowerCAmelCase						:      Optional[int]=20      ,      _lowerCAmelCase						:      Tuple=False      ):
  			if not is_cython_available():
  				raise ImportError("""Graphormer preprocessing needs Cython (pyximport)"""      )
  
  			A							  =      spatial_pos_max
  			A							  =      on_the_fly_processing
  
  
  
  
  
  
  
  		def __call__(self						:      Optional[int]      ,      _lowerCAmelCase						:      List[dict]      ):
  			if self.on_the_fly_processing:
  				A							  =      [preprocess_item(_lowerCAmelCase      ) for i in features]
  
  			if not isinstance(features[0]      ,      _lowerCAmelCase      ):
  				A							  =      [vars(_lowerCAmelCase      ) for f in features]
  			A							  =      {}
  
  			A							  =      max(len(i["""input_nodes"""]      ) for i in features      )
  			A							  =      len(features[0]["""input_nodes"""][0]      )
  			A							  =      len(features[0]["""attn_edge_type"""][0][0]      )
  			A							  =      max(len(i["""input_edges"""][0][0]      ) for i in features      )
  			A							  =      len(features[0]["""input_edges"""][0][0][0]      )
  			A							  =      len(_lowerCAmelCase      )
  
  			A							  =      torch.zeros(_lowerCAmelCase      ,      max_node_num + 1      ,      max_node_num + 1      ,      dtype=torch.float      )
  			A							  =      torch.zeros(_lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase      ,      dtype=torch.long      )
  			A							  =      torch.zeros(_lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase      ,      dtype=torch.long      )
  			A							  =      torch.zeros(_lowerCAmelCase      ,      _lowerCAmelCase      ,      dtype=torch.long      )
  			A							  =      torch.zeros(_lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase      ,      dtype=torch.long      )
  			A							  =      torch.zeros(
  			    _lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase      ,      dtype=torch.long      )
  
  			for ix, f in enumerate(_lowerCAmelCase      ):
  				for k in ["attn_bias", "attn_edge_type", "spatial_pos", "in_degree", "input_nodes", "input_edges"]:
  					A							  =      torch.tensor(f[k]      )
  
  				if len(f["""attn_bias"""][1:, 1:][f["""spatial_pos"""] >= self.spatial_pos_max]      ) > 0:
  					A							  =      float("""-inf"""      )
  
  				A							  =      f['attn_bias']
  				A							  =      f[
  				    'attn_edge_type'
  				]
  				A							  =      f['spatial_pos']
  				A							  =      f['in_degree']
  				A							  =      f['input_nodes']
  				A							  =      f['input_edges']
  
  			A							  =      batch['in_degree']
  
  			A							  =      features[0]['labels']
  			if len(_lowerCAmelCase      ) == 1:  # one task
  				if isinstance(sample[0]      ,      _lowerCAmelCase      ):  # regression
  					A							  =      torch.from_numpy(np.concatenate([i["""labels"""] for i in features]      )      )
  				else:  # binary classification
  					A							  =      torch.from_numpy(np.concatenate([i["""labels"""] for i in features]      )      )
  			else:  # multi task classification, left to float to keep the NaNs
  				A							  =      torch.from_numpy(np.stack([i["""labels"""] for i in features]      ,      axis=0      )      )
  
  			return batch
  
  
style_context: 
  
  
  
  
  
  '''simple docstring'''
  
  
  import tempfile
  import unittest
  
  import numpy as np
  import torch
  
  from diffusers import DiffusionPipeline
  from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor
  
  
  
  class 							__UpperCAmelCase  (       unittest.TestCase						):
  
  
  
  
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		def        A     (self						:      Optional[int]      ,      _lowerCAmelCase						:      bool = False      ):
  			A							  =      10
  
  			if only_cross_attention:
  				A							  =      12
  			else:
  				# when only cross attention is not set, the cross attention dim must be the same as the query dim
  				A							  =      query_dim
  
  			return {
  			    "query_dim": query_dim,
  			    "cross_attention_dim": cross_attention_dim,
  			    "heads": 2,
  			    "dim_head": 4,
  			    "added_kv_proj_dim": 6,
  			    "norm_num_groups": 1,
  			    "only_cross_attention": only_cross_attention,
  			    "processor": AttnAddedKVProcessor(),
  			}
  
  
  
  
  
  
  
  		def        A     (self						:      Optional[int]      ,      _lowerCAmelCase						:      Dict      ,      _lowerCAmelCase						:      Any      ):
  			A							  =      2
  
  			A							  =      torch.rand(_lowerCAmelCase      ,      _lowerCAmelCase      ,      3      ,      2      )
  			A							  =      torch.rand(_lowerCAmelCase      ,      4      ,      _lowerCAmelCase      )
  			A							  =      None
  
  			return {
  			    "hidden_states": hidden_states,
  			    "encoder_hidden_states": encoder_hidden_states,
  			    "attention_mask": attention_mask,
  			}
  
  
  
  
  
  
  
  		def        A     (self						:      Optional[Any]      ):
  			# self and cross attention
  
  			torch.manual_seed(0      )
  
  			A							  =      self.get_constructor_arguments(only_cross_attention=_lowerCAmelCase      )
  			A							  =      Attention(**_lowerCAmelCase      )
  
  			self.assertTrue(attn.to_k is not None      )
  			self.assertTrue(attn.to_v is not None      )
  
  			A							  =      self.get_forward_arguments(
  			    query_dim=constructor_args["""query_dim"""]      ,      added_kv_proj_dim=constructor_args["""added_kv_proj_dim"""]      )
  
  			A							  =      attn(**_lowerCAmelCase      )
  
  			# only self attention
  
  			torch.manual_seed(0      )
  
  			A							  =      self.get_constructor_arguments(only_cross_attention=_lowerCAmelCase      )
  			A							  =      Attention(**_lowerCAmelCase      )
  
  			self.assertTrue(attn.to_k is None      )
  			self.assertTrue(attn.to_v is None      )
  
  			A							  =      self.get_forward_arguments(
  			    query_dim=constructor_args["""query_dim"""]      ,      added_kv_proj_dim=constructor_args["""added_kv_proj_dim"""]      )
  
  			A							  =      attn(**_lowerCAmelCase      )
  
  			self.assertTrue((only_cross_attn_out != self_and_cross_attn_out).all()      )
  
  
  
  
  
  
  class 							__UpperCAmelCase  (       unittest.TestCase						):
  
  
  
  
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		def        A     (self						:      Any      ):
  			A							  =      DiffusionPipeline.from_pretrained("""hf-internal-testing/tiny-stable-diffusion-pipe"""      ,      safety_checker=_lowerCAmelCase      )
  
  			A							  =      pipe(
  			    """foo"""      ,      num_inference_steps=2      ,      generator=torch.Generator("""cpu"""      ).manual_seed(0      )      ,      output_type="""np"""      ,      ).images
  
  			# the initial conversion succeeds
  			A							  =      DiffusionPipeline.from_pretrained(
  			    """hf-internal-testing/tiny-stable-diffusion-pipe"""      ,      device_map="""sequential"""      ,      safety_checker=_lowerCAmelCase      )
  
  			A							  =      pipe(
  			    """foo"""      ,      num_inference_steps=2      ,      generator=torch.Generator("""cpu"""      ).manual_seed(0      )      ,      output_type="""np"""      ,      ).images
  
  			with tempfile.TemporaryDirectory() as tmpdir:
  				# save the converted model
  				pipe.save_pretrained(_lowerCAmelCase      )
  
  				# can also load the converted weights
  				A							  =      DiffusionPipeline.from_pretrained(_lowerCAmelCase      ,      device_map="""sequential"""      ,      safety_checker=_lowerCAmelCase      )
  
  			A							  =      pipe(
  			    """foo"""      ,      num_inference_steps=2      ,      generator=torch.Generator("""cpu"""      ).manual_seed(0      )      ,      output_type="""np"""      ,      ).images
  
  			self.assertTrue(np.allclose(_lowerCAmelCase      ,      _lowerCAmelCase      ,      atol=1e-5      )      )
  			self.assertTrue(np.allclose(_lowerCAmelCase      ,      _lowerCAmelCase      ,      atol=1e-5      )      )
  
  