





'''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,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
)
@flax.struct.dataclass
class       a_   (				_lowerCAmelCase							):
     __A      =			42
     __A      =			42







class       a_   (				nn.Module							):
     __A      =			42
     __A      =			(16, 32, 96, 256)
     __A      =			jnp.floataa





     def 							lowercase__      ( self		:					Optional[Any]						):





        """simple docstring"""




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

        lowercase_			:List[Any]    			= []
        for i in range(len(self.block_out_channels						) - 1						):
           lowercase_			:Dict    			= self.block_out_channels[i]
           lowercase_			:List[str]    			= self.block_out_channels[i + 1]
           lowercase_			:Optional[int]    			= nn.Conv(
               lowercase	,      kernel_size=(3, 3)	,      padding=((1, 1), (1, 1))	,      dtype=self.dtype	,      )
           blocks.append(lowercase						)
           lowercase_			:Union[str, Any]    			= nn.Conv(
               lowercase	,      kernel_size=(3, 3)	,      strides=(2, 2)	,      padding=((1, 1), (1, 1))	,      dtype=self.dtype	,      )
           blocks.append(lowercase						)
        lowercase_			:Tuple    			= blocks

        lowercase_			:List[str]    			= nn.Conv(
            self.conditioning_embedding_channels	,      kernel_size=(3, 3)	,      padding=((1, 1), (1, 1))	,      kernel_init=nn.initializers.zeros_init()	,      bias_init=nn.initializers.zeros_init()	,      dtype=self.dtype	,      )







     def __call__( self		:					Optional[int]	,      lowercase		:					Optional[int]						):





        """simple docstring"""




        lowercase_			:Union[str, Any]    			= self.conv_in(lowercase						)
        lowercase_			:List[Any]    			= nn.silu(lowercase						)

        for block in self.blocks:
           lowercase_			:Dict    			= block(lowercase						)
           lowercase_			:Optional[int]    			= nn.silu(lowercase						)

        lowercase_			:Union[str, Any]    			= self.conv_out(lowercase						)

        return embedding







@flax_register_to_config
class       a_   (				nn.Module	,						_lowerCAmelCase	,						_lowerCAmelCase							):
     __A      =			32
     __A      =			4
     __A      =			(
         "CrossAttnDownBlock2D",
         "CrossAttnDownBlock2D",
         "CrossAttnDownBlock2D",
         "DownBlock2D",
     )
     __A      =			False
     __A      =			(320, 640, 1_280, 1_280)
     __A      =			2
     __A      =			8
     __A      =			None
     __A      =			1_280
     __A      =			0.0
     __A      =			False
     __A      =			jnp.floataa
     __A      =			True
     __A      =			0
     __A      =			"rgb"
     __A      =			(16, 32, 96, 256)





     def 							lowercase__      ( self		:					int	,      lowercase		:					jax.random.KeyArray						):





        """simple docstring"""




        lowercase_			:List[Any]    			= (1, self.in_channels, self.sample_size, self.sample_size)
        lowercase_			:int    			= jnp.zeros(lowercase	,      dtype=jnp.floataa						)
        lowercase_			:List[str]    			= jnp.ones((1,)	,      dtype=jnp.intaa						)
        lowercase_			:Any    			= jnp.zeros((1, 1, self.cross_attention_dim)	,      dtype=jnp.floataa						)
        lowercase_			:int    			= (1, 3, self.sample_size * 8, self.sample_size * 8)
        lowercase_			:List[str]    			= jnp.zeros(lowercase	,      dtype=jnp.floataa						)

        lowercase_			,							lowercase_			:List[str]    			= jax.random.split(lowercase						)
        lowercase_			:List[Any]    			= {"params": params_rng, "dropout": dropout_rng}

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





     def 							lowercase__      ( self		:					Union[str, Any]						):





        """simple docstring"""




        lowercase_			:int    			= self.block_out_channels
        lowercase_			:int    			= block_out_channels[0] * 4

        # 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.
        lowercase_			:Optional[Any]    			= self.num_attention_heads or self.attention_head_dim

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

        # time
        lowercase_			:List[str]    			= FlaxTimesteps(
            block_out_channels[0]	,      flip_sin_to_cos=self.flip_sin_to_cos	,      freq_shift=self.config.freq_shift						)
        lowercase_			:Union[str, Any]    			= FlaxTimestepEmbedding(lowercase	,      dtype=self.dtype						)

        lowercase_			:Tuple    			= FlaxControlNetConditioningEmbedding(
            conditioning_embedding_channels=block_out_channels[0]	,      block_out_channels=self.conditioning_embedding_out_channels	,      )

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

        if isinstance(lowercase	,      lowercase						):
           lowercase_			:Any    			= (num_attention_heads,) * len(self.down_block_types						)

        # down
        lowercase_			:Tuple    			= []
        lowercase_			:Any    			= []

        lowercase_			:str    			= block_out_channels[0]

        lowercase_			:Tuple    			= nn.Conv(
            lowercase	,      kernel_size=(1, 1)	,      padding="VALID"	,      kernel_init=nn.initializers.zeros_init()	,      bias_init=nn.initializers.zeros_init()	,      dtype=self.dtype	,      )
        controlnet_down_blocks.append(lowercase						)

        for i, down_block_type in enumerate(self.down_block_types						):
           lowercase_			:Any    			= output_channel
           lowercase_			:Any    			= block_out_channels[i]
           lowercase_			:Union[str, Any]    			= i == len(lowercase						) - 1

           if down_block_type == "CrossAttnDownBlock2D":
              lowercase_			:List[Any]    			= FlaxCrossAttnDownBlockaD(
                  in_channels=lowercase	,      out_channels=lowercase	,      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]	,      dtype=self.dtype	,      )
           else:
              lowercase_			:Optional[int]    			= FlaxDownBlockaD(
                  in_channels=lowercase	,      out_channels=lowercase	,      dropout=self.dropout	,      num_layers=self.layers_per_block	,      add_downsample=not is_final_block	,      dtype=self.dtype	,      )

           down_blocks.append(lowercase						)

           for _ in range(self.layers_per_block						):
              lowercase_			:Dict    			= nn.Conv(
                  lowercase	,      kernel_size=(1, 1)	,      padding="VALID"	,      kernel_init=nn.initializers.zeros_init()	,      bias_init=nn.initializers.zeros_init()	,      dtype=self.dtype	,      )
              controlnet_down_blocks.append(lowercase						)

           if not is_final_block:
              lowercase_			:int    			= nn.Conv(
                  lowercase	,      kernel_size=(1, 1)	,      padding="VALID"	,      kernel_init=nn.initializers.zeros_init()	,      bias_init=nn.initializers.zeros_init()	,      dtype=self.dtype	,      )
              controlnet_down_blocks.append(lowercase						)

        lowercase_			:List[str]    			= down_blocks
        lowercase_			:Union[str, Any]    			= controlnet_down_blocks

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

        lowercase_			:Dict    			= nn.Conv(
            lowercase	,      kernel_size=(1, 1)	,      padding="VALID"	,      kernel_init=nn.initializers.zeros_init()	,      bias_init=nn.initializers.zeros_init()	,      dtype=self.dtype	,      )







     def __call__( self		:					Union[str, Any]	,      lowercase		:					Dict	,      lowercase		:					Tuple	,      lowercase		:					Union[str, Any]	,      lowercase		:					int	,      lowercase		:					float = 1.0	,      lowercase		:					bool = True	,      lowercase		:					bool = False	,      ):





        """simple docstring"""




        lowercase_			:List[str]    			= self.controlnet_conditioning_channel_order
        if channel_order == "bgr":
           lowercase_			:str    			= jnp.flip(lowercase	,      axis=1						)

        # 1. time
        if not isinstance(lowercase	,      jnp.ndarray						):
           lowercase_			:int    			= jnp.array([timesteps]	,      dtype=jnp.intaa						)
        elif isinstance(lowercase	,      jnp.ndarray						) and len(timesteps.shape						) == 0:
           lowercase_			:Dict    			= timesteps.astype(dtype=jnp.floataa						)
           lowercase_			:int    			= jnp.expand_dims(lowercase	,      0						)

        lowercase_			:Tuple    			= self.time_proj(lowercase						)
        lowercase_			:Tuple    			= self.time_embedding(lowercase						)

        # 2. pre-process
        lowercase_			:Union[str, Any]    			= jnp.transpose(lowercase	,      (0, 2, 3, 1)						)
        lowercase_			:Any    			= self.conv_in(lowercase						)

        lowercase_			:int    			= jnp.transpose(lowercase	,      (0, 2, 3, 1)						)
        lowercase_			:Any    			= self.controlnet_cond_embedding(lowercase						)
        sample += controlnet_cond

        # 3. down
        lowercase_			:Union[str, Any]    			= (sample,)
        for down_block in self.down_blocks:
           if isinstance(lowercase	,      lowercase						):
              lowercase_			,							lowercase_			:List[str]    			= down_block(lowercase	,      lowercase	,      lowercase	,      deterministic=not train						)
           else:
              lowercase_			,							lowercase_			:Any    			= down_block(lowercase	,      lowercase	,      deterministic=not train						)
           down_block_res_samples += res_samples

        # 4. mid
        lowercase_			:Any    			= self.mid_block(lowercase	,      lowercase	,      lowercase	,      deterministic=not train						)

        # 5. contronet blocks
        lowercase_			:Optional[int]    			= ()
        for down_block_res_sample, controlnet_block in zip(lowercase	,      self.controlnet_down_blocks						):
           lowercase_			:int    			= controlnet_block(lowercase						)
           controlnet_down_block_res_samples += (down_block_res_sample,)

        lowercase_			:Optional[int]    			= controlnet_down_block_res_samples

        lowercase_			:str    			= self.controlnet_mid_block(lowercase						)

        # 6. scaling
        lowercase_			:str    			= [sample * conditioning_scale for sample in down_block_res_samples]
        mid_block_res_sample *= conditioning_scale

        if not return_dict:
           return (down_block_res_samples, mid_block_res_sample)

        return FlaxControlNetOutput(
            down_block_res_samples=lowercase	,      mid_block_res_sample=lowercase						)









'''simple docstring'''


import unittest

from transformers.testing_utils import CaptureStdout
from transformers.tools.python_interpreter import evaluate
def 							UpperCAmelCase_  (						__lowerCamelCase					:					List[Any] ):
   return x + 2


class       a_   (				unittest.TestCase							):





     def 							lowercase__      ( self		:					Optional[Any]						):





        """simple docstring"""




        lowercase_			:Optional[int]    			= "x = 3"
        lowercase_			:Any    			= {}
        lowercase_			:Any    			= evaluate(lowercase	,      {}	,      state=lowercase						)
        assert result == 3
        self.assertDictEqual(lowercase	,      {"x": 3}						)

        lowercase_			:Dict    			= "x = y"
        lowercase_			:Dict    			= {"y": 5}
        lowercase_			:str    			= evaluate(lowercase	,      {}	,      state=lowercase						)
        # evaluate returns the value of the last assignment.
        assert result == 5
        self.assertDictEqual(lowercase	,      {"x": 5, "y": 5}						)





     def 							lowercase__      ( self		:					List[Any]						):





        """simple docstring"""




        lowercase_			:Dict    			= "y = add_two(x)"
        lowercase_			:Optional[Any]    			= {"x": 3}
        lowercase_			:int    			= evaluate(lowercase	,      {"add_two": add_two}	,      state=lowercase						)
        assert result == 5
        self.assertDictEqual(lowercase	,      {"x": 3, "y": 5}						)

        # Won't work without the tool
        with CaptureStdout() as out:
           lowercase_			:List[Any]    			= evaluate(lowercase	,      {}	,      state=lowercase						)
        assert result is None
        assert "tried to execute add_two" in out.out





     def 							lowercase__      ( self		:					int						):





        """simple docstring"""




        lowercase_			:Optional[int]    			= "x = 3"
        lowercase_			:Any    			= {}
        lowercase_			:Any    			= evaluate(lowercase	,      {}	,      state=lowercase						)
        assert result == 3
        self.assertDictEqual(lowercase	,      {"x": 3}						)





     def 							lowercase__      ( self		:					Any						):





        """simple docstring"""




        lowercase_			:Optional[Any]    			= "test_dict = {'x': x, 'y': add_two(x)}"
        lowercase_			:List[str]    			= {"x": 3}
        lowercase_			:Any    			= evaluate(lowercase	,      {"add_two": add_two}	,      state=lowercase						)
        self.assertDictEqual(lowercase	,      {"x": 3, "y": 5}						)
        self.assertDictEqual(lowercase	,      {"x": 3, "test_dict": {"x": 3, "y": 5}}						)





     def 							lowercase__      ( self		:					Optional[Any]						):





        """simple docstring"""




        lowercase_			:Optional[int]    			= "x = 3\ny = 5"
        lowercase_			:Union[str, Any]    			= {}
        lowercase_			:Optional[int]    			= evaluate(lowercase	,      {}	,      state=lowercase						)
        # evaluate returns the value of the last assignment.
        assert result == 5
        self.assertDictEqual(lowercase	,      {"x": 3, "y": 5}						)





     def 							lowercase__      ( self		:					List[Any]						):





        """simple docstring"""




        lowercase_			:List[str]    			= "text = f'This is x: {x}.'"
        lowercase_			:int    			= {"x": 3}
        lowercase_			:int    			= evaluate(lowercase	,      {}	,      state=lowercase						)
        # evaluate returns the value of the last assignment.
        assert result == "This is x: 3."
        self.assertDictEqual(lowercase	,      {"x": 3, "text": "This is x: 3."}						)





     def 							lowercase__      ( self		:					Optional[int]						):





        """simple docstring"""




        lowercase_			:List[Any]    			= "if x <= 3:\n    y = 2\nelse:\n    y = 5"
        lowercase_			:List[str]    			= {"x": 3}
        lowercase_			:Union[str, Any]    			= evaluate(lowercase	,      {}	,      state=lowercase						)
        # evaluate returns the value of the last assignment.
        assert result == 2
        self.assertDictEqual(lowercase	,      {"x": 3, "y": 2}						)

        lowercase_			:List[str]    			= {"x": 8}
        lowercase_			:Tuple    			= evaluate(lowercase	,      {}	,      state=lowercase						)
        # evaluate returns the value of the last assignment.
        assert result == 5
        self.assertDictEqual(lowercase	,      {"x": 8, "y": 5}						)





     def 							lowercase__      ( self		:					Dict						):





        """simple docstring"""




        lowercase_			:Any    			= "test_list = [x, add_two(x)]"
        lowercase_			:Optional[int]    			= {"x": 3}
        lowercase_			:List[str]    			= evaluate(lowercase	,      {"add_two": add_two}	,      state=lowercase						)
        self.assertListEqual(lowercase	,      [3, 5]						)
        self.assertDictEqual(lowercase	,      {"x": 3, "test_list": [3, 5]}						)





     def 							lowercase__      ( self		:					List[Any]						):





        """simple docstring"""




        lowercase_			:str    			= "y = x"
        lowercase_			:Any    			= {"x": 3}
        lowercase_			:Tuple    			= evaluate(lowercase	,      {}	,      state=lowercase						)
        assert result == 3
        self.assertDictEqual(lowercase	,      {"x": 3, "y": 3}						)





     def 							lowercase__      ( self		:					Any						):





        """simple docstring"""




        lowercase_			:List[Any]    			= "test_list = [x, add_two(x)]\ntest_list[1]"
        lowercase_			:Dict    			= {"x": 3}
        lowercase_			:Any    			= evaluate(lowercase	,      {"add_two": add_two}	,      state=lowercase						)
        assert result == 5
        self.assertDictEqual(lowercase	,      {"x": 3, "test_list": [3, 5]}						)

        lowercase_			:Union[str, Any]    			= "test_dict = {'x': x, 'y': add_two(x)}\ntest_dict['y']"
        lowercase_			:List[Any]    			= {"x": 3}
        lowercase_			:List[Any]    			= evaluate(lowercase	,      {"add_two": add_two}	,      state=lowercase						)
        assert result == 5
        self.assertDictEqual(lowercase	,      {"x": 3, "test_dict": {"x": 3, "y": 5}}						)







     def 							lowercase__      ( self		:					Any						):





        """simple docstring"""




        lowercase_			:Dict    			= "x = 0\nfor i in range(3):\n    x = i"
        lowercase_			:Any    			= {}
        lowercase_			:Tuple    			= evaluate(lowercase	,      {"range": range}	,      state=lowercase						)
        assert result == 2
        self.assertDictEqual(lowercase	,      {"x": 2, "i": 2}						)


