





'''simple docstring'''







from typing import List, Optional, Tuple, Union

import PIL
import torch
from torchvision import transforms

from diffusers.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
from diffusers.schedulers import DDIMScheduler
from diffusers.utils import randn_tensor


snake_case_   :				int		=     transforms.Compose(
    [
        transforms.Resize((2_56, 2_56)),
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5]),
    ]
)





def 				lowerCamelCase_					(     SCREAMING_SNAKE_CASE__					:	List[Any]      )       ->		List[Any]:






						if isinstance(SCREAMING_SNAKE_CASE__,	torch.Tensor      ):
												return image
						elif isinstance(SCREAMING_SNAKE_CASE__,	PIL.Image.Image      ):
												UpperCAmelCase_	:					List[str]     =     [image]

						UpperCAmelCase_	:					Dict     =     [trans(img.convert('''RGB'''      )      ) for img in image]
						UpperCAmelCase_	:					Tuple     =     torch.stack(SCREAMING_SNAKE_CASE__      )
						return image


class     __a					(_lowerCamelCase   ):

							def __init__(							self							:   List[str]				,   __magic_name__							:   List[Any]				,   __magic_name__							:   List[str]      )    ->				List[Any]:







													"""simple docstring"""
													super().__init__()

													# make sure scheduler can always be converted to DDIM
													UpperCAmelCase_	:					Tuple     =     DDIMScheduler.from_config(scheduler.config      )

													self.register_modules(unet=_UpperCamelCase				,   scheduler=_UpperCamelCase      )

							def 			UpperCAmelCase__		(							self							:   int				,   __magic_name__							:   Optional[Any]      )    ->				Union[str, Any]:







													"""simple docstring"""
													if strength < 0 or strength > 1:
																			raise ValueError(F"""The value of strength should in [0.0, 1.0] but is {strength}"""      )

							def 			UpperCAmelCase__		(							self							:   Optional[Any]				,   __magic_name__							:   Dict				,   __magic_name__							:   List[str]				,   __magic_name__							:   List[str]      )    ->				Union[str, Any]:







													"""simple docstring"""
													# get the original timestep using init_timestep
													UpperCAmelCase_	:					Optional[Any]     =     min(int(num_inference_steps * strength      )				,   _UpperCamelCase      )

													UpperCAmelCase_	:					List[Any]     =     max(num_inference_steps - init_timestep				,   0      )
													UpperCAmelCase_	:					Dict     =     self.scheduler.timesteps[t_start:]

													return timesteps, num_inference_steps - t_start

							def 			UpperCAmelCase__		(							self							:   str				,   __magic_name__							:   Tuple				,   __magic_name__							:   List[str]				,   __magic_name__							:   int				,   __magic_name__							:   List[str]				,   __magic_name__							:   int				,   __magic_name__							:   Dict=None      )    ->				List[str]:







													"""simple docstring"""
													if not isinstance(_UpperCamelCase				,   (torch.Tensor, PIL.Image.Image, list)      ):
																			raise ValueError(
																			    F"""`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(_UpperCamelCase      )}"""      )

													UpperCAmelCase_	:					List[str]     =     image.to(device=_UpperCamelCase				,   dtype=_UpperCamelCase      )

													if isinstance(_UpperCamelCase				,   _UpperCamelCase      ) and len(_UpperCamelCase      ) != batch_size:
																			raise ValueError(
																			    F"""You have passed a list of generators of length {len(_UpperCamelCase      )}, but requested an effective batch"""
																			    F""" size of {batch_size}. Make sure the batch size matches the length of the generators."""      )

													UpperCAmelCase_	:					Optional[int]     =     init_latents.shape
													UpperCAmelCase_	:					List[Any]     =     randn_tensor(_UpperCamelCase				,   generator=_UpperCamelCase				,   device=_UpperCamelCase				,   dtype=_UpperCamelCase      )

													# get latents
													print('''add noise to latents at timestep'''				,   _UpperCamelCase      )
													UpperCAmelCase_	:					Optional[Any]     =     self.scheduler.add_noise(_UpperCamelCase				,   _UpperCamelCase				,   _UpperCamelCase      )
													UpperCAmelCase_	:					Dict     =     init_latents

													return latents



							@torch.no_grad()
							def __call__(							self							:   List[str]				,   __magic_name__							:   Union[torch.FloatTensor, PIL.Image.Image] = None				,   __magic_name__							:   float = 0.8				,   __magic_name__							:   int = 1				,   __magic_name__							:   Optional[Union[torch.Generator, List[torch.Generator]]] = None				,   __magic_name__							:   float = 0.0				,   __magic_name__							:   int = 50				,   __magic_name__							:   Optional[bool] = None				,   __magic_name__							:   Optional[str] = "pil"				,   __magic_name__							:   bool = True				,   )    ->				Union[ImagePipelineOutput, Tuple]:







													"""simple docstring"""
													self.check_inputs(_UpperCamelCase      )

													# 2. Preprocess image
													UpperCAmelCase_	:					Optional[int]     =     preprocess(_UpperCamelCase      )

													# 3. set timesteps
													self.scheduler.set_timesteps(_UpperCamelCase				,   device=self.device      )
													UpperCAmelCase_  ,				UpperCAmelCase_	:					Tuple     =     self.get_timesteps(_UpperCamelCase				,   _UpperCamelCase				,   self.device      )
													UpperCAmelCase_	:					Any     =     timesteps[:1].repeat(_UpperCamelCase      )

													# 4. Prepare latent variables
													UpperCAmelCase_	:					List[Any]     =     self.prepare_latents(_UpperCamelCase				,   _UpperCamelCase				,   _UpperCamelCase				,   self.unet.dtype				,   self.device				,   _UpperCamelCase      )
													UpperCAmelCase_	:					Optional[int]     =     latents

													# 5. Denoising loop
													for t in self.progress_bar(_UpperCamelCase      ):
																			# 1. predict noise model_output
																			UpperCAmelCase_	:					Optional[int]     =     self.unet(_UpperCamelCase				,   _UpperCamelCase      ).sample

																			# 2. predict previous mean of image x_t-1 and add variance depending on eta
																			# eta corresponds to η in paper and should be between [0, 1]
																			# do x_t -> x_t-1
																			UpperCAmelCase_	:					List[Any]     =     self.scheduler.step(
																			    _UpperCamelCase				,   _UpperCamelCase				,   _UpperCamelCase				,   eta=_UpperCamelCase				,   use_clipped_model_output=_UpperCamelCase				,   generator=_UpperCamelCase				,   ).prev_sample

													UpperCAmelCase_	:					Optional[int]     =     (image / 2 + 0.5).clamp(0				,   1      )
													UpperCAmelCase_	:					Optional[Any]     =     image.cpu().permute(0				,   2				,   3				,   1      ).numpy()
													if output_type == "pil":
																			UpperCAmelCase_	:					List[Any]     =     self.numpy_to_pil(_UpperCamelCase      )

													if not return_dict:
																			return (image, latent_timestep.item())

													return ImagePipelineOutput(images=_UpperCamelCase      )












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   ):
   lowerCAmelCase__        = 42





@flax_register_to_config
class A_					(			nn.Module ,      _lowerCamelCase ,      _lowerCamelCase   ):
   lowerCAmelCase__        = 32
   lowerCAmelCase__        = 4
   lowerCAmelCase__        = 4
   lowerCAmelCase__        = (
       "CrossAttnDownBlock2D",
       "CrossAttnDownBlock2D",
       "CrossAttnDownBlock2D",
       "DownBlock2D",
   )
   lowerCAmelCase__        = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
   lowerCAmelCase__        = False
   lowerCAmelCase__        = (320, 640, 1280, 1280)
   lowerCAmelCase__        = 2
   lowerCAmelCase__        = 8
   lowerCAmelCase__        = None
   lowerCAmelCase__        = 1280
   lowerCAmelCase__        = 0.0
   lowerCAmelCase__        = False
   lowerCAmelCase__        = jnp.floataa
   lowerCAmelCase__        = True
   lowerCAmelCase__        = 0
   lowerCAmelCase__        = False




   def _lowerCAmelCase      (self	:Tuple    ,   _UpperCamelCase	:jax.random.KeyArray )->					FrozenDict:
         # init input tensors
         __A          =				(1, self.in_channels, self.sample_size, self.sample_size)
         __A          =				jnp.zeros(_UpperCamelCase    ,   dtype=jnp.floataa )
         __A          =				jnp.ones((1,)    ,   dtype=jnp.intaa )
         __A          =				jnp.zeros((1, 1, self.cross_attention_dim)    ,   dtype=jnp.floataa )

         __A   ,						__A          =				jax.random.split(_UpperCamelCase )
         __A          =				{'''params''': params_rng, '''dropout''': dropout_rng}

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




   def _lowerCAmelCase      (self	:Tuple )->					Optional[int]:
         __A          =				self.block_out_channels
         __A          =				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          =				self.num_attention_heads or self.attention_head_dim

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

         # time
         __A          =				FlaxTimesteps(
             block_out_channels[0]    ,   flip_sin_to_cos=self.flip_sin_to_cos    ,   freq_shift=self.config.freq_shift )
         __A          =				FlaxTimestepEmbedding(_UpperCamelCase    ,   dtype=self.dtype )

         __A          =				self.only_cross_attention
         if isinstance(_UpperCamelCase    ,   _UpperCamelCase ):
               __A          =				(only_cross_attention,) * len(self.down_block_types )

         if isinstance(_UpperCamelCase    ,   _UpperCamelCase ):
               __A          =				(num_attention_heads,) * len(self.down_block_types )

         # down
         __A          =				[]
         __A          =				block_out_channels[0]
         for i, down_block_type in enumerate(self.down_block_types ):
               __A          =				output_channel
               __A          =				block_out_channels[i]
               __A          =				i == len(_UpperCamelCase ) - 1

               if down_block_type == "CrossAttnDownBlock2D":
                     __A          =				FlaxCrossAttnDownBlockaD(
                         in_channels=_UpperCamelCase    ,   out_channels=_UpperCamelCase    ,   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          =				FlaxDownBlockaD(
                         in_channels=_UpperCamelCase    ,   out_channels=_UpperCamelCase    ,   dropout=self.dropout    ,   num_layers=self.layers_per_block    ,   add_downsample=not is_final_block    ,   dtype=self.dtype    ,   )

               down_blocks.append(_UpperCamelCase )
         __A          =				down_blocks

         # mid
         __A          =				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          =				[]
         __A          =				list(reversed(_UpperCamelCase ) )
         __A          =				list(reversed(_UpperCamelCase ) )
         __A          =				list(reversed(_UpperCamelCase ) )
         __A          =				reversed_block_out_channels[0]
         for i, up_block_type in enumerate(self.up_block_types ):
               __A          =				output_channel
               __A          =				reversed_block_out_channels[i]
               __A          =				reversed_block_out_channels[min(i + 1    ,   len(_UpperCamelCase ) - 1 )]

               __A          =				i == len(_UpperCamelCase ) - 1

               if up_block_type == "CrossAttnUpBlock2D":
                     __A          =				FlaxCrossAttnUpBlockaD(
                         in_channels=_UpperCamelCase    ,   out_channels=_UpperCamelCase    ,   prev_output_channel=_UpperCamelCase    ,   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          =				FlaxUpBlockaD(
                         in_channels=_UpperCamelCase    ,   out_channels=_UpperCamelCase    ,   prev_output_channel=_UpperCamelCase    ,   num_layers=self.layers_per_block + 1    ,   add_upsample=not is_final_block    ,   dropout=self.dropout    ,   dtype=self.dtype    ,   )

               up_blocks.append(_UpperCamelCase )
               __A          =				output_channel
         __A          =				up_blocks

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






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

         __A          =				self.time_proj(_UpperCamelCase )
         __A          =				self.time_embedding(_UpperCamelCase )

         # 2. pre-process
         __A          =				jnp.transpose(_UpperCamelCase    ,   (0, 2, 3, 1) )
         __A          =				self.conv_in(_UpperCamelCase )

         # 3. down
         __A          =				(sample,)
         for down_block in self.down_blocks:
               if isinstance(_UpperCamelCase    ,   _UpperCamelCase ):
                     __A   ,						__A          =				down_block(_UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase    ,   deterministic=not train )
               else:
                     __A   ,						__A          =				down_block(_UpperCamelCase    ,   _UpperCamelCase    ,   deterministic=not train )
               down_block_res_samples += res_samples

         if down_block_additional_residuals is not None:
               __A          =				()

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

               __A          =				new_down_block_res_samples

         # 4. mid
         __A          =				self.mid_block(_UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase    ,   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          =				down_block_res_samples[-(self.layers_per_block + 1) :]
               __A          =				down_block_res_samples[: -(self.layers_per_block + 1)]
               if isinstance(_UpperCamelCase    ,   _UpperCamelCase ):
                     __A          =				up_block(
                         _UpperCamelCase    ,   temb=_UpperCamelCase    ,   encoder_hidden_states=_UpperCamelCase    ,   res_hidden_states_tuple=_UpperCamelCase    ,   deterministic=not train    ,   )
               else:
                     __A          =				up_block(_UpperCamelCase    ,   temb=_UpperCamelCase    ,   res_hidden_states_tuple=_UpperCamelCase    ,   deterministic=not train )

        # 6. post-process
         __A          =				self.conv_norm_out(_UpperCamelCase )
         __A          =				nn.silu(_UpperCamelCase )
         __A          =				self.conv_out(_UpperCamelCase )
         __A          =				jnp.transpose(_UpperCamelCase    ,   (0, 3, 1, 2) )

         if not return_dict:
               return (sample,)

         return FlaxUNetaDConditionOutput(sample=_UpperCamelCase )





