



'''simple docstring'''






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 				lowerCamelCase__	(						_A			,      _A			,      _A=[]    ):
  a						: Any					       =					size[0] - overlap_pixels * 2
  a						: int					       =					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
  a						: Optional[int]					       =					np.ones((size_y, size_x)			,      dtype=np.uinta    ) * 255
  a						: Optional[int]					       =					np.pad(_A			,      mode='linear_ramp'			,      pad_width=_A			,      end_values=0    )

  if "l" in remove_borders:
    a						: List[Any]					       =					mask[:, overlap_pixels : mask.shape[1]]
  if "r" in remove_borders:
    a						: int					       =					mask[:, 0 : mask.shape[1] - overlap_pixels]
  if "t" in remove_borders:
    a						: int					       =					mask[overlap_pixels : mask.shape[0], :]
  if "b" in remove_borders:
    a						: Optional[int]					       =					mask[0 : mask.shape[0] - overlap_pixels, :]
  return mask
def 				lowerCamelCase__	(						_A			,      _A			,      _A    ):
  return max(_A			,      min(_A			,      _A    )    )
def 				lowerCamelCase__	(						_A			,      _A			,      _A    ):
  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 				lowerCamelCase__	(						_A			,      _A			,      _A    ):
  a						: Optional[int]					       =					list(_A    )
  rect[0] -= overlap
  rect[1] -= overlap
  rect[2] += overlap
  rect[3] += overlap
  a						: Optional[int]					       =					clamp_rect(_A			,      [0, 0]			,      [image_size[0], image_size[1]]    )
  return rect
def 				lowerCamelCase__	(						_A			,      _A			,      _A			,      _A    ):
  a						: Optional[int]					       =					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 				lowerCamelCase__	(						_A			,      _A    ):
  a						: Any					       =					(original_image_slice * 4, 0, tile.size[0], tile.size[1])
  a						: Optional[Any]					       =					tile.crop(_A    )
  return tile
def 				lowerCamelCase__	(						_A			,      _A    ):
  a						: Dict					       =					n % d
  return n - divisor

class 							a__( lowerCamelCase__     ):







    def __init__(				self			:  Tuple						,      __snake_case			:  AutoencoderKL						,      __snake_case			:  CLIPTextModel						,      __snake_case			:  CLIPTokenizer						,      __snake_case			:  UNetaDConditionModel						,      __snake_case			:  DDPMScheduler						,      __snake_case			:  Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]						,      __snake_case			:  int = 3_50						,      ):
      super().__init__(
          vae=__snake_case						,      text_encoder=__snake_case						,      tokenizer=__snake_case						,      unet=__snake_case						,      low_res_scheduler=__snake_case						,      scheduler=__snake_case						,      max_noise_level=__snake_case						,      )







    def       lowercase_	(				self			:  List[Any]						,      __snake_case			:  List[str]						,      __snake_case			:  Any						,      __snake_case			:  Tuple						,      __snake_case			:  Dict						,      __snake_case			:  List[Any]						,      __snake_case			:  Optional[Any]						,      __snake_case			:  List[str]						,      **__snake_case			:  List[str]     ):
      torch.manual_seed(0     )
      a						: Tuple					       =					(
          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     ),
      )
      a						: Dict					       =					add_overlap_rect(__snake_case						,      __snake_case						,      image.size     )
      a						: Tuple					       =					image.crop(__snake_case     )
      a						: Tuple					       =					((crop_rect[0] + ((crop_rect[2] - crop_rect[0]) / 2)) / image.size[0]) * tile.size[0]
      a						: List[str]					       =					translated_slice_x - (original_image_slice / 2)
      a						: Dict					       =					max(0						,      __snake_case     )
      a						: Optional[int]					       =					squeeze_tile(__snake_case						,      __snake_case						,      __snake_case						,      __snake_case     )
      a						: Tuple					       =					to_input.size
      a						: Union[str, Any]					       =					to_input.resize((tile_size, tile_size)						,      Image.BICUBIC     )
      a						: Any					       =					super(__snake_case						,      self     ).__call__(image=__snake_case						,      **__snake_case     ).images[0]
      a						: str					       =					upscaled_tile.resize((orig_input_size[0] * 4, orig_input_size[1] * 4)						,      Image.BICUBIC     )
      a						: str					       =					unsqueeze_tile(__snake_case						,      __snake_case     )
      a						: Dict					       =					upscaled_tile.resize((tile.size[0] * 4, tile.size[1] * 4)						,      Image.BICUBIC     )
      a						: Union[str, Any]					       =					[]
      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'     )
      a						: Optional[Any]					       =					Image.fromarray(
          make_transparency_mask(
              (upscaled_tile.size[0], upscaled_tile.size[1])						,      tile_border * 4						,      remove_borders=__snake_case     )						,      mode='L'						,      )
      final_image.paste(
          __snake_case						,      (crop_rect_with_overlap[0] * 4, crop_rect_with_overlap[1] * 4)						,      __snake_case     )







    @torch.no_grad()
    def __call__(				self			:  Tuple						,      __snake_case			:  Union[str, List[str]]						,      __snake_case			:  Union[PIL.Image.Image, List[PIL.Image.Image]]						,      __snake_case			:  int = 75						,      __snake_case			:  float = 9.0						,      __snake_case			:  int = 50						,      __snake_case			:  Optional[Union[str, List[str]]] = None						,      __snake_case			:  Optional[int] = 1						,      __snake_case			:  float = 0.0						,      __snake_case			:  Optional[torch.Generator] = None						,      __snake_case			:  Optional[torch.FloatTensor] = None						,      __snake_case			:  Optional[Callable[[int, int, torch.FloatTensor], None]] = None						,      __snake_case			:  int = 1						,      __snake_case			:  int = 1_28						,      __snake_case			:  int = 32						,      __snake_case			:  int = 32						,      ):
      a						: Tuple					       =					Image.new('RGB'						,      (image.size[0] * 4, image.size[1] * 4)     )
      a						: Optional[int]					       =					math.ceil(image.size[0] / tile_size     )
      a						: Tuple					       =					math.ceil(image.size[1] / tile_size     )
      a						: Tuple					       =					tcx * tcy
      a						: Dict					       =					0
      for y in range(__snake_case     ):
        for x in range(__snake_case     ):
          self._process_tile(
              __snake_case						,      __snake_case						,      __snake_case						,      __snake_case						,      __snake_case						,      __snake_case						,      __snake_case						,      prompt=__snake_case						,      num_inference_steps=__snake_case						,      guidance_scale=__snake_case						,      noise_level=__snake_case						,      negative_prompt=__snake_case						,      num_images_per_prompt=__snake_case						,      eta=__snake_case						,      generator=__snake_case						,      latents=__snake_case						,      )
          current_count += 1
          if callback is not None:
            callback({'progress': current_count / total_tile_count, 'image': final_image}     )
      return final_image
def 				lowerCamelCase__	(						):
  # Run a demo
  a						: List[str]					       =					'stabilityai/stable-diffusion-x4-upscaler'
  a						: Optional[int]					       =					StableDiffusionTiledUpscalePipeline.from_pretrained(_A			,      revision='fp16'			,      torch_dtype=torch.floataa    )
  a						: Dict					       =					pipe.to('cuda'    )
  a						: Optional[int]					       =					Image.open('../../docs/source/imgs/diffusers_library.jpg'    )

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

  a						: List[Any]					       =					pipe(image=_A			,      prompt='Black font, white background, vector'			,      noise_level=40			,      callback=_A    )
  final_image.save('diffusers_library.jpg'    )


if __name__ == "__main__":
     main()




'''simple docstring'''






from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)

@flax.struct.dataclass
class 							a__( lowerCamelCase__     ):
    lowercase__			=   42

@flax_register_to_config
class 							a__( nn.Module		,			lowerCamelCase__		,			lowerCamelCase__     ):
    lowercase__			=   32
    lowercase__			=   4
    lowercase__			=   4
    lowercase__			=   (
        "CrossAttnDownBlock2D",
        "CrossAttnDownBlock2D",
        "CrossAttnDownBlock2D",
        "DownBlock2D",
    )
    lowercase__			=   ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
    lowercase__			=   False
    lowercase__			=   (3_20, 6_40, 12_80, 12_80)
    lowercase__			=   2
    lowercase__			=   8
    lowercase__			=   None
    lowercase__			=   12_80
    lowercase__			=   0.0
    lowercase__			=   False
    lowercase__			=   jnp.floataa
    lowercase__			=   True
    lowercase__			=   0
    lowercase__			=   False







    def       lowercase_	(				self			:  List[Any]						,      __snake_case			:  jax.random.KeyArray     ):
      # init input tensors
      a						: Dict					       =					(1, self.in_channels, self.sample_size, self.sample_size)
      a						: str					       =					jnp.zeros(__snake_case						,      dtype=jnp.floataa     )
      a						: Union[str, Any]					       =					jnp.ones((1,)						,      dtype=jnp.intaa     )
      a						: List[Any]					       =					jnp.zeros((1, 1, self.cross_attention_dim)						,      dtype=jnp.floataa     )

      a       ,     a						: Optional[int]					       =					jax.random.split(__snake_case     )
      a						: Any					       =					{'params': params_rng, 'dropout': dropout_rng}

      return self.init(__snake_case						,      __snake_case						,      __snake_case						,      __snake_case     )["params"]







    def       lowercase_	(				self			:  Union[str, Any]     ):
      a						: int					       =					self.block_out_channels
      a						: Tuple					       =					block_out_channels[0] * 4

      if self.num_attention_heads is not None:
        raise ValueError(
            'At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19.'     )

      # If `num_attention_heads` is not defined (which is the case for most models)
      # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
      # The reason for this behavior is to correct for incorrectly named variables that were introduced
      # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
      # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
      # which is why we correct for the naming here.
      a						: str					       =					self.num_attention_heads or self.attention_head_dim

      # input
      a						: List[Any]					       =					nn.Conv(
          block_out_channels[0]						,      kernel_size=(3, 3)						,      strides=(1, 1)						,      padding=((1, 1), (1, 1))						,      dtype=self.dtype						,      )

      # time
      a						: List[str]					       =					FlaxTimesteps(
          block_out_channels[0]						,      flip_sin_to_cos=self.flip_sin_to_cos						,      freq_shift=self.config.freq_shift     )
      a						: Optional[int]					       =					FlaxTimestepEmbedding(__snake_case						,      dtype=self.dtype     )

      a						: Optional[Any]					       =					self.only_cross_attention
      if isinstance(__snake_case						,      __snake_case     ):
        a						: int					       =					(only_cross_attention,) * len(self.down_block_types     )

      if isinstance(__snake_case						,      __snake_case     ):
        a						: str					       =					(num_attention_heads,) * len(self.down_block_types     )

      # down
      a						: int					       =					[]
      a						: Dict					       =					block_out_channels[0]
      for i, down_block_type in enumerate(self.down_block_types     ):
        a						: List[Any]					       =					output_channel
        a						: Dict					       =					block_out_channels[i]
        a						: Union[str, Any]					       =					i == len(__snake_case     ) - 1

        if down_block_type == "CrossAttnDownBlock2D":
          a						: Dict					       =					FlaxCrossAttnDownBlockaD(
              in_channels=__snake_case						,      out_channels=__snake_case						,      dropout=self.dropout						,      num_layers=self.layers_per_block						,      num_attention_heads=num_attention_heads[i]						,      add_downsample=not is_final_block						,      use_linear_projection=self.use_linear_projection						,      only_cross_attention=only_cross_attention[i]						,      use_memory_efficient_attention=self.use_memory_efficient_attention						,      dtype=self.dtype						,      )
        else:
          a						: List[str]					       =					FlaxDownBlockaD(
              in_channels=__snake_case						,      out_channels=__snake_case						,      dropout=self.dropout						,      num_layers=self.layers_per_block						,      add_downsample=not is_final_block						,      dtype=self.dtype						,      )

        down_blocks.append(__snake_case     )
      a						: Dict					       =					down_blocks

      # mid
      a						: Union[str, Any]					       =					FlaxUNetMidBlockaDCrossAttn(
          in_channels=block_out_channels[-1]						,      dropout=self.dropout						,      num_attention_heads=num_attention_heads[-1]						,      use_linear_projection=self.use_linear_projection						,      use_memory_efficient_attention=self.use_memory_efficient_attention						,      dtype=self.dtype						,      )

      # up
      a						: List[Any]					       =					[]
      a						: Dict					       =					list(reversed(__snake_case     )     )
      a						: List[str]					       =					list(reversed(__snake_case     )     )
      a						: List[str]					       =					list(reversed(__snake_case     )     )
      a						: Optional[int]					       =					reversed_block_out_channels[0]
      for i, up_block_type in enumerate(self.up_block_types     ):
        a						: Dict					       =					output_channel
        a						: Dict					       =					reversed_block_out_channels[i]
        a						: Dict					       =					reversed_block_out_channels[min(i + 1						,      len(__snake_case     ) - 1     )]

        a						: List[Any]					       =					i == len(__snake_case     ) - 1

        if up_block_type == "CrossAttnUpBlock2D":
          a						: List[Any]					       =					FlaxCrossAttnUpBlockaD(
              in_channels=__snake_case						,      out_channels=__snake_case						,      prev_output_channel=__snake_case						,      num_layers=self.layers_per_block + 1						,      num_attention_heads=reversed_num_attention_heads[i]						,      add_upsample=not is_final_block						,      dropout=self.dropout						,      use_linear_projection=self.use_linear_projection						,      only_cross_attention=only_cross_attention[i]						,      use_memory_efficient_attention=self.use_memory_efficient_attention						,      dtype=self.dtype						,      )
        else:
          a						: Dict					       =					FlaxUpBlockaD(
              in_channels=__snake_case						,      out_channels=__snake_case						,      prev_output_channel=__snake_case						,      num_layers=self.layers_per_block + 1						,      add_upsample=not is_final_block						,      dropout=self.dropout						,      dtype=self.dtype						,      )

        up_blocks.append(__snake_case     )
        a						: Dict					       =					output_channel
      a						: int					       =					up_blocks

      # out
      a						: int					       =					nn.GroupNorm(num_groups=32						,      epsilon=1e-5     )
      a						: Any					       =					nn.Conv(
          self.out_channels						,      kernel_size=(3, 3)						,      strides=(1, 1)						,      padding=((1, 1), (1, 1))						,      dtype=self.dtype						,      )







    def __call__(				self			:  Optional[Any]						,      __snake_case			:  Dict						,      __snake_case			:  List[str]						,      __snake_case			:  Tuple						,      __snake_case			:  Dict=None						,      __snake_case			:  Any=None						,      __snake_case			:  bool = True						,      __snake_case			:  bool = False						,      ):
      # 1. time
      if not isinstance(__snake_case						,      jnp.ndarray     ):
        a						: Tuple					       =					jnp.array([timesteps]						,      dtype=jnp.intaa     )
      elif isinstance(__snake_case						,      jnp.ndarray     ) and len(timesteps.shape     ) == 0:
        a						: Tuple					       =					timesteps.astype(dtype=jnp.floataa     )
        a						: Optional[int]					       =					jnp.expand_dims(__snake_case						,      0     )

      a						: Optional[Any]					       =					self.time_proj(__snake_case     )
      a						: Optional[int]					       =					self.time_embedding(__snake_case     )

      # 2. pre-process
      a						: str					       =					jnp.transpose(__snake_case						,      (0, 2, 3, 1)     )
      a						: Optional[int]					       =					self.conv_in(__snake_case     )

      # 3. down
      a						: Optional[Any]					       =					(sample,)
      for down_block in self.down_blocks:
        if isinstance(__snake_case						,      __snake_case     ):
          a       ,     a						: Any					       =					down_block(__snake_case						,      __snake_case						,      __snake_case						,      deterministic=not train     )
        else:
          a       ,     a						: Optional[Any]					       =					down_block(__snake_case						,      __snake_case						,      deterministic=not train     )
        down_block_res_samples += res_samples

      if down_block_additional_residuals is not None:
        a						: Tuple					       =					()

        for down_block_res_sample, down_block_additional_residual in zip(
            __snake_case						,      __snake_case     ):
          down_block_res_sample += down_block_additional_residual
          new_down_block_res_samples += (down_block_res_sample,)

        a						: Tuple					       =					new_down_block_res_samples

      # 4. mid
      a						: int					       =					self.mid_block(__snake_case						,      __snake_case						,      __snake_case						,      deterministic=not train     )

      if mid_block_additional_residual is not None:
        sample += mid_block_additional_residual

      # 5. up
      for up_block in self.up_blocks:
        a						: int					       =					down_block_res_samples[-(self.layers_per_block + 1) :]
        a						: List[Any]					       =					down_block_res_samples[: -(self.layers_per_block + 1)]
        if isinstance(__snake_case						,      __snake_case     ):
          a						: Optional[int]					       =					up_block(
              __snake_case						,      temb=__snake_case						,      encoder_hidden_states=__snake_case						,      res_hidden_states_tuple=__snake_case						,      deterministic=not train						,      )
        else:
          a						: Optional[Any]					       =					up_block(__snake_case						,      temb=__snake_case						,      res_hidden_states_tuple=__snake_case						,      deterministic=not train     )

        # 6. post-process
      a						: Any					       =					self.conv_norm_out(__snake_case     )
      a						: List[str]					       =					nn.silu(__snake_case     )
      a						: Optional[int]					       =					self.conv_out(__snake_case     )
      a						: Any					       =					jnp.transpose(__snake_case						,      (0, 3, 1, 2)     )

      if not return_dict:
        return (sample,)

      return FlaxUNetaDConditionOutput(sample=__snake_case     )