


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  lowerCAmelCase__(    snake_case__						):






    '''simple docstring'''





    A_  :     jnp.ndarray





@flax_register_to_config
class  lowerCAmelCase__(    nn.Module    ,    snake_case__    ,    snake_case__						):






    '''simple docstring'''





    A_  :     int			= 3_2
    A_  :     int			= 4
    A_  :     int			= 4
    A_  :     Tuple[str]			= (
        "CrossAttnDownBlock2D",
        "CrossAttnDownBlock2D",
        "CrossAttnDownBlock2D",
        "DownBlock2D",
    )
    A_  :     Tuple[str]			= ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
    A_  :     Union[bool, Tuple[bool]]			= False
    A_  :     Tuple[int]			= (3_2_0, 6_4_0, 1_2_8_0, 1_2_8_0)
    A_  :     int			= 2
    A_  :     Union[int, Tuple[int]]			= 8
    A_  :     Optional[Union[int, Tuple[int]]]			= None
    A_  :     int			= 1_2_8_0
    A_  :     float			= 0.0
    A_  :     bool			= False
    A_  :     jnp.dtype			= jnp.floataa
    A_  :     bool			= True
    A_  :     int			= 0
    A_  :     bool			= False





    def 					_lowerCamelCase (	self				:   Tuple       ,   __snake_case				:   jax.random.KeyArray	):





       '''simple docstring'''
       # init input tensors
       UpperCAmelCase_							:      Union[str, Any]						=   (1, self.in_channels, self.sample_size, self.sample_size)
       UpperCAmelCase_							:      Optional[int]						=   jnp.zeros(__snake_case       ,   dtype=jnp.floataa	)
       UpperCAmelCase_							:      int						=   jnp.ones((1,)       ,   dtype=jnp.intaa	)
       UpperCAmelCase_							:      Dict						=   jnp.zeros((1, 1, self.cross_attention_dim)       ,   dtype=jnp.floataa	)

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

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





    def 					_lowerCamelCase (	self				:   Tuple	):





       '''simple docstring'''
       UpperCAmelCase_							:      List[str]						=   self.block_out_channels
       UpperCAmelCase_							:      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.
       UpperCAmelCase_							:      Optional[int]						=   self.num_attention_heads or self.attention_head_dim

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

       # time
       UpperCAmelCase_							:      Union[str, Any]						=   FlaxTimesteps(
           block_out_channels[0]       ,   flip_sin_to_cos=self.flip_sin_to_cos       ,   freq_shift=self.config.freq_shift	)
       UpperCAmelCase_							:      Tuple						=   FlaxTimestepEmbedding(__snake_case       ,   dtype=self.dtype	)

       UpperCAmelCase_							:      Union[str, Any]						=   self.only_cross_attention
       if isinstance(__snake_case       ,   __snake_case	):
          UpperCAmelCase_							:      Optional[Any]						=   (only_cross_attention,) * len(self.down_block_types	)

       if isinstance(__snake_case       ,   __snake_case	):
          UpperCAmelCase_							:      Optional[Any]						=   (num_attention_heads,) * len(self.down_block_types	)

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

          if down_block_type == "CrossAttnDownBlock2D":
             UpperCAmelCase_							:      List[str]						=   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:
             UpperCAmelCase_							:      List[Any]						=   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	)
       UpperCAmelCase_							:      Dict						=   down_blocks

       # mid
       UpperCAmelCase_							:      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
       UpperCAmelCase_							:      int						=   []
       UpperCAmelCase_							:      Dict						=   list(reversed(__snake_case	)	)
       UpperCAmelCase_							:      Optional[Any]						=   list(reversed(__snake_case	)	)
       UpperCAmelCase_							:      Any						=   list(reversed(__snake_case	)	)
       UpperCAmelCase_							:      Dict						=   reversed_block_out_channels[0]
       for i, up_block_type in enumerate(self.up_block_types	):
          UpperCAmelCase_							:      List[str]						=   output_channel
          UpperCAmelCase_							:      List[Any]						=   reversed_block_out_channels[i]
          UpperCAmelCase_							:      Dict						=   reversed_block_out_channels[min(i + 1       ,   len(__snake_case	) - 1	)]

          UpperCAmelCase_							:      int						=   i == len(__snake_case	) - 1

          if up_block_type == "CrossAttnUpBlock2D":
             UpperCAmelCase_							:      List[str]						=   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:
             UpperCAmelCase_							:      str						=   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	)
          UpperCAmelCase_							:      Any						=   output_channel
       UpperCAmelCase_							:      str						=   up_blocks

       # out
       UpperCAmelCase_							:      Dict						=   nn.GroupNorm(num_groups=32       ,   epsilon=1E-5	)
       UpperCAmelCase_							:      Union[str, 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				:   Union[str, Any]       ,   __snake_case				:   Tuple       ,   __snake_case				:   List[Any]       ,   __snake_case				:   str=None       ,   __snake_case				:   List[Any]=None       ,   __snake_case				:   bool = True       ,   __snake_case				:   bool = False       ,   ):





       '''simple docstring'''
       # 1. time
       if not isinstance(__snake_case       ,   jnp.ndarray	):
          UpperCAmelCase_							:      List[Any]						=   jnp.array([timesteps]       ,   dtype=jnp.intaa	)
       elif isinstance(__snake_case       ,   jnp.ndarray	) and len(timesteps.shape	) == 0:
          UpperCAmelCase_							:      Tuple						=   timesteps.astype(dtype=jnp.floataa	)
          UpperCAmelCase_							:      List[str]						=   jnp.expand_dims(__snake_case       ,   0	)

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

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

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

       if down_block_additional_residuals is not None:
          UpperCAmelCase_							:      Dict						=   ()

          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,)

          UpperCAmelCase_							:      Tuple						=   new_down_block_res_samples

       # 4. mid
       UpperCAmelCase_							:      Tuple						=   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:
          UpperCAmelCase_							:      List[str]						=   down_block_res_samples[-(self.layers_per_block + 1) :]
          UpperCAmelCase_							:      Optional[Any]						=   down_block_res_samples[: -(self.layers_per_block + 1)]
          if isinstance(__snake_case       ,   __snake_case	):
             UpperCAmelCase_							:      Optional[int]						=   up_block(
                 __snake_case       ,   temb=__snake_case       ,   encoder_hidden_states=__snake_case       ,   res_hidden_states_tuple=__snake_case       ,   deterministic=not train       ,   )
          else:
             UpperCAmelCase_							:      List[Any]						=   up_block(__snake_case       ,   temb=__snake_case       ,   res_hidden_states_tuple=__snake_case       ,   deterministic=not train	)

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

       if not return_dict:
          return (sample,)

       return FlaxUNetaDConditionOutput(sample=__snake_case	)



# This code is adapted from OpenAI's release
# https://github.com/openai/human-eval/blob/master/human_eval/execution.py

import contextlib
import faulthandler
import io
import multiprocessing
import os
import platform
import signal
import tempfile
def      snake_case_						( __lowercase			,       __lowercase			,       __lowercase			,       __lowercase       ):
   UpperCAmelCase_							:      int						=   multiprocessing.Manager()
   UpperCAmelCase_							:      Union[str, Any]						=   manager.list()

   UpperCAmelCase_							:      int						=   multiprocessing.Process(target=__lowercase			,       args=(check_program, result, timeout)       )
   p.start()
   p.join(timeout=timeout + 1       )
   if p.is_alive():
      p.kill()

   if not result:
      result.append('''timed out'''       )

   return {
       "task_id": task_id,
       "passed": result[0] == "passed",
       "result": result[0],
       "completion_id": completion_id,
   }
def      snake_case_						( __lowercase			,       __lowercase			,       __lowercase       ):
   with create_tempdir():
      # These system calls are needed when cleaning up tempdir.
      import os
      import shutil

      UpperCAmelCase_							:      str						=   shutil.rmtree
      UpperCAmelCase_							:      Tuple						=   os.rmdir
      UpperCAmelCase_							:      Dict						=   os.chdir

      # Disable functionalities that can make destructive changes to the test.
      reliability_guard()

      # Run program.
      try:
         UpperCAmelCase_							:      Optional[int]						=   {}
         with swallow_io():
            with time_limit(__lowercase       ):
               exec(__lowercase			,       __lowercase       )
         result.append('''passed'''       )
      except TimeoutException:
         result.append('''timed out'''       )
      except BaseException as e:
         result.append(F'''failed: {e}'''       )

      # Needed for cleaning up.
      UpperCAmelCase_							:      Optional[int]						=   rmtree
      UpperCAmelCase_							:      Optional[Any]						=   rmdir
      UpperCAmelCase_							:      Optional[Any]						=   chdir
@contextlib.contextmanager
def      snake_case_						( __lowercase       ):
   def signal_handler(__lowercase			,       __lowercase       ):
      raise TimeoutException('''Timed out!'''       )

   signal.setitimer(signal.ITIMER_REAL			,       __lowercase       )
   signal.signal(signal.SIGALRM			,       __lowercase       )
   try:
      yield
   finally:
      signal.setitimer(signal.ITIMER_REAL			,       0       )
@contextlib.contextmanager
def      snake_case_						( ):
   UpperCAmelCase_							:      Optional[Any]						=   WriteOnlyStringIO()
   with contextlib.redirect_stdout(__lowercase       ):
      with contextlib.redirect_stderr(__lowercase       ):
         with redirect_stdin(__lowercase       ):
            yield
@contextlib.contextmanager
def      snake_case_						( ):
   with tempfile.TemporaryDirectory() as dirname:
      with chdir(__lowercase       ):
         yield dirname





class  lowerCAmelCase__(    snake_case__						):






    '''simple docstring'''





    pass





class  lowerCAmelCase__(    io.StringIO						):






    '''simple docstring'''





    def 					_lowerCamelCase (	self				:   Dict       ,   *__snake_case				:   List[Any]       ,   **__snake_case				:   int	):





       '''simple docstring'''
       raise OSError





    def 					_lowerCamelCase (	self				:   Dict       ,   *__snake_case				:   int       ,   **__snake_case				:   Any	):





       '''simple docstring'''
       raise OSError





    def 					_lowerCamelCase (	self				:   int       ,   *__snake_case				:   List[str]       ,   **__snake_case				:   Optional[Any]	):





       '''simple docstring'''
       raise OSError





    def 					_lowerCamelCase (	self				:   Union[str, Any]       ,   *__snake_case				:   Optional[Any]       ,   **__snake_case				:   List[Any]	):





       '''simple docstring'''
       return False





class  lowerCAmelCase__(    contextlib._RedirectStream						):  # type: ignore






    '''simple docstring'''





    A_  :     Optional[Any]			= 'stdin'





@contextlib.contextmanager
def      snake_case_						( __lowercase       ):
   if root == ".":
      yield
      return
   UpperCAmelCase_							:      Tuple						=   os.getcwd()
   os.chdir(__lowercase       )
   try:
      yield
   except BaseException as exc:
      raise exc
   finally:
      os.chdir(__lowercase       )


def      snake_case_						( __lowercase=None       ):

   if maximum_memory_bytes is not None:
      import resource

      resource.setrlimit(resource.RLIMIT_AS			,       (maximum_memory_bytes, maximum_memory_bytes)       )
      resource.setrlimit(resource.RLIMIT_DATA			,       (maximum_memory_bytes, maximum_memory_bytes)       )
      if not platform.uname().system == "Darwin":
         resource.setrlimit(resource.RLIMIT_STACK			,       (maximum_memory_bytes, maximum_memory_bytes)       )

   faulthandler.disable()

   import builtins

   UpperCAmelCase_							:      Any						=   None
   UpperCAmelCase_							:      Any						=   None

   import os

   UpperCAmelCase_							:      Union[str, Any]						=   '''1'''

   UpperCAmelCase_							:      List[str]						=   None
   UpperCAmelCase_							:      Dict						=   None
   UpperCAmelCase_							:      Any						=   None
   UpperCAmelCase_							:      Optional[int]						=   None
   UpperCAmelCase_							:      List[str]						=   None
   UpperCAmelCase_							:      List[Any]						=   None
   UpperCAmelCase_							:      Union[str, Any]						=   None
   UpperCAmelCase_							:      str						=   None
   UpperCAmelCase_							:      Dict						=   None
   UpperCAmelCase_							:      int						=   None
   UpperCAmelCase_							:      Dict						=   None
   UpperCAmelCase_							:      List[Any]						=   None
   UpperCAmelCase_							:      List[Any]						=   None
   UpperCAmelCase_							:      Optional[Any]						=   None
   UpperCAmelCase_							:      Optional[Any]						=   None
   UpperCAmelCase_							:      Dict						=   None
   UpperCAmelCase_							:      Optional[Any]						=   None
   UpperCAmelCase_							:      Optional[Any]						=   None
   UpperCAmelCase_							:      Optional[Any]						=   None
   UpperCAmelCase_							:      List[Any]						=   None
   UpperCAmelCase_							:      int						=   None
   UpperCAmelCase_							:      Union[str, Any]						=   None
   UpperCAmelCase_							:      List[Any]						=   None
   UpperCAmelCase_							:      Dict						=   None
   UpperCAmelCase_							:      Union[str, Any]						=   None
   UpperCAmelCase_							:      Dict						=   None
   UpperCAmelCase_							:      Any						=   None

   import shutil

   UpperCAmelCase_							:      Union[str, Any]						=   None
   UpperCAmelCase_							:      Dict						=   None
   UpperCAmelCase_							:      Tuple						=   None

   import subprocess

   UpperCAmelCase_							:      Dict						=   None  # type: ignore

   UpperCAmelCase_							:      Union[str, Any]						=   None

   import sys

   UpperCAmelCase_							:      Optional[Any]						=   None
   UpperCAmelCase_							:      List[Any]						=   None
   UpperCAmelCase_							:      List[Any]						=   None
   UpperCAmelCase_							:      List[str]						=   None
   UpperCAmelCase_							:      Dict						=   None