

from ...configuration_utils import PretrainedConfig
from ...utils import logging


_SCREAMING_SNAKE_CASE  :						Any         =    logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE  :						Optional[Any]         =    {
    '''google/switch-base-8''': '''https://huggingface.co/google/switch-base-8/blob/main/config.json''',
}






class  UpperCAmelCase__	( A__ ):






  """simple docstring"""


  a     =      "switch_transformers"
  a     =      ["past_key_values"]
  a     =      {"hidden_size": "d_model", "num_attention_heads": "num_heads", "num_hidden_layers": "num_layers"}






  def __init__(							self		:	str			,     __lowerCamelCase		:	Any=3_2128			,     __lowerCamelCase		:	Any=768			,     __lowerCamelCase		:	str=64			,     __lowerCamelCase		:	Optional[Any]=2048			,     __lowerCamelCase		:	List[str]=64			,     __lowerCamelCase		:	Any=12			,     __lowerCamelCase		:	Optional[Any]=3			,     __lowerCamelCase		:	Optional[int]=12			,     __lowerCamelCase		:	Dict=3			,     __lowerCamelCase		:	List[str]=12			,     __lowerCamelCase		:	int=8			,     __lowerCamelCase		:	List[str]=False			,     __lowerCamelCase		:	Tuple=0.01			,     __lowerCamelCase		:	Union[str, Any]="float32"			,     __lowerCamelCase		:	Optional[int]=False			,     __lowerCamelCase		:	int=32			,     __lowerCamelCase		:	Any=128			,     __lowerCamelCase		:	int=0.1			,     __lowerCamelCase		:	List[str]=1e-6			,     __lowerCamelCase		:	Dict=0.001			,     __lowerCamelCase		:	str=0.001			,     __lowerCamelCase		:	Tuple=1.0			,     __lowerCamelCase		:	Tuple="relu"			,     __lowerCamelCase		:	Optional[int]=True			,     __lowerCamelCase		:	str=False			,     __lowerCamelCase		:	int=True			,     __lowerCamelCase		:	List[str]=0			,     __lowerCamelCase		:	Optional[Any]=1			,     **__lowerCamelCase		:	Dict			,     )	->							List[str]:
       SCREAMING_SNAKE_CASE__					=					vocab_size
       SCREAMING_SNAKE_CASE__					=					d_model
       SCREAMING_SNAKE_CASE__					=					d_kv
       SCREAMING_SNAKE_CASE__					=					d_ff

       SCREAMING_SNAKE_CASE__					=					num_sparse_encoder_layers

       SCREAMING_SNAKE_CASE__					=					num_layers
       SCREAMING_SNAKE_CASE__					=					(
           num_decoder_layers if num_decoder_layers is not None else self.num_layers
       )  # default = symmetry
       SCREAMING_SNAKE_CASE__					=					num_sparse_decoder_layers

       # This tells us, each how many encoder layer we'll have to set a sparse layer.
       if self.num_sparse_encoder_layers > 0:
            SCREAMING_SNAKE_CASE__					=					self.num_layers // self.num_sparse_encoder_layers
       else:
            SCREAMING_SNAKE_CASE__					=					self.num_layers  # HACK: this will create 0 sparse layers

       # This tells us, each how many encoder layer we'll have to set a sparse layer.
       if self.num_sparse_decoder_layers > 0:
            SCREAMING_SNAKE_CASE__					=					self.num_decoder_layers // self.num_sparse_decoder_layers
       else:
            SCREAMING_SNAKE_CASE__					=					self.num_decoder_layers  # HACK: this will create 0 sparse layers

       SCREAMING_SNAKE_CASE__					=					num_heads
       SCREAMING_SNAKE_CASE__					=					num_experts
       SCREAMING_SNAKE_CASE__					=					expert_capacity
       SCREAMING_SNAKE_CASE__					=					router_bias
       SCREAMING_SNAKE_CASE__					=					router_jitter_noise
       if router_dtype not in ["float32", "float16", "bfloat16"]:
            raise ValueError(f'''`router_dtype` must be one of \'float32\', \'float16\' or \'bfloat16\', got {router_dtype}'''      )
       SCREAMING_SNAKE_CASE__					=					router_dtype

       SCREAMING_SNAKE_CASE__					=					router_ignore_padding_tokens
       SCREAMING_SNAKE_CASE__					=					relative_attention_num_buckets
       SCREAMING_SNAKE_CASE__					=					relative_attention_max_distance

       SCREAMING_SNAKE_CASE__					=					dropout_rate
       SCREAMING_SNAKE_CASE__					=					layer_norm_epsilon
       SCREAMING_SNAKE_CASE__					=					initializer_factor
       SCREAMING_SNAKE_CASE__					=					feed_forward_proj
       SCREAMING_SNAKE_CASE__					=					use_cache
       SCREAMING_SNAKE_CASE__					=					add_router_probs

       SCREAMING_SNAKE_CASE__					=					router_z_loss_coef
       SCREAMING_SNAKE_CASE__					=					router_aux_loss_coef

       SCREAMING_SNAKE_CASE__					=					self.feed_forward_proj.split('''-'''      )
       SCREAMING_SNAKE_CASE__					=					act_info[-1]
       SCREAMING_SNAKE_CASE__					=					act_info[0] == '''gated'''

       if len(__lowerCamelCase      ) > 1 and act_info[0] != "gated" or len(__lowerCamelCase      ) > 2:
            raise ValueError(
                f'''`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer.'''
                '''Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. '''
                '''\'gated-gelu\' or \'relu\''''      )

       # for backwards compatibility
       if feed_forward_proj == "gated-gelu":
            SCREAMING_SNAKE_CASE__					=					'''gelu_new'''

       super().__init__(
           pad_token_id=__lowerCamelCase			,     eos_token_id=__lowerCamelCase			,     is_encoder_decoder=__lowerCamelCase			,     **__lowerCamelCase			,     )





import math
from typing import Callable, List, Optional, Union

import numpy as np
import PIL
import torch
from PIL import Image
from transformers import CLIPTextModel, CLIPTokenizer

from diffusers.models import AutoencoderKL, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale import StableDiffusionUpscalePipeline
from diffusers.schedulers import DDIMScheduler, DDPMScheduler, LMSDiscreteScheduler, PNDMScheduler
def        UpperCAmelCase_     (		_A   ,	_A   ,	_A=[]       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					size[0] - overlap_pixels * 2
     SCREAMING_SNAKE_CASE__					=					size[1] - overlap_pixels * 2
     for letter in ["l", "r"]:
          if letter in remove_borders:
               size_x += overlap_pixels
     for letter in ["t", "b"]:
          if letter in remove_borders:
               size_y += overlap_pixels
     SCREAMING_SNAKE_CASE__					=					np.ones((size_y, size_x)   ,	dtype=np.uinta       ) * 2_55
     SCREAMING_SNAKE_CASE__					=					np.pad(_A   ,	mode='''linear_ramp'''   ,	pad_width=_A   ,	end_values=0       )

     if "l" in remove_borders:
          SCREAMING_SNAKE_CASE__					=					mask[:, overlap_pixels : mask.shape[1]]
     if "r" in remove_borders:
          SCREAMING_SNAKE_CASE__					=					mask[:, 0 : mask.shape[1] - overlap_pixels]
     if "t" in remove_borders:
          SCREAMING_SNAKE_CASE__					=					mask[overlap_pixels : mask.shape[0], :]
     if "b" in remove_borders:
          SCREAMING_SNAKE_CASE__					=					mask[0 : mask.shape[0] - overlap_pixels, :]
     return mask
def        UpperCAmelCase_     (		_A   ,	_A   ,	_A       ):





     '''simple docstring'''




     return max(_A   ,	min(_A   ,	_A       )       )
def        UpperCAmelCase_     (		_A   ,	_A   ,	_A       ):





     '''simple docstring'''




     return (
         clamp(rect[0]   ,	min[0]   ,	max[0]       ),
         clamp(rect[1]   ,	min[1]   ,	max[1]       ),
         clamp(rect[2]   ,	min[0]   ,	max[0]       ),
         clamp(rect[3]   ,	min[1]   ,	max[1]       ),
     )
def        UpperCAmelCase_     (		_A   ,	_A   ,	_A       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					list(_A       )
     rect[0] -= overlap
     rect[1] -= overlap
     rect[2] += overlap
     rect[3] += overlap
     SCREAMING_SNAKE_CASE__					=					clamp_rect(_A   ,	[0, 0]   ,	[image_size[0], image_size[1]]       )
     return rect
def        UpperCAmelCase_     (		_A   ,	_A   ,	_A   ,	_A       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					Image.new('''RGB'''   ,	(tile.size[0] + original_slice, tile.size[1])       )
     result.paste(
         original_image.resize((tile.size[0], tile.size[1])   ,	Image.BICUBIC       ).crop(
             (slice_x, 0, slice_x + original_slice, tile.size[1])       )   ,	(0, 0)   ,	)
     result.paste(_A   ,	(original_slice, 0)       )
     return result
def        UpperCAmelCase_     (		_A   ,	_A       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					(original_image_slice * 4, 0, tile.size[0], tile.size[1])
     SCREAMING_SNAKE_CASE__					=					tile.crop(_A       )
     return tile
def        UpperCAmelCase_     (		_A   ,	_A       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					n % d
     return n - divisor






class  UpperCAmelCase__	( A__ ):






  """simple docstring"""


  def __init__(							self		:	int			,     __lowerCamelCase		:	AutoencoderKL			,     __lowerCamelCase		:	CLIPTextModel			,     __lowerCamelCase		:	CLIPTokenizer			,     __lowerCamelCase		:	UNetaDConditionModel			,     __lowerCamelCase		:	DDPMScheduler			,     __lowerCamelCase		:	Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]			,     __lowerCamelCase		:	int = 350			,     )	->							int:
       super().__init__(
           vae=__lowerCamelCase			,     text_encoder=__lowerCamelCase			,     tokenizer=__lowerCamelCase			,     unet=__lowerCamelCase			,     low_res_scheduler=__lowerCamelCase			,     scheduler=__lowerCamelCase			,     max_noise_level=__lowerCamelCase			,     )






  def    lowercase_	(							self		:	Tuple			,     __lowerCamelCase		:	Optional[int]			,     __lowerCamelCase		:	Tuple			,     __lowerCamelCase		:	Any			,     __lowerCamelCase		:	List[Any]			,     __lowerCamelCase		:	str			,     __lowerCamelCase		:	List[str]			,     __lowerCamelCase		:	str			,     **__lowerCamelCase		:	Tuple      )	->							Dict:
       torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					(
           min(image.size[0] - (tile_size + original_image_slice)			,     x * tile_size      ),
           min(image.size[1] - (tile_size + original_image_slice)			,     y * tile_size      ),
           min(image.size[0]			,     (x + 1) * tile_size      ),
           min(image.size[1]			,     (y + 1) * tile_size      ),
       )
       SCREAMING_SNAKE_CASE__					=					add_overlap_rect(__lowerCamelCase			,     __lowerCamelCase			,     image.size      )
       SCREAMING_SNAKE_CASE__					=					image.crop(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					((crop_rect[0] + ((crop_rect[2] - crop_rect[0]) / 2)) / image.size[0]) * tile.size[0]
       SCREAMING_SNAKE_CASE__					=					translated_slice_x - (original_image_slice / 2)
       SCREAMING_SNAKE_CASE__					=					max(0			,     __lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					squeeze_tile(__lowerCamelCase			,     __lowerCamelCase			,     __lowerCamelCase			,     __lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					to_input.size
       SCREAMING_SNAKE_CASE__					=					to_input.resize((tile_size, tile_size)			,     Image.BICUBIC      )
       SCREAMING_SNAKE_CASE__					=					super(__lowerCamelCase			,     self      ).__call__(image=__lowerCamelCase			,     **__lowerCamelCase      ).images[0]
       SCREAMING_SNAKE_CASE__					=					upscaled_tile.resize((orig_input_size[0] * 4, orig_input_size[1] * 4)			,     Image.BICUBIC      )
       SCREAMING_SNAKE_CASE__					=					unsqueeze_tile(__lowerCamelCase			,     __lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					upscaled_tile.resize((tile.size[0] * 4, tile.size[1] * 4)			,     Image.BICUBIC      )
       SCREAMING_SNAKE_CASE__					=					[]
       if x == 0:
            remove_borders.append('''l'''      )
       elif crop_rect[2] == image.size[0]:
            remove_borders.append('''r'''      )
       if y == 0:
            remove_borders.append('''t'''      )
       elif crop_rect[3] == image.size[1]:
            remove_borders.append('''b'''      )
       SCREAMING_SNAKE_CASE__					=					Image.fromarray(
           make_transparency_mask(
               (upscaled_tile.size[0], upscaled_tile.size[1])			,     tile_border * 4			,     remove_borders=__lowerCamelCase      )			,     mode='''L'''			,     )
       final_image.paste(
           __lowerCamelCase			,     (crop_rect_with_overlap[0] * 4, crop_rect_with_overlap[1] * 4)			,     __lowerCamelCase      )







  @torch.no_grad()
  def __call__(							self		:	List[str]			,     __lowerCamelCase		:	Union[str, List[str]]			,     __lowerCamelCase		:	Union[PIL.Image.Image, List[PIL.Image.Image]]			,     __lowerCamelCase		:	int = 75			,     __lowerCamelCase		:	float = 9.0			,     __lowerCamelCase		:	int = 50			,     __lowerCamelCase		:	Optional[Union[str, List[str]]] = None			,     __lowerCamelCase		:	Optional[int] = 1			,     __lowerCamelCase		:	float = 0.0			,     __lowerCamelCase		:	Optional[torch.Generator] = None			,     __lowerCamelCase		:	Optional[torch.FloatTensor] = None			,     __lowerCamelCase		:	Optional[Callable[[int, int, torch.FloatTensor], None]] = None			,     __lowerCamelCase		:	int = 1			,     __lowerCamelCase		:	int = 128			,     __lowerCamelCase		:	int = 32			,     __lowerCamelCase		:	int = 32			,     )	->							List[Any]:
       SCREAMING_SNAKE_CASE__					=					Image.new('''RGB'''			,     (image.size[0] * 4, image.size[1] * 4)      )
       SCREAMING_SNAKE_CASE__					=					math.ceil(image.size[0] / tile_size      )
       SCREAMING_SNAKE_CASE__					=					math.ceil(image.size[1] / tile_size      )
       SCREAMING_SNAKE_CASE__					=					tcx * tcy
       SCREAMING_SNAKE_CASE__					=					0
       for y in range(__lowerCamelCase      ):
            for x in range(__lowerCamelCase      ):
                 self._process_tile(
                     __lowerCamelCase			,     __lowerCamelCase			,     __lowerCamelCase			,     __lowerCamelCase			,     __lowerCamelCase			,     __lowerCamelCase			,     __lowerCamelCase			,     prompt=__lowerCamelCase			,     num_inference_steps=__lowerCamelCase			,     guidance_scale=__lowerCamelCase			,     noise_level=__lowerCamelCase			,     negative_prompt=__lowerCamelCase			,     num_images_per_prompt=__lowerCamelCase			,     eta=__lowerCamelCase			,     generator=__lowerCamelCase			,     latents=__lowerCamelCase			,     )
                 current_count += 1
                 if callback is not None:
                      callback({'''progress''': current_count / total_tile_count, '''image''': final_image}      )
       return final_image



def        UpperCAmelCase_     (		):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					'''stabilityai/stable-diffusion-x4-upscaler'''
     SCREAMING_SNAKE_CASE__					=					StableDiffusionTiledUpscalePipeline.from_pretrained(_A   ,	revision='''fp16'''   ,	torch_dtype=torch.floataa       )
     SCREAMING_SNAKE_CASE__					=					pipe.to('''cuda'''       )
     SCREAMING_SNAKE_CASE__					=					Image.open('''../../docs/source/imgs/diffusers_library.jpg'''       )

     def callback(_A       ):
          print(F'''progress: {obj["progress"]:.4f}'''       )
          obj["image"].save('''diffusers_library_progress.jpg'''       )

     SCREAMING_SNAKE_CASE__					=					pipe(image=_A   ,	prompt='''Black font, white background, vector'''   ,	noise_level=40   ,	callback=_A       )
     final_image.save('''diffusers_library.jpg'''       )


if __name__ == "__main__":
       main()


