


"""simple docstring"""




import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp



def 						__a      (  a,		a,		a,		a,		a = 4_0_9_6					):





       """simple docstring"""







       _a   , _a   , _a          =				key.shape[-3:]
       _a          =				value.shape[-1]
       _a          =				min(a,		a					)
       _a          =				query / jnp.sqrt(a					)

       @functools.partial(jax.checkpoint,		prevent_cse=a					)
       def summarize_chunk(a,		a,		a					):
              _a          =				jnp.einsum("...qhd,...khd->...qhk",		a,		a,		precision=a					)

              _a          =				jnp.max(a,		axis=-1,		keepdims=a					)
              _a          =				jax.lax.stop_gradient(a					)
              _a          =				jnp.exp(attn_weights - max_score					)

              _a          =				jnp.einsum("...vhf,...qhv->...qhf",		a,		a,		precision=a					)
              _a          =				jnp.einsum("...qhk->...qh",		a					)

              return (exp_values, exp_weights.sum(axis=-1					), max_score)

       def chunk_scanner(a					):
              # julienne key array
              _a          =				jax.lax.dynamic_slice(
                  operand=a,		start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0],		slice_sizes=list(key.shape[:-3]					) + [key_chunk_size, num_heads, k_features],		)

              # julienne value array
              _a          =				jax.lax.dynamic_slice(
                  operand=a,		start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0],		slice_sizes=list(value.shape[:-3]					) + [key_chunk_size, num_heads, v_features],		)

              return summarize_chunk(a,		a,		a					)

       _a   , _a   , _a          =				jax.lax.map(f=a,		xs=jnp.arange(0,		a,		a					)					)

       _a          =				jnp.max(a,		axis=0,		keepdims=a					)
       _a          =				jnp.exp(chunk_max - global_max					)

       chunk_values *= jnp.expand_dims(a,		axis=-1					)
       chunk_weights *= max_diffs

       _a          =				chunk_values.sum(axis=0					)
       _a          =				jnp.expand_dims(a,		-1					).sum(axis=0					)

       return all_values / all_weights



def 						__a      (  a,		a,		a,		a=jax.lax.Precision.HIGHEST,		a = 1_0_2_4,		a = 4_0_9_6					):





       """simple docstring"""







       _a   , _a   , _a          =				query.shape[-3:]

       def chunk_scanner(a,		a					):
              # julienne query array
              _a          =				jax.lax.dynamic_slice(
                  operand=a,		start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0],		slice_sizes=list(query.shape[:-3]					) + [min(a,		a					), num_heads, q_features],		)

              return (
                  chunk_idx + query_chunk_size,  # unused ignore it
                  _query_chunk_attention(
                      query=a,		key=a,		value=a,		precision=a,		key_chunk_size=a					),
              )

       _a   , _a          =				jax.lax.scan(
           f=a,		init=0,		xs=a,		length=math.ceil(num_q / query_chunk_size					)					)

       return jnp.concatenate(a,		axis=-3					)  # fuse the chunked result back





class __snake_case		(					nn.Module      ):







   """simple docstring"""




   lowerCAmelCase_				:   int
   lowerCAmelCase_				:   int     	=			8
   lowerCAmelCase_				:   int     	=			64
   lowerCAmelCase_				:   float     	=			0.0
   lowerCAmelCase_				:   bool     	=			False
   lowerCAmelCase_				:   jnp.dtype     	=			jnp.floataa




   def SCREAMING_SNAKE_CASE_      (		self							:Any       ):
          _a          =				self.dim_head * self.heads
          _a          =				self.dim_head**-0.5

          # Weights were exported with old names {to_q, to_k, to_v, to_out}
          _a          =				nn.Dense(UpperCamelCase__				,  use_bias=UpperCamelCase__				,  dtype=self.dtype				,  name="to_q"       )
          _a          =				nn.Dense(UpperCamelCase__				,  use_bias=UpperCamelCase__				,  dtype=self.dtype				,  name="to_k"       )
          _a          =				nn.Dense(UpperCamelCase__				,  use_bias=UpperCamelCase__				,  dtype=self.dtype				,  name="to_v"       )

          _a          =				nn.Dense(self.query_dim				,  dtype=self.dtype				,  name="to_out_0"       )
          _a          =				nn.Dropout(rate=self.dropout       )




   def SCREAMING_SNAKE_CASE_      (		self							:List[str]				,  UpperCamelCase__							:Tuple       ):
          _a   , _a   , _a          =				tensor.shape
          _a          =				self.heads
          _a          =				tensor.reshape(UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__				,  dim // head_size       )
          _a          =				jnp.transpose(UpperCamelCase__				,  (0, 2, 1, 3)       )
          _a          =				tensor.reshape(batch_size * head_size				,  UpperCamelCase__				,  dim // head_size       )
          return tensor




   def SCREAMING_SNAKE_CASE_      (		self							:int				,  UpperCamelCase__							:int       ):
          _a   , _a   , _a          =				tensor.shape
          _a          =				self.heads
          _a          =				tensor.reshape(batch_size // head_size				,  UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__       )
          _a          =				jnp.transpose(UpperCamelCase__				,  (0, 2, 1, 3)       )
          _a          =				tensor.reshape(batch_size // head_size				,  UpperCamelCase__				,  dim * head_size       )
          return tensor




   def __call__(		self							:Optional[int]				,  UpperCamelCase__							:Optional[int]				,  UpperCamelCase__							:int=None				,  UpperCamelCase__							:Optional[Any]=True       ):
          _a          =				hidden_states if context is None else context

          _a          =				self.query(UpperCamelCase__       )
          _a          =				self.key(UpperCamelCase__       )
          _a          =				self.value(UpperCamelCase__       )

          _a          =				self.reshape_heads_to_batch_dim(UpperCamelCase__       )
          _a          =				self.reshape_heads_to_batch_dim(UpperCamelCase__       )
          _a          =				self.reshape_heads_to_batch_dim(UpperCamelCase__       )

          if self.use_memory_efficient_attention:
                 _a          =				query_states.transpose(1				,  0				,  2       )
                 _a          =				key_states.transpose(1				,  0				,  2       )
                 _a          =				value_states.transpose(1				,  0				,  2       )

                 # this if statement create a chunk size for each layer of the unet
                 # the chunk size is equal to the query_length dimension of the deepest layer of the unet

                 _a          =				query_states.shape[-3]
                 if flatten_latent_dim % 64 == 0:
                        _a          =				int(flatten_latent_dim / 64       )
                 elif flatten_latent_dim % 16 == 0:
                        _a          =				int(flatten_latent_dim / 16       )
                 elif flatten_latent_dim % 4 == 0:
                        _a          =				int(flatten_latent_dim / 4       )
                 else:
                        _a          =				int(UpperCamelCase__       )

                 _a          =				jax_memory_efficient_attention(
                     UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__				,  query_chunk_size=UpperCamelCase__				,  key_chunk_size=4_096 * 4       )

                 _a          =				hidden_states.transpose(1				,  0				,  2       )
          else:
                 # compute attentions
                 _a          =				jnp.einsum("b i d, b j d->b i j"				,  UpperCamelCase__				,  UpperCamelCase__       )
                 _a          =				attention_scores * self.scale
                 _a          =				nn.softmax(UpperCamelCase__				,  axis=2       )

                 # attend to values
                 _a          =				jnp.einsum("b i j, b j d -> b i d"				,  UpperCamelCase__				,  UpperCamelCase__       )

          _a          =				self.reshape_batch_dim_to_heads(UpperCamelCase__       )
          _a          =				self.proj_attn(UpperCamelCase__       )
          return self.dropout_layer(UpperCamelCase__				,  deterministic=UpperCamelCase__       )






class __snake_case		(					nn.Module      ):







   """simple docstring"""




   lowerCAmelCase_				:   int
   lowerCAmelCase_				:   int
   lowerCAmelCase_				:   int
   lowerCAmelCase_				:   float     	=			0.0
   lowerCAmelCase_				:   bool     	=			False
   lowerCAmelCase_				:   jnp.dtype     	=			jnp.floataa
   lowerCAmelCase_				:   bool     	=			False




   def SCREAMING_SNAKE_CASE_      (		self							:List[str]       ):
          # self attention (or cross_attention if only_cross_attention is True)
          _a          =				FlaxAttention(
              self.dim				,  self.n_heads				,  self.d_head				,  self.dropout				,  self.use_memory_efficient_attention				,  dtype=self.dtype       )
          # cross attention
          _a          =				FlaxAttention(
              self.dim				,  self.n_heads				,  self.d_head				,  self.dropout				,  self.use_memory_efficient_attention				,  dtype=self.dtype       )
          _a          =				FlaxFeedForward(dim=self.dim				,  dropout=self.dropout				,  dtype=self.dtype       )
          _a          =				nn.LayerNorm(epsilon=1E-5				,  dtype=self.dtype       )
          _a          =				nn.LayerNorm(epsilon=1E-5				,  dtype=self.dtype       )
          _a          =				nn.LayerNorm(epsilon=1E-5				,  dtype=self.dtype       )
          _a          =				nn.Dropout(rate=self.dropout       )




   def __call__(		self							:Union[str, Any]				,  UpperCamelCase__							:int				,  UpperCamelCase__							:str				,  UpperCamelCase__							:List[Any]=True       ):
          # self attention
          _a          =				hidden_states
          if self.only_cross_attention:
                 _a          =				self.attna(self.norma(UpperCamelCase__       )				,  UpperCamelCase__				,  deterministic=UpperCamelCase__       )
          else:
                 _a          =				self.attna(self.norma(UpperCamelCase__       )				,  deterministic=UpperCamelCase__       )
          _a          =				hidden_states + residual

          # cross attention
          _a          =				hidden_states
          _a          =				self.attna(self.norma(UpperCamelCase__       )				,  UpperCamelCase__				,  deterministic=UpperCamelCase__       )
          _a          =				hidden_states + residual

          # feed forward
          _a          =				hidden_states
          _a          =				self.ff(self.norma(UpperCamelCase__       )				,  deterministic=UpperCamelCase__       )
          _a          =				hidden_states + residual

          return self.dropout_layer(UpperCamelCase__				,  deterministic=UpperCamelCase__       )






class __snake_case		(					nn.Module      ):







   """simple docstring"""




   lowerCAmelCase_				:   int
   lowerCAmelCase_				:   int
   lowerCAmelCase_				:   int
   lowerCAmelCase_				:   int     	=			1
   lowerCAmelCase_				:   float     	=			0.0
   lowerCAmelCase_				:   bool     	=			False
   lowerCAmelCase_				:   bool     	=			False
   lowerCAmelCase_				:   jnp.dtype     	=			jnp.floataa
   lowerCAmelCase_				:   bool     	=			False




   def SCREAMING_SNAKE_CASE_      (		self							:str       ):
          _a          =				nn.GroupNorm(num_groups=32				,  epsilon=1E-5       )

          _a          =				self.n_heads * self.d_head
          if self.use_linear_projection:
                 _a          =				nn.Dense(UpperCamelCase__				,  dtype=self.dtype       )
          else:
                 _a          =				nn.Conv(
                     UpperCamelCase__				,  kernel_size=(1, 1)				,  strides=(1, 1)				,  padding="VALID"				,  dtype=self.dtype				,  )

          _a          =				[
              FlaxBasicTransformerBlock(
                  UpperCamelCase__				,  self.n_heads				,  self.d_head				,  dropout=self.dropout				,  only_cross_attention=self.only_cross_attention				,  dtype=self.dtype				,  use_memory_efficient_attention=self.use_memory_efficient_attention				,  )
              for _ in range(self.depth       )
          ]

          if self.use_linear_projection:
                 _a          =				nn.Dense(UpperCamelCase__				,  dtype=self.dtype       )
          else:
                 _a          =				nn.Conv(
                     UpperCamelCase__				,  kernel_size=(1, 1)				,  strides=(1, 1)				,  padding="VALID"				,  dtype=self.dtype				,  )

          _a          =				nn.Dropout(rate=self.dropout       )




   def __call__(		self							:Union[str, Any]				,  UpperCamelCase__							:str				,  UpperCamelCase__							:str				,  UpperCamelCase__							:Any=True       ):
          _a   , _a   , _a   , _a          =				hidden_states.shape
          _a          =				hidden_states
          _a          =				self.norm(UpperCamelCase__       )
          if self.use_linear_projection:
                 _a          =				hidden_states.reshape(UpperCamelCase__				,  height * width				,  UpperCamelCase__       )
                 _a          =				self.proj_in(UpperCamelCase__       )
          else:
                 _a          =				self.proj_in(UpperCamelCase__       )
                 _a          =				hidden_states.reshape(UpperCamelCase__				,  height * width				,  UpperCamelCase__       )

          for transformer_block in self.transformer_blocks:
                 _a          =				transformer_block(UpperCamelCase__				,  UpperCamelCase__				,  deterministic=UpperCamelCase__       )

          if self.use_linear_projection:
                 _a          =				self.proj_out(UpperCamelCase__       )
                 _a          =				hidden_states.reshape(UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__       )
          else:
                 _a          =				hidden_states.reshape(UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__       )
                 _a          =				self.proj_out(UpperCamelCase__       )

          _a          =				hidden_states + residual
          return self.dropout_layer(UpperCamelCase__				,  deterministic=UpperCamelCase__       )






class __snake_case		(					nn.Module      ):







   """simple docstring"""




   lowerCAmelCase_				:   int
   lowerCAmelCase_				:   float     	=			0.0
   lowerCAmelCase_				:   jnp.dtype     	=			jnp.floataa




   def SCREAMING_SNAKE_CASE_      (		self							:Tuple       ):
          # The second linear layer needs to be called
          # net_2 for now to match the index of the Sequential layer
          _a          =				FlaxGEGLU(self.dim				,  self.dropout				,  self.dtype       )
          _a          =				nn.Dense(self.dim				,  dtype=self.dtype       )




   def __call__(		self							:Any				,  UpperCamelCase__							:Optional[int]				,  UpperCamelCase__							:Union[str, Any]=True       ):
          _a          =				self.net_a(UpperCamelCase__				,  deterministic=UpperCamelCase__       )
          _a          =				self.net_a(UpperCamelCase__       )
          return hidden_states






class __snake_case		(					nn.Module      ):







   """simple docstring"""




   lowerCAmelCase_				:   int
   lowerCAmelCase_				:   float     	=			0.0
   lowerCAmelCase_				:   jnp.dtype     	=			jnp.floataa




   def SCREAMING_SNAKE_CASE_      (		self							:Union[str, Any]       ):
          _a          =				self.dim * 4
          _a          =				nn.Dense(inner_dim * 2				,  dtype=self.dtype       )
          _a          =				nn.Dropout(rate=self.dropout       )




   def __call__(		self							:Optional[int]				,  UpperCamelCase__							:Dict				,  UpperCamelCase__							:Union[str, Any]=True       ):
          _a          =				self.proj(UpperCamelCase__       )
          _a   , _a          =				jnp.split(UpperCamelCase__				,  2				,  axis=2       )
          return self.dropout_layer(hidden_linear * nn.gelu(UpperCamelCase__       )				,  deterministic=UpperCamelCase__       )




"""simple docstring"""




from typing import Union

import torch
from PIL import Image
from torchvision import transforms as tfms
from tqdm.auto import tqdm
from transformers import CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)





class __snake_case		(					_SCREAMING_SNAKE_CASE      ):







   """simple docstring"""




   def __init__(		self							:List[Any]				,  UpperCamelCase__							:AutoencoderKL				,  UpperCamelCase__							:CLIPTextModel				,  UpperCamelCase__							:CLIPTokenizer				,  UpperCamelCase__							:UNetaDConditionModel				,  UpperCamelCase__							:Union[PNDMScheduler, LMSDiscreteScheduler, DDIMScheduler]				,  ):
          super().__init__()

          self.register_modules(vae=UpperCamelCase__				,  text_encoder=UpperCamelCase__				,  tokenizer=UpperCamelCase__				,  unet=UpperCamelCase__				,  scheduler=UpperCamelCase__       )




   def SCREAMING_SNAKE_CASE_      (		self							:Tuple				,  UpperCamelCase__							:str       ):
          with torch.no_grad():
                 _a          =				self.vae.encode(tfms.ToTensor()(UpperCamelCase__       ).unsqueeze(0       ).to(self.device       ) * 2 - 1       )
                 _a          =				0.18215 * latent.latent_dist.sample()
          return latent




   def SCREAMING_SNAKE_CASE_      (		self							:List[Any]				,  UpperCamelCase__							:Any       ):
          _a          =				(1 / 0.18215) * latent
          with torch.no_grad():
                 _a          =				self.vae.decode(UpperCamelCase__       ).sample
          _a          =				(img / 2 + 0.5).clamp(0				,  1       )
          _a          =				img.detach().cpu().permute(0				,  2				,  3				,  1       ).numpy()
          _a          =				(img * 255).round().astype("uint8"       )
          return Image.fromarray(img[0]       )




   def SCREAMING_SNAKE_CASE_      (		self							:List[str]				,  UpperCamelCase__							:List[Any]       ):
          _a          =				self.tokenizer(
              UpperCamelCase__				,  padding="max_length"				,  max_length=self.tokenizer.model_max_length				,  truncation=UpperCamelCase__				,  return_tensors="pt"				,  )

          _a          =				self.text_encoder(text_input.input_ids.to(self.device       )       )[0]

          _a          =				self.tokenizer(
              ""				,  padding="max_length"				,  max_length=self.tokenizer.model_max_length				,  truncation=UpperCamelCase__				,  return_tensors="pt"				,  )

          _a          =				self.text_encoder(uncond_input.input_ids.to(self.device       )       )[0]

          return torch.cat([uncond_embedding, text_embedding]       )




   def __call__(		self							:Tuple				,  UpperCamelCase__							:Image.Image				,  UpperCamelCase__							:str				,  UpperCamelCase__							:float = 0.3				,  UpperCamelCase__							:float = 0.6				,  UpperCamelCase__							:float = 0.5				,  UpperCamelCase__							:int = 42				,  UpperCamelCase__							:int = 50				,  UpperCamelCase__							:float = 7.5				,  ):
          _a          =				steps - int(kmin * steps       )
          _a          =				steps - int(kmax * steps       )

          _a          =				self.prep_text(UpperCamelCase__       )

          self.scheduler.set_timesteps(UpperCamelCase__       )

          _a   , _a          =				img.size
          _a          =				self.encode(UpperCamelCase__       )

          torch.manual_seed(UpperCamelCase__       )
          _a          =				torch.randn(
              (1, self.unet.config.in_channels, height // 8, width // 8)				,  ).to(self.device       )

          _a          =				self.scheduler.add_noise(
              UpperCamelCase__				,  UpperCamelCase__				,  timesteps=self.scheduler.timesteps[tmax]				,  )

          _a          =				torch.cat([latents] * 2       )

          _a          =				self.scheduler.scale_model_input(UpperCamelCase__				,  self.scheduler.timesteps[tmax]       )

          with torch.no_grad():
                 _a          =				self.unet(
                     UpperCamelCase__				,  self.scheduler.timesteps[tmax]				,  encoder_hidden_states=UpperCamelCase__				,  ).sample

          _a   , _a          =				pred.chunk(2       )
          _a          =				pred_uncond + guidance_scale * (pred_text - pred_uncond)

          _a          =				self.scheduler.step(UpperCamelCase__				,  self.scheduler.timesteps[tmax]				,  UpperCamelCase__       ).prev_sample

          for i, t in enumerate(tqdm(self.scheduler.timesteps       )       ):
                 if i > tmax:
                        if i < tmin:  # layout generation phase
                               _a          =				self.scheduler.add_noise(
                                   UpperCamelCase__				,  UpperCamelCase__				,  timesteps=UpperCamelCase__				,  )

                               _a          =				(mix_factor * latents) + (
                                   1 - mix_factor
                               ) * orig_latents  # interpolating between layout noise and conditionally generated noise to preserve layout sematics
                               _a          =				torch.cat([input] * 2       )

                        else:  # content generation phase
                               _a          =				torch.cat([latents] * 2       )

                        _a          =				self.scheduler.scale_model_input(UpperCamelCase__				,  UpperCamelCase__       )

                        with torch.no_grad():
                               _a          =				self.unet(
                                   UpperCamelCase__				,  UpperCamelCase__				,  encoder_hidden_states=UpperCamelCase__				,  ).sample

                        _a   , _a          =				pred.chunk(2       )
                        _a          =				pred_uncond + guidance_scale * (pred_text - pred_uncond)

                        _a          =				self.scheduler.step(UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__       ).prev_sample

          return self.decode(UpperCamelCase__       )
