


'''simple docstring'''






import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion_safe import StableDiffusionPipelineSafe as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import require_torch_gpu






class  __magic_name__							(		unittest.TestCase    ):






       def       lowerCAmelCase		(      self)						->  Optional[Any]:




          '''simple docstring'''
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()






       @property
       def       lowerCAmelCase		(      self)						->  Any:




          '''simple docstring'''
          _UpperCAmelCase       :       Dict							   =1
          _UpperCAmelCase       :       Dict							   =3
          _UpperCAmelCase       :       Dict							   =(3_2, 3_2)

          _UpperCAmelCase       :       int							   =floats_tensor((batch_size, num_channels) + sizes  ,							rng=random.Random(0)).to(snake_case)
          return image






       @property
       def       lowerCAmelCase		(      self)						->  Optional[int]:




          '''simple docstring'''
          torch.manual_seed(0)
          _UpperCAmelCase       :       str							   =UNetaDConditionModel(
              block_out_channels=(3_2, 6_4)  ,							layers_per_block=2  ,							sample_size=3_2  ,							in_channels=4  ,							out_channels=4  ,							down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')  ,							up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')  ,							cross_attention_dim=3_2  ,							)
          return model






       @property
       def       lowerCAmelCase		(      self)						->  List[Any]:




          '''simple docstring'''
          torch.manual_seed(0)
          _UpperCAmelCase       :       Optional[Any]							   =AutoencoderKL(
              block_out_channels=[3_2, 6_4]  ,							in_channels=3  ,							out_channels=3  ,							down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']  ,							up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']  ,							latent_channels=4  ,							)
          return model






       @property
       def       lowerCAmelCase		(      self)						->  Tuple:




          '''simple docstring'''
          torch.manual_seed(0)
          _UpperCAmelCase       :       int							   =CLIPTextConfig(
              bos_token_id=0  ,							eos_token_id=2  ,							hidden_size=3_2  ,							intermediate_size=3_7  ,							layer_norm_eps=1E-0_5  ,							num_attention_heads=4  ,							num_hidden_layers=5  ,							pad_token_id=1  ,							vocab_size=1_0_0_0  ,							)
          return CLIPTextModel(snake_case)






       @property
       def       lowerCAmelCase		(      self)						->  Union[str, Any]:




          '''simple docstring'''
          def extract(*snake_case  ,							**snake_case):






             class  __magic_name__							:






                    def __init__(      self)						->  Any:




                       '''simple docstring'''
                       _UpperCAmelCase       :       Optional[int]							   =torch.ones([0])






                    def       lowerCAmelCase		(      self  ,							snake_case)						->  Tuple:




                       '''simple docstring'''
                       self.pixel_values.to(snake_case)
                       return self






             return Out()

          return extract






       def       lowerCAmelCase		(      self)						->  Optional[int]:




          '''simple docstring'''
          _UpperCAmelCase       :       str							   ='cpu'  # ensure determinism for the device-dependent torch.Generator
          _UpperCAmelCase       :       Any							   =self.dummy_cond_unet
          _UpperCAmelCase       :       Union[str, Any]							   =DDIMScheduler(
              beta_start=0.0_00_85  ,							beta_end=0.0_12  ,							beta_schedule='scaled_linear'  ,							clip_sample=snake_case  ,							set_alpha_to_one=snake_case  ,							)

          _UpperCAmelCase       :       str							   =self.dummy_vae
          _UpperCAmelCase       :       Dict							   =self.dummy_text_encoder
          _UpperCAmelCase       :       str							   =CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')

          # make sure here that pndm scheduler skips prk
          _UpperCAmelCase       :       int							   =StableDiffusionPipeline(
              unet=snake_case  ,							scheduler=snake_case  ,							vae=snake_case  ,							text_encoder=snake_case  ,							tokenizer=snake_case  ,							safety_checker=snake_case  ,							feature_extractor=self.dummy_extractor  ,							)
          _UpperCAmelCase       :       Union[str, Any]							   =sd_pipe.to(snake_case)
          sd_pipe.set_progress_bar_config(disable=snake_case)

          _UpperCAmelCase       :       Any							   ='A painting of a squirrel eating a burger'

          _UpperCAmelCase       :       Optional[int]							   =torch.Generator(device=snake_case).manual_seed(0)
          _UpperCAmelCase       :       Optional[int]							   =sd_pipe([prompt]  ,							generator=snake_case  ,							guidance_scale=6.0  ,							num_inference_steps=2  ,							output_type='np')
          _UpperCAmelCase       :       List[str]							   =output.images

          _UpperCAmelCase       :       int							   =torch.Generator(device=snake_case).manual_seed(0)
          _UpperCAmelCase       :       int							   =sd_pipe(
              [prompt]  ,							generator=snake_case  ,							guidance_scale=6.0  ,							num_inference_steps=2  ,							output_type='np'  ,							return_dict=snake_case  ,							)[0]

          _UpperCAmelCase       :       List[Any]							   =image[0, -3:, -3:, -1]
          _UpperCAmelCase       :       Union[str, Any]							   =image_from_tuple[0, -3:, -3:, -1]

          assert image.shape == (1, 6_4, 6_4, 3)
          _UpperCAmelCase       :       str							   =np.array([0.57_56, 0.61_18, 0.50_05, 0.50_41, 0.54_71, 0.47_26, 0.49_76, 0.48_65, 0.48_64])

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2
          assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1E-2






       def       lowerCAmelCase		(      self)						->  str:




          '''simple docstring'''
          _UpperCAmelCase       :       int							   ='cpu'  # ensure determinism for the device-dependent torch.Generator
          _UpperCAmelCase       :       Tuple							   =self.dummy_cond_unet
          _UpperCAmelCase       :       int							   =PNDMScheduler(skip_prk_steps=snake_case)
          _UpperCAmelCase       :       Tuple							   =self.dummy_vae
          _UpperCAmelCase       :       List[Any]							   =self.dummy_text_encoder
          _UpperCAmelCase       :       List[str]							   =CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')

          # make sure here that pndm scheduler skips prk
          _UpperCAmelCase       :       Dict							   =StableDiffusionPipeline(
              unet=snake_case  ,							scheduler=snake_case  ,							vae=snake_case  ,							text_encoder=snake_case  ,							tokenizer=snake_case  ,							safety_checker=snake_case  ,							feature_extractor=self.dummy_extractor  ,							)
          _UpperCAmelCase       :       List[Any]							   =sd_pipe.to(snake_case)
          sd_pipe.set_progress_bar_config(disable=snake_case)

          _UpperCAmelCase       :       Optional[Any]							   ='A painting of a squirrel eating a burger'
          _UpperCAmelCase       :       Union[str, Any]							   =torch.Generator(device=snake_case).manual_seed(0)
          _UpperCAmelCase       :       Tuple							   =sd_pipe([prompt]  ,							generator=snake_case  ,							guidance_scale=6.0  ,							num_inference_steps=2  ,							output_type='np')

          _UpperCAmelCase       :       int							   =output.images

          _UpperCAmelCase       :       str							   =torch.Generator(device=snake_case).manual_seed(0)
          _UpperCAmelCase       :       List[str]							   =sd_pipe(
              [prompt]  ,							generator=snake_case  ,							guidance_scale=6.0  ,							num_inference_steps=2  ,							output_type='np'  ,							return_dict=snake_case  ,							)[0]

          _UpperCAmelCase       :       Union[str, Any]							   =image[0, -3:, -3:, -1]
          _UpperCAmelCase       :       Tuple							   =image_from_tuple[0, -3:, -3:, -1]

          assert image.shape == (1, 6_4, 6_4, 3)
          _UpperCAmelCase       :       Tuple							   =np.array([0.51_25, 0.57_16, 0.48_28, 0.50_60, 0.56_50, 0.47_68, 0.51_85, 0.48_95, 0.49_93])

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2
          assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1E-2






       def       lowerCAmelCase		(      self)						->  Optional[int]:




          '''simple docstring'''
          _UpperCAmelCase       :       Optional[Any]							   =StableDiffusionPipeline.from_pretrained(
              'hf-internal-testing/tiny-stable-diffusion-lms-pipe'  ,							safety_checker=snake_case)
          assert isinstance(snake_case  ,							snake_case)
          assert isinstance(pipe.scheduler  ,							snake_case)
          assert pipe.safety_checker is None

          _UpperCAmelCase       :       Union[str, Any]							   =pipe('example prompt'  ,							num_inference_steps=2).images[0]
          assert image is not None

          # check that there's no error when saving a pipeline with one of the models being None
          with tempfile.TemporaryDirectory() as tmpdirname:
             pipe.save_pretrained(snake_case)
             _UpperCAmelCase       :       Dict							   =StableDiffusionPipeline.from_pretrained(snake_case)

          # sanity check that the pipeline still works
          assert pipe.safety_checker is None
          _UpperCAmelCase       :       List[str]							   =pipe('example prompt'  ,							num_inference_steps=2).images[0]
          assert image is not None






       @unittest.skipIf(torch_device != 'cuda'  ,							'This test requires a GPU')
       def       lowerCAmelCase		(      self)						->  Any:




          '''simple docstring'''
          _UpperCAmelCase       :       str							   =self.dummy_cond_unet
          _UpperCAmelCase       :       str							   =PNDMScheduler(skip_prk_steps=snake_case)
          _UpperCAmelCase       :       Optional[Any]							   =self.dummy_vae
          _UpperCAmelCase       :       str							   =self.dummy_text_encoder
          _UpperCAmelCase       :       Optional[int]							   =CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')

          # put models in fp16
          _UpperCAmelCase       :       List[str]							   =unet.half()
          _UpperCAmelCase       :       List[str]							   =vae.half()
          _UpperCAmelCase       :       Union[str, Any]							   =bert.half()

          # make sure here that pndm scheduler skips prk
          _UpperCAmelCase       :       Dict							   =StableDiffusionPipeline(
              unet=snake_case  ,							scheduler=snake_case  ,							vae=snake_case  ,							text_encoder=snake_case  ,							tokenizer=snake_case  ,							safety_checker=snake_case  ,							feature_extractor=self.dummy_extractor  ,							)
          _UpperCAmelCase       :       Optional[Any]							   =sd_pipe.to(snake_case)
          sd_pipe.set_progress_bar_config(disable=snake_case)

          _UpperCAmelCase       :       List[Any]							   ='A painting of a squirrel eating a burger'
          _UpperCAmelCase       :       Union[str, Any]							   =sd_pipe([prompt]  ,							num_inference_steps=2  ,							output_type='np').images

          assert image.shape == (1, 6_4, 6_4, 3)






@nightly
@require_torch_gpu
class  __magic_name__							(		unittest.TestCase    ):






       def       lowerCAmelCase		(      self)						->  Optional[int]:




          '''simple docstring'''
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()






       def       lowerCAmelCase		(      self)						->  Optional[int]:




          '''simple docstring'''
          _UpperCAmelCase       :       int							   =StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'  ,							safety_checker=snake_case)
          _UpperCAmelCase       :       List[str]							   =LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config)
          _UpperCAmelCase       :       Any							   =sd_pipe.to(snake_case)
          sd_pipe.set_progress_bar_config(disable=snake_case)

          _UpperCAmelCase       :       Any							   =(
              'portrait of girl with smokey eyes makeup in abandoned hotel, grange clothes, redshift, wide high angle'
              ' coloured polaroid photograph with flash, kodak film, hyper real, stunning moody cinematography, with'
              ' anamorphic lenses, by maripol, fallen angels by wong kar - wai, style of suspiria and neon demon and'
              ' children from bahnhof zoo, detailed '
          )
          _UpperCAmelCase       :       str							   =4_0_0_3_6_6_0_3_4_6
          _UpperCAmelCase       :       Optional[int]							   =7

          # without safety guidance (sld_guidance_scale = 0)
          _UpperCAmelCase       :       List[Any]							   =torch.manual_seed(snake_case)
          _UpperCAmelCase       :       int							   =sd_pipe(
              [prompt]  ,							generator=snake_case  ,							guidance_scale=snake_case  ,							num_inference_steps=5_0  ,							output_type='np'  ,							width=5_1_2  ,							height=5_1_2  ,							sld_guidance_scale=0  ,							)

          _UpperCAmelCase       :       Dict							   =output.images
          _UpperCAmelCase       :       Dict							   =image[0, -3:, -3:, -1]
          _UpperCAmelCase       :       Dict							   =[0.22_78, 0.22_31, 0.22_49, 0.23_33, 0.23_03, 0.18_85, 0.22_73, 0.21_44, 0.21_76]

          assert image.shape == (1, 5_1_2, 5_1_2, 3)

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2

          # without safety guidance (strong configuration)
          _UpperCAmelCase       :       List[Any]							   =torch.manual_seed(snake_case)
          _UpperCAmelCase       :       str							   =sd_pipe(
              [prompt]  ,							generator=snake_case  ,							guidance_scale=snake_case  ,							num_inference_steps=5_0  ,							output_type='np'  ,							width=5_1_2  ,							height=5_1_2  ,							sld_guidance_scale=2_0_0_0  ,							sld_warmup_steps=7  ,							sld_threshold=0.0_25  ,							sld_momentum_scale=0.5  ,							sld_mom_beta=0.7  ,							)

          _UpperCAmelCase       :       Tuple							   =output.images
          _UpperCAmelCase       :       Optional[Any]							   =image[0, -3:, -3:, -1]
          _UpperCAmelCase       :       str							   =[0.23_83, 0.22_76, 0.2_36, 0.21_92, 0.21_86, 0.20_53, 0.19_71, 0.19_01, 0.17_19]

          assert image.shape == (1, 5_1_2, 5_1_2, 3)

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2






       def       lowerCAmelCase		(      self)						->  List[str]:




          '''simple docstring'''
          _UpperCAmelCase       :       str							   =StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'  ,							safety_checker=snake_case)
          _UpperCAmelCase       :       int							   =LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config)
          _UpperCAmelCase       :       List[Any]							   =sd_pipe.to(snake_case)
          sd_pipe.set_progress_bar_config(disable=snake_case)

          _UpperCAmelCase       :       str							   ='padme amidala taking a bath artwork, safe for work, no nudity'
          _UpperCAmelCase       :       Tuple							   =2_7_3_4_9_7_1_7_5_5
          _UpperCAmelCase       :       Union[str, Any]							   =7

          _UpperCAmelCase       :       List[Any]							   =torch.manual_seed(snake_case)
          _UpperCAmelCase       :       int							   =sd_pipe(
              [prompt]  ,							generator=snake_case  ,							guidance_scale=snake_case  ,							num_inference_steps=5_0  ,							output_type='np'  ,							width=5_1_2  ,							height=5_1_2  ,							sld_guidance_scale=0  ,							)

          _UpperCAmelCase       :       Any							   =output.images
          _UpperCAmelCase       :       Any							   =image[0, -3:, -3:, -1]
          _UpperCAmelCase       :       Optional[int]							   =[0.35_02, 0.36_22, 0.33_96, 0.36_42, 0.34_78, 0.33_18, 0.35, 0.33_48, 0.32_97]

          assert image.shape == (1, 5_1_2, 5_1_2, 3)

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2

          _UpperCAmelCase       :       Optional[int]							   =torch.manual_seed(snake_case)
          _UpperCAmelCase       :       Dict							   =sd_pipe(
              [prompt]  ,							generator=snake_case  ,							guidance_scale=snake_case  ,							num_inference_steps=5_0  ,							output_type='np'  ,							width=5_1_2  ,							height=5_1_2  ,							sld_guidance_scale=2_0_0_0  ,							sld_warmup_steps=7  ,							sld_threshold=0.0_25  ,							sld_momentum_scale=0.5  ,							sld_mom_beta=0.7  ,							)

          _UpperCAmelCase       :       Dict							   =output.images
          _UpperCAmelCase       :       List[Any]							   =image[0, -3:, -3:, -1]
          _UpperCAmelCase       :       int							   =[0.55_31, 0.52_06, 0.48_95, 0.51_56, 0.51_82, 0.47_51, 0.48_02, 0.48_03, 0.44_43]

          assert image.shape == (1, 5_1_2, 5_1_2, 3)

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2






       def       lowerCAmelCase		(      self)						->  Tuple:




          '''simple docstring'''
          _UpperCAmelCase       :       int							   =StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5')
          _UpperCAmelCase       :       Any							   =sd_pipe.to(snake_case)
          sd_pipe.set_progress_bar_config(disable=snake_case)

          _UpperCAmelCase       :       Optional[int]							   =(
              'the four horsewomen of the apocalypse, painting by tom of finland, gaston bussiere, craig mullins, j. c.'
              ' leyendecker'
          )
          _UpperCAmelCase       :       Dict							   =1_0_4_4_3_5_5_2_3_4
          _UpperCAmelCase       :       int							   =1_2

          _UpperCAmelCase       :       Union[str, Any]							   =torch.manual_seed(snake_case)
          _UpperCAmelCase       :       List[str]							   =sd_pipe(
              [prompt]  ,							generator=snake_case  ,							guidance_scale=snake_case  ,							num_inference_steps=5_0  ,							output_type='np'  ,							width=5_1_2  ,							height=5_1_2  ,							sld_guidance_scale=0  ,							)

          _UpperCAmelCase       :       str							   =output.images
          _UpperCAmelCase       :       Optional[int]							   =image[0, -3:, -3:, -1]
          _UpperCAmelCase       :       Tuple							   =np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])

          assert image.shape == (1, 5_1_2, 5_1_2, 3)

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-7

          _UpperCAmelCase       :       str							   =torch.manual_seed(snake_case)
          _UpperCAmelCase       :       Optional[int]							   =sd_pipe(
              [prompt]  ,							generator=snake_case  ,							guidance_scale=snake_case  ,							num_inference_steps=5_0  ,							output_type='np'  ,							width=5_1_2  ,							height=5_1_2  ,							sld_guidance_scale=2_0_0_0  ,							sld_warmup_steps=7  ,							sld_threshold=0.0_25  ,							sld_momentum_scale=0.5  ,							sld_mom_beta=0.7  ,							)

          _UpperCAmelCase       :       Tuple							   =output.images
          _UpperCAmelCase       :       List[str]							   =image[0, -3:, -3:, -1]
          _UpperCAmelCase       :       List[str]							   =np.array([0.58_18, 0.62_85, 0.68_35, 0.60_19, 0.6_25, 0.67_54, 0.60_96, 0.63_34, 0.65_61])
          assert image.shape == (1, 5_1_2, 5_1_2, 3)

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2





'''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  __magic_name__							(		lowerCAmelCase    ):
       UpperCAmelCase										=42
       UpperCAmelCase										=42






class  __magic_name__							(		nn.Module    ):
       UpperCAmelCase										=42
       UpperCAmelCase										=(1_6, 3_2, 9_6, 2_5_6)
       UpperCAmelCase										=jnp.floataa






       def       lowerCAmelCase		(      self)						->  Optional[int]:




          '''simple docstring'''
          _UpperCAmelCase       :       str							   =nn.Conv(
              self.block_out_channels[0]  ,							kernel_size=(3, 3)  ,							padding=((1, 1), (1, 1))  ,							dtype=self.dtype  ,							)

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

          _UpperCAmelCase       :       Tuple							   =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  ,							snake_case)						->  List[str]:




          '''simple docstring'''
          _UpperCAmelCase       :       int							   =self.conv_in(snake_case)
          _UpperCAmelCase       :       Any							   =nn.silu(snake_case)

          for block in self.blocks:
             _UpperCAmelCase       :       Optional[Any]							   =block(snake_case)
             _UpperCAmelCase       :       Union[str, Any]							   =nn.silu(snake_case)

          _UpperCAmelCase       :       str							   =self.conv_out(snake_case)

          return embedding






@flax_register_to_config
class  __magic_name__							(		nn.Module ,lowerCAmelCase ,lowerCAmelCase    ):
       UpperCAmelCase										=3_2
       UpperCAmelCase										=4
       UpperCAmelCase										=(
           "CrossAttnDownBlock2D",
           "CrossAttnDownBlock2D",
           "CrossAttnDownBlock2D",
           "DownBlock2D",
       )
       UpperCAmelCase										=False
       UpperCAmelCase										=(3_2_0, 6_4_0, 1_2_8_0, 1_2_8_0)
       UpperCAmelCase										=2
       UpperCAmelCase										=8
       UpperCAmelCase										=None
       UpperCAmelCase										=1_2_8_0
       UpperCAmelCase										=0.0
       UpperCAmelCase										=False
       UpperCAmelCase										=jnp.floataa
       UpperCAmelCase										=True
       UpperCAmelCase										=0
       UpperCAmelCase										="rgb"
       UpperCAmelCase										=(1_6, 3_2, 9_6, 2_5_6)






       def       lowerCAmelCase		(      self  ,							snake_case)						->  FrozenDict:




          '''simple docstring'''
          # init input tensors
          _UpperCAmelCase       :       Any							   =(1, self.in_channels, self.sample_size, self.sample_size)
          _UpperCAmelCase       :       Optional[Any]							   =jnp.zeros(snake_case  ,							dtype=jnp.floataa)
          _UpperCAmelCase       :       Optional[int]							   =jnp.ones((1,)  ,							dtype=jnp.intaa)
          _UpperCAmelCase       :       str							   =jnp.zeros((1, 1, self.cross_attention_dim)  ,							dtype=jnp.floataa)
          _UpperCAmelCase       :       Optional[Any]							   =(1, 3, self.sample_size * 8, self.sample_size * 8)
          _UpperCAmelCase       :       int							   =jnp.zeros(snake_case  ,							dtype=jnp.floataa)

          _UpperCAmelCase   ,	_UpperCAmelCase       :       List[Any]							   =jax.random.split(snake_case)
          _UpperCAmelCase       :       str							   ={'params': params_rng, 'dropout': dropout_rng}

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






       def       lowerCAmelCase		(      self)						->  Tuple:




          '''simple docstring'''
          _UpperCAmelCase       :       Optional[Any]							   =self.block_out_channels
          _UpperCAmelCase       :       Tuple							   =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.
          _UpperCAmelCase       :       Optional[Any]							   =self.num_attention_heads or self.attention_head_dim

          # input
          _UpperCAmelCase       :       Tuple							   =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       :       str							   =FlaxTimestepEmbedding(snake_case  ,							dtype=self.dtype)

          _UpperCAmelCase       :       Optional[Any]							   =FlaxControlNetConditioningEmbedding(
              conditioning_embedding_channels=block_out_channels[0]  ,							block_out_channels=self.conditioning_embedding_out_channels  ,							)

          _UpperCAmelCase       :       Optional[int]							   =self.only_cross_attention
          if isinstance(snake_case  ,							snake_case):
             _UpperCAmelCase       :       Dict							   =(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       :       int							   =[]
          _UpperCAmelCase       :       Optional[int]							   =[]

          _UpperCAmelCase       :       List[str]							   =block_out_channels[0]

          _UpperCAmelCase       :       int							   =nn.Conv(
              snake_case  ,							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(snake_case)

          for i, down_block_type in enumerate(self.down_block_types):
             _UpperCAmelCase       :       Tuple							   =output_channel
             _UpperCAmelCase       :       Dict							   =block_out_channels[i]
             _UpperCAmelCase       :       str							   =i == len(snake_case) - 1

             if down_block_type == "CrossAttnDownBlock2D":
                _UpperCAmelCase       :       Tuple							   =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]  ,							dtype=self.dtype  ,							)
             else:
                _UpperCAmelCase       :       Optional[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)

             for _ in range(self.layers_per_block):
                _UpperCAmelCase       :       Tuple							   =nn.Conv(
                    snake_case  ,							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(snake_case)

             if not is_final_block:
                _UpperCAmelCase       :       List[str]							   =nn.Conv(
                    snake_case  ,							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(snake_case)

          _UpperCAmelCase       :       List[Any]							   =down_blocks
          _UpperCAmelCase       :       Optional[Any]							   =controlnet_down_blocks

          # mid
          _UpperCAmelCase       :       int							   =block_out_channels[-1]
          _UpperCAmelCase       :       Optional[Any]							   =FlaxUNetMidBlockaDCrossAttn(
              in_channels=snake_case  ,							dropout=self.dropout  ,							num_attention_heads=num_attention_heads[-1]  ,							use_linear_projection=self.use_linear_projection  ,							dtype=self.dtype  ,							)

          _UpperCAmelCase       :       Optional[int]							   =nn.Conv(
              snake_case  ,							kernel_size=(1, 1)  ,							padding='VALID'  ,							kernel_init=nn.initializers.zeros_init()  ,							bias_init=nn.initializers.zeros_init()  ,							dtype=self.dtype  ,							)






       def __call__(      self  ,							snake_case  ,							snake_case  ,							snake_case  ,							snake_case  ,							snake_case = 1.0  ,							snake_case = True  ,							snake_case = False  ,							)						->  Union[FlaxControlNetOutput, Tuple]:




          '''simple docstring'''
          _UpperCAmelCase       :       Union[str, Any]							   =self.controlnet_conditioning_channel_order
          if channel_order == "bgr":
             _UpperCAmelCase       :       Optional[int]							   =jnp.flip(snake_case  ,							axis=1)

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

          _UpperCAmelCase       :       int							   =self.time_proj(snake_case)
          _UpperCAmelCase       :       Any							   =self.time_embedding(snake_case)

          # 2. pre-process
          _UpperCAmelCase       :       str							   =jnp.transpose(snake_case  ,							(0, 2, 3, 1))
          _UpperCAmelCase       :       Any							   =self.conv_in(snake_case)

          _UpperCAmelCase       :       List[str]							   =jnp.transpose(snake_case  ,							(0, 2, 3, 1))
          _UpperCAmelCase       :       Optional[int]							   =self.controlnet_cond_embedding(snake_case)
          sample += controlnet_cond

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

          # 4. mid
          _UpperCAmelCase       :       List[Any]							   =self.mid_block(snake_case  ,							snake_case  ,							snake_case  ,							deterministic=not train)

          # 5. contronet blocks
          _UpperCAmelCase       :       Union[str, Any]							   =()
          for down_block_res_sample, controlnet_block in zip(snake_case  ,							self.controlnet_down_blocks):
             _UpperCAmelCase       :       List[str]							   =controlnet_block(snake_case)
             controlnet_down_block_res_samples += (down_block_res_sample,)

          _UpperCAmelCase       :       Optional[int]							   =controlnet_down_block_res_samples

          _UpperCAmelCase       :       List[str]							   =self.controlnet_mid_block(snake_case)

          # 6. scaling
          _UpperCAmelCase       :       Tuple							   =[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=snake_case  ,							mid_block_res_sample=snake_case)

