code: 
  
  from functools import partial
  from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload
  
  import torch
  import torch.nn as nn
  import torch.types
  
  def 			snake_case__    (					SCREAMING_SNAKE_CASE_	:					torch.Tensor			,  SCREAMING_SNAKE_CASE_	:					torch.Tensor			,  SCREAMING_SNAKE_CASE_	:					bool    ):
  
  
  
       '''simple docstring'''
  
  
  
  
  
       if not inplace:
            lowercase__					:     str   = ma + ma
       else:
            ma += ma
  
       return ma
  
  def 			snake_case__    (					SCREAMING_SNAKE_CASE_	:					torch.Tensor			,  SCREAMING_SNAKE_CASE_	:					List[int]    ):
  
  
  
       '''simple docstring'''
  
  
  
  
  
       lowercase__					:     str   = -1 * len(SCREAMING_SNAKE_CASE_    )
       lowercase__					:     int   = list(range(len(tensor.shape[:zero_index]    )    )    )
       return tensor.permute(first_inds + [zero_index + i for i in inds]    )
  
  def 			snake_case__    (					SCREAMING_SNAKE_CASE_	:					torch.Tensor			,  SCREAMING_SNAKE_CASE_	:					int    ):
  
  
  
       '''simple docstring'''
  
  
  
  
  
       return t.reshape(t.shape[:-no_dims] + (-1,)    )
  
  def 			snake_case__    (					SCREAMING_SNAKE_CASE_	:					torch.Tensor			,  SCREAMING_SNAKE_CASE_	:					torch.Tensor			,  SCREAMING_SNAKE_CASE_	:					int			,  SCREAMING_SNAKE_CASE_	:					float = 1E-4    ):
  
  
  
       '''simple docstring'''
  
  
  
  
  
       lowercase__					:     List[Any]   = mask.expand(*value.shape    )
       return torch.sum(mask * value			,  dim=SCREAMING_SNAKE_CASE_    ) / (eps + torch.sum(SCREAMING_SNAKE_CASE_			,  dim=SCREAMING_SNAKE_CASE_    ))
  
  def 			snake_case__    (					SCREAMING_SNAKE_CASE_	:					torch.Tensor			,  SCREAMING_SNAKE_CASE_	:					torch.types.Number = 2.3125			,  SCREAMING_SNAKE_CASE_	:					torch.types.Number = 21.6875			,  SCREAMING_SNAKE_CASE_	:					int = 64    ):
  
  
  
       '''simple docstring'''
  
  
  
  
  
       lowercase__					:     Dict   = torch.linspace(SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_			,  no_bins - 1			,  device=pts.device    )
       lowercase__					:     str   = torch.sqrt(torch.sum((pts.unsqueeze(-2    ) - pts.unsqueeze(-3    )) ** 2			,  dim=-1    )    )
       return torch.bucketize(SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_    )
  
  def 			snake_case__    (					SCREAMING_SNAKE_CASE_	:					Callable[[list], Any]			,  SCREAMING_SNAKE_CASE_	:					List[dict]    ):
  
  
  
       '''simple docstring'''
  
  
  
  
  
       lowercase__					:     Tuple   = dicts[0]
       lowercase__					:     str   = {}
       for k, v in first.items():
            lowercase__					:     List[Any]   = [d[k] for d in dicts]
            if isinstance(SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_    ):
                 lowercase__					:     Optional[Any]   = dict_multimap(SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_    )
            else:
                 lowercase__					:     Union[str, Any]   = fn(SCREAMING_SNAKE_CASE_    )
  
       return new_dict
  
  def 			snake_case__    (					SCREAMING_SNAKE_CASE_	:					torch.Tensor			,  SCREAMING_SNAKE_CASE_	:					torch.Tensor    ):
  
  
  
       '''simple docstring'''
  
  
  
  
  
       lowercase__					:     Optional[Any]   = v_bins.view(((1,) * len(x.shape    )) + (len(SCREAMING_SNAKE_CASE_    ),)    )
       lowercase__					:     Tuple   = x[..., None] - reshaped_bins
       lowercase__					:     Any   = torch.argmin(torch.abs(SCREAMING_SNAKE_CASE_    )			,  dim=-1    )
       return nn.functional.one_hot(SCREAMING_SNAKE_CASE_			,  num_classes=len(SCREAMING_SNAKE_CASE_    )    ).float()
  
  def 			snake_case__    (					SCREAMING_SNAKE_CASE_	:					torch.Tensor			,  SCREAMING_SNAKE_CASE_	:					torch.Tensor			,  SCREAMING_SNAKE_CASE_	:					int = 0			,  SCREAMING_SNAKE_CASE_	:					int = 0    ):
  
  
  
       '''simple docstring'''
  
  
  
  
  
       lowercase__					:     List[Union[slice, torch.Tensor]]   = []
       for i, s in enumerate(data.shape[:no_batch_dims]    ):
            lowercase__					:     int   = torch.arange(SCREAMING_SNAKE_CASE_    )
            lowercase__					:     Dict   = r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape    ) - i - 1)))    )
            ranges.append(SCREAMING_SNAKE_CASE_    )
  
       lowercase__					:     List[Union[slice, torch.Tensor]]   = [slice(SCREAMING_SNAKE_CASE_    ) for _ in range(len(data.shape    ) - no_batch_dims    )]
       lowercase__					:     Optional[Any]   = inds
       ranges.extend(SCREAMING_SNAKE_CASE_    )
       # Matt note: Editing this to get around the behaviour of using a list as an array index changing
       # in recent Numpy versions
       return data[tuple(SCREAMING_SNAKE_CASE_    )]
  
  
  snake_case_    						=				TypeVar('''T''')
  
  def 			snake_case__    (					SCREAMING_SNAKE_CASE_	:					Callable[[T], Any]			,  SCREAMING_SNAKE_CASE_	:					Dict[Any, Union[dict, list, tuple, T]]			,  SCREAMING_SNAKE_CASE_	:					Type[T]    ):
  
  
  
       '''simple docstring'''
  
  
  
  
  
       lowercase__					:     Dict[Any, Union[dict, list, tuple, Any]]   = {}
       for k, v in dic.items():
            if isinstance(SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_    ):
                 lowercase__					:     List[Any]   = dict_map(SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_    )
            else:
                 lowercase__					:     List[Any]   = tree_map(SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_    )
  
       return new_dict
  
  @overload
  def 			snake_case__    (					SCREAMING_SNAKE_CASE_	:					Callable[[T], Any]			,  SCREAMING_SNAKE_CASE_	:					T			,  SCREAMING_SNAKE_CASE_	:					Type[T]    ):
  
  
  
       '''simple docstring'''
  
  
  
  
  
       ...
  
  @overload
  def 			snake_case__    (					SCREAMING_SNAKE_CASE_	:					Callable[[T], Any]			,  SCREAMING_SNAKE_CASE_	:					dict			,  SCREAMING_SNAKE_CASE_	:					Type[T]    ):
  
  
  
       '''simple docstring'''
  
  
  
  
  
       ...
  
  @overload
  def 			snake_case__    (					SCREAMING_SNAKE_CASE_	:					Callable[[T], Any]			,  SCREAMING_SNAKE_CASE_	:					list			,  SCREAMING_SNAKE_CASE_	:					Type[T]    ):
  
  
  
       '''simple docstring'''
  
  
  
  
  
       ...
  
  @overload
  def 			snake_case__    (					SCREAMING_SNAKE_CASE_	:					Callable[[T], Any]			,  SCREAMING_SNAKE_CASE_	:					tuple			,  SCREAMING_SNAKE_CASE_	:					Type[T]    ):
  
  
  
       '''simple docstring'''
  
  
  
  
  
       ...
  
  def 			snake_case__    (					SCREAMING_SNAKE_CASE_	:					Any			,  SCREAMING_SNAKE_CASE_	:					int			,  SCREAMING_SNAKE_CASE_	:					str    ):
  
  
  
       '''simple docstring'''
  
  
  
  
  
       if isinstance(SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_    ):
            return dict_map(SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_    )
       elif isinstance(SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_    ):
            return [tree_map(SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_    ) for x in tree]
       elif isinstance(SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_    ):
            return tuple(tree_map(SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_    ) for x in tree    )
       elif isinstance(SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_    ):
            return fn(SCREAMING_SNAKE_CASE_    )
       else:
            print(type(SCREAMING_SNAKE_CASE_    )    )
            raise ValueError('Not supported'    )
  
  snake_case_    						=				partial(tree_map, leaf_type=torch.Tensor)
  
  
style_context: 
  
  from dataclasses import dataclass
  from typing import List, Optional, Union
  
  import numpy as np
  import PIL
  
  from ...utils import BaseOutput, is_invisible_watermark_available, is_torch_available, is_transformers_available
  
  
  @dataclass
  class    SCREAMING_SNAKE_CASE__      (__snake_case  ):
      __lowerCamelCase		: Union[List[PIL.Image.Image], np.ndarray]
  
  
  if is_transformers_available() and is_torch_available() and is_invisible_watermark_available():
         from .pipeline_stable_diffusion_xl import StableDiffusionXLPipeline
         from .pipeline_stable_diffusion_xl_imgaimg import StableDiffusionXLImgaImgPipeline
  
  