



"""simple docstring"""






import os
import re
import shutil
import sys
import tempfile
import unittest

import black


_lowercase				:				Dict   		=		os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, "utils"))

import check_copies  # noqa: E402


# This is the reference code that will be used in the tests.
# If BertLMPredictionHead is changed in modeling_bert.py, this code needs to be manually updated.
_lowercase				:				Optional[int]   		=		"    def __init__(self, config):\n        super().__init__()\n        self.transform = BertPredictionHeadTransform(config)\n\n        # The output weights are the same as the input embeddings, but there is\n        # an output-only bias for each token.\n        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)\n\n        self.bias = nn.Parameter(torch.zeros(config.vocab_size))\n\n        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`\n        self.decoder.bias = self.bias\n\n    def forward(self, hidden_states):\n        hidden_states = self.transform(hidden_states)\n        hidden_states = self.decoder(hidden_states)\n        return hidden_states\n"

class  __SCREAMING_SNAKE_CASE			(   unittest.TestCase     ):


					'''simple docstring'''





					def 	snake_case ( self							:						List[str]  )->      Optional[Any]:
								lowerCamelCase__	:	Dict				=tempfile.mkdtemp()
								os.makedirs(os.path.join(self.transformer_dir,						'''models/bert/'''  )  )
								lowerCamelCase__	:	str				=self.transformer_dir
								shutil.copy(
								    os.path.join(lowercase__,						'''src/transformers/models/bert/modeling_bert.py'''  ),						os.path.join(self.transformer_dir,						'''models/bert/modeling_bert.py'''  ),						)



					def 	snake_case ( self							:						Tuple  )->      List[str]:
								lowerCamelCase__	:	Optional[Any]				="""src/transformers"""
								shutil.rmtree(self.transformer_dir  )



					def 	snake_case ( self							:						Any,						lowerCamelCase							:						int,						lowerCamelCase							:						Optional[Any],						lowerCamelCase							:						Optional[int],						lowerCamelCase							:						List[str]=None  )->      Optional[int]:
								lowerCamelCase__	:	Optional[Any]				=comment + F'''\nclass {class_name}(nn.Module):\n''' + class_code
								if overwrite_result is not None:
											lowerCamelCase__	:	List[Any]				=comment + F'''\nclass {class_name}(nn.Module):\n''' + overwrite_result
								lowerCamelCase__	:	str				=black.Mode(target_versions={black.TargetVersion.PYaa},						line_length=119  )
								lowerCamelCase__	:	Optional[int]				=black.format_str(lowercase__,						mode=lowercase__  )
								lowerCamelCase__	:	Any				=os.path.join(self.transformer_dir,						'''new_code.py'''  )
								with open(lowercase__,						'''w''',						newline='''\n'''  ) as f:
											f.write(lowercase__  )
								if overwrite_result is None:
											self.assertTrue(len(check_copies.is_copy_consistent(lowercase__  )  ) == 0  )
								else:
											check_copies.is_copy_consistent(f.name,						overwrite=lowercase__  )
											with open(lowercase__,						'''r'''  ) as f:
														self.assertTrue(f.read(),						lowercase__  )



					def 	snake_case ( self							:						List[Any]  )->      int:
								lowerCamelCase__	:	Optional[Any]				=check_copies.find_code_in_transformers('''models.bert.modeling_bert.BertLMPredictionHead'''  )
								self.assertEqual(lowercase__,						lowercase__  )



					def 	snake_case ( self							:						int  )->      List[Any]:
								# Base copy consistency
								self.check_copy_consistency(
								    '''# Copied from transformers.models.bert.modeling_bert.BertLMPredictionHead''',						'''BertLMPredictionHead''',						REFERENCE_CODE + '''\n''',						)

								# With no empty line at the end
								self.check_copy_consistency(
								    '''# Copied from transformers.models.bert.modeling_bert.BertLMPredictionHead''',						'''BertLMPredictionHead''',						lowercase__,						)

								# Copy consistency with rename
								self.check_copy_consistency(
								    '''# Copied from transformers.models.bert.modeling_bert.BertLMPredictionHead with Bert->TestModel''',						'''TestModelLMPredictionHead''',						re.sub('''Bert''',						'''TestModel''',						lowercase__  ),						)

								# Copy consistency with a really long name
								lowerCamelCase__	:	List[Any]				="""TestModelWithAReallyLongNameBecauseSomePeopleLikeThatForSomeReason"""
								self.check_copy_consistency(
								    F'''# Copied from transformers.models.bert.modeling_bert.BertLMPredictionHead with Bert->{long_class_name}''',						F'''{long_class_name}LMPredictionHead''',						re.sub('''Bert''',						lowercase__,						lowercase__  ),						)

								# Copy consistency with overwrite
								self.check_copy_consistency(
								    '''# Copied from transformers.models.bert.modeling_bert.BertLMPredictionHead with Bert->TestModel''',						'''TestModelLMPredictionHead''',						lowercase__,						overwrite_result=re.sub('''Bert''',						'''TestModel''',						lowercase__  ),						)




					def 	snake_case ( self							:						Any  )->      int:
								lowerCamelCase__	:	List[Any]				=check_copies.LOCALIZED_READMES["""README_zh-hans.md"""]

								lowerCamelCase__	:	List[Any]				=(
								    """1. **[ALBERT](https://huggingface.co/transformers/model_doc/albert.html)** (from Google Research and the"""
								    """ Toyota Technological Institute at Chicago) released with the paper [ALBERT: A Lite BERT for"""
								    """ Self-supervised Learning of Language Representations](https://arxiv.org/abs/1909.11942), by Zhenzhong"""
								    """ Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, Radu Soricut.\n1."""
								    """ **[DistilBERT](https://huggingface.co/transformers/model_doc/distilbert.html)** (from HuggingFace),"""
								    """ released together with the paper [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and"""
								    """ lighter](https://arxiv.org/abs/1910.01108) by Victor Sanh, Lysandre Debut and Thomas Wolf. The same"""
								    """ method has been applied to compress GPT2 into"""
								    """ [DistilGPT2](https://github.com/huggingface/transformers/tree/main/examples/distillation), RoBERTa into"""
								    """ [DistilRoBERTa](https://github.com/huggingface/transformers/tree/main/examples/distillation),"""
								    """ Multilingual BERT into"""
								    """ [DistilmBERT](https://github.com/huggingface/transformers/tree/main/examples/distillation) and a German"""
								    """ version of DistilBERT.\n1. **[ELECTRA](https://huggingface.co/transformers/model_doc/electra.html)**"""
								    """ (from Google Research/Stanford University) released with the paper [ELECTRA: Pre-training text encoders"""
								    """ as discriminators rather than generators](https://arxiv.org/abs/2003.10555) by Kevin Clark, Minh-Thang"""
								    """ Luong, Quoc V. Le, Christopher D. Manning."""
								)
								lowerCamelCase__	:	Optional[int]				=(
								    """1. **[ALBERT](https://huggingface.co/transformers/model_doc/albert.html)** (来自 Google Research and the"""
								    """ Toyota Technological Institute at Chicago) 伴随论文 [ALBERT: A Lite BERT for Self-supervised Learning of"""
								    """ Language Representations](https://arxiv.org/abs/1909.11942), 由 Zhenzhong Lan, Mingda Chen, Sebastian"""
								    """ Goodman, Kevin Gimpel, Piyush Sharma, Radu Soricut 发布。\n"""
								)
								lowerCamelCase__	:	Tuple				=(
								    """1. **[ALBERT](https://huggingface.co/transformers/model_doc/albert.html)** (来自 Google Research and the"""
								    """ Toyota Technological Institute at Chicago) 伴随论文 [ALBERT: A Lite BERT for Self-supervised Learning of"""
								    """ Language Representations](https://arxiv.org/abs/1909.11942), 由 Zhenzhong Lan, Mingda Chen, Sebastian"""
								    """ Goodman, Kevin Gimpel, Piyush Sharma, Radu Soricut 发布。\n1."""
								    """ **[DistilBERT](https://huggingface.co/transformers/model_doc/distilbert.html)** (来自 HuggingFace) 伴随论文"""
								    """ [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and"""
								    """ lighter](https://arxiv.org/abs/1910.01108) 由 Victor Sanh, Lysandre Debut and Thomas Wolf 发布。 The same"""
								    """ method has been applied to compress GPT2 into"""
								    """ [DistilGPT2](https://github.com/huggingface/transformers/tree/main/examples/distillation), RoBERTa into"""
								    """ [DistilRoBERTa](https://github.com/huggingface/transformers/tree/main/examples/distillation),"""
								    """ Multilingual BERT into"""
								    """ [DistilmBERT](https://github.com/huggingface/transformers/tree/main/examples/distillation) and a German"""
								    """ version of DistilBERT.\n1. **[ELECTRA](https://huggingface.co/transformers/model_doc/electra.html)** (来自"""
								    """ Google Research/Stanford University) 伴随论文 [ELECTRA: Pre-training text encoders as discriminators rather"""
								    """ than generators](https://arxiv.org/abs/2003.10555) 由 Kevin Clark, Minh-Thang Luong, Quoc V. Le,"""
								    """ Christopher D. Manning 发布。\n"""
								)

								lowerCamelCase__	:	Any				=check_copies.convert_to_localized_md(
								    lowercase__,						lowercase__,						localized_readme['''format_model_list''']  )

								self.assertFalse(lowercase__  )
								self.assertEqual(lowercase__,						lowercase__  )

								lowerCamelCase__	:	Tuple				=check_copies.convert_to_localized_md(
								    lowercase__,						lowercase__,						localized_readme['''format_model_list''']  )

								# Check whether the number of models is equal to README.md after conversion.
								self.assertTrue(lowercase__  )

								lowerCamelCase__	:	str				=(
								    """1. **[ALBERT](https://huggingface.co/transformers/model_doc/albert.html)** (from Google Research and the"""
								    """ Toyota Technological Institute at Chicago) released with the paper [ALBERT: A Lite BERT for"""
								    """ Self-supervised Learning of Language Representations](https://arxiv.org/abs/1909.11942), by Zhenzhong"""
								    """ Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, Radu Soricut."""
								)
								lowerCamelCase__	:	Union[str, Any]				=(
								    """1. **[ALBERT](https://huggingface.co/transformers/main/model_doc/albert.html)** (来自 Google Research and"""
								    """ the Toyota Technological Institute at Chicago) 伴随论文 [ALBERT: A Lite BERT for Self-supervised Learning of"""
								    """ Language Representations](https://arxiv.org/abs/1909.11942), 由 Zhenzhong Lan, Mingda Chen, Sebastian"""
								    """ Goodman, Kevin Gimpel, Piyush Sharma, Radu Soricut 发布。\n"""
								)
								lowerCamelCase__	:	str				=(
								    """1. **[ALBERT](https://huggingface.co/transformers/model_doc/albert.html)** (来自 Google Research and the"""
								    """ Toyota Technological Institute at Chicago) 伴随论文 [ALBERT: A Lite BERT for Self-supervised Learning of"""
								    """ Language Representations](https://arxiv.org/abs/1909.11942), 由 Zhenzhong Lan, Mingda Chen, Sebastian"""
								    """ Goodman, Kevin Gimpel, Piyush Sharma, Radu Soricut 发布。\n"""
								)

								lowerCamelCase__	:	Dict				=check_copies.convert_to_localized_md(
								    lowercase__,						lowercase__,						localized_readme['''format_model_list''']  )

								# Check if the model link is synchronized.
								self.assertEqual(lowercase__,						lowercase__  )









"""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  __SCREAMING_SNAKE_CASE			(   lowerCAmelCase_     ):


					'''simple docstring'''





					_a         =							42
					_a         =							42





class  __SCREAMING_SNAKE_CASE			(   nn.Module     ):


					'''simple docstring'''





					_a         =							42
					_a         =							(1_6, 3_2, 9_6, 2_5_6)
					_a         =							jnp.floataa



					def 	snake_case ( self							:						Tuple  )->      int:
								lowerCamelCase__	:	Tuple				=nn.Conv(
								    self.block_out_channels[0],						kernel_size=(3, 3),						padding=((1, 1), (1, 1)),						dtype=self.dtype,						)

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

								lowerCamelCase__	:	Optional[int]				=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							:						Any,						lowerCamelCase							:						int  )->      List[str]:
								lowerCamelCase__	:	Tuple				=self.conv_in(lowerCamelCase  )
								lowerCamelCase__	:	Dict				=nn.silu(lowerCamelCase  )

								for block in self.blocks:
											lowerCamelCase__	:	str				=block(lowerCamelCase  )
											lowerCamelCase__	:	List[str]				=nn.silu(lowerCamelCase  )

								lowerCamelCase__	:	Any				=self.conv_out(lowerCamelCase  )

								return embedding





@flax_register_to_config
class  __SCREAMING_SNAKE_CASE			(   nn.Module	,     lowerCAmelCase_	,     lowerCAmelCase_     ):


					'''simple docstring'''





					_a         =							3_2
					_a         =							4
					_a         =							(
					    "CrossAttnDownBlock2D",
					    "CrossAttnDownBlock2D",
					    "CrossAttnDownBlock2D",
					    "DownBlock2D",
					)
					_a         =							False
					_a         =							(3_2_0, 6_4_0, 1_2_8_0, 1_2_8_0)
					_a         =							2
					_a         =							8
					_a         =							None
					_a         =							1_2_8_0
					_a         =							0.0
					_a         =							False
					_a         =							jnp.floataa
					_a         =							True
					_a         =							0
					_a         =							"rgb"
					_a         =							(1_6, 3_2, 9_6, 2_5_6)



					def 	snake_case ( self							:						str,						lowerCamelCase							:						jax.random.KeyArray  )->      FrozenDict:
								# init input tensors
								lowerCamelCase__	:	int				=(1, self.in_channels, self.sample_size, self.sample_size)
								lowerCamelCase__	:	int				=jnp.zeros(lowerCamelCase,						dtype=jnp.floataa  )
								lowerCamelCase__	:	Union[str, Any]				=jnp.ones((1,),						dtype=jnp.intaa  )
								lowerCamelCase__	:	str				=jnp.zeros((1, 1, self.cross_attention_dim),						dtype=jnp.floataa  )
								lowerCamelCase__	:	Any				=(1, 3, self.sample_size * 8, self.sample_size * 8)
								lowerCamelCase__	:	Optional[Any]				=jnp.zeros(lowerCamelCase,						dtype=jnp.floataa  )

								lowerCamelCase__			,     lowerCamelCase__	:	List[Any]				=jax.random.split(lowerCamelCase  )
								lowerCamelCase__	:	Dict				={'''params''': params_rng, '''dropout''': dropout_rng}

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



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

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

								# time
								lowerCamelCase__	:	str				=FlaxTimesteps(
								    block_out_channels[0],						flip_sin_to_cos=self.flip_sin_to_cos,						freq_shift=self.config.freq_shift  )
								lowerCamelCase__	:	Dict				=FlaxTimestepEmbedding(lowerCamelCase,						dtype=self.dtype  )

								lowerCamelCase__	:	List[Any]				=FlaxControlNetConditioningEmbedding(
								    conditioning_embedding_channels=block_out_channels[0],						block_out_channels=self.conditioning_embedding_out_channels,						)

								lowerCamelCase__	:	Dict				=self.only_cross_attention
								if isinstance(lowerCamelCase,						lowerCamelCase  ):
											lowerCamelCase__	:	int				=(only_cross_attention,) * len(self.down_block_types  )

								if isinstance(lowerCamelCase,						lowerCamelCase  ):
											lowerCamelCase__	:	List[str]				=(num_attention_heads,) * len(self.down_block_types  )

								# down
								lowerCamelCase__	:	Union[str, Any]				=[]
								lowerCamelCase__	:	Dict				=[]

								lowerCamelCase__	:	List[Any]				=block_out_channels[0]

								lowerCamelCase__	:	List[Any]				=nn.Conv(
								    lowerCamelCase,						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(lowerCamelCase  )

								for i, down_block_type in enumerate(self.down_block_types  ):
											lowerCamelCase__	:	List[Any]				=output_channel
											lowerCamelCase__	:	str				=block_out_channels[i]
											lowerCamelCase__	:	Dict				=i == len(lowerCamelCase  ) - 1

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

											down_blocks.append(lowerCamelCase  )

											for _ in range(self.layers_per_block  ):
														lowerCamelCase__	:	Any				=nn.Conv(
														    lowerCamelCase,						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(lowerCamelCase  )

											if not is_final_block:
														lowerCamelCase__	:	Any				=nn.Conv(
														    lowerCamelCase,						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(lowerCamelCase  )

								lowerCamelCase__	:	int				=down_blocks
								lowerCamelCase__	:	List[str]				=controlnet_down_blocks

								# mid
								lowerCamelCase__	:	Tuple				=block_out_channels[-1]
								lowerCamelCase__	:	List[Any]				=FlaxUNetMidBlockaDCrossAttn(
								    in_channels=lowerCamelCase,						dropout=self.dropout,						num_attention_heads=num_attention_heads[-1],						use_linear_projection=self.use_linear_projection,						dtype=self.dtype,						)

								lowerCamelCase__	:	List[str]				=nn.Conv(
								    lowerCamelCase,						kernel_size=(1, 1),						padding='''VALID''',						kernel_init=nn.initializers.zeros_init(),						bias_init=nn.initializers.zeros_init(),						dtype=self.dtype,						)




					def __call__( self							:						int,						lowerCamelCase							:						List[Any],						lowerCamelCase							:						Tuple,						lowerCamelCase							:						Union[str, Any],						lowerCamelCase							:						str,						lowerCamelCase							:						float = 1.0,						lowerCamelCase							:						bool = True,						lowerCamelCase							:						bool = False,						)->      Union[FlaxControlNetOutput, Tuple]:
								lowerCamelCase__	:	int				=self.controlnet_conditioning_channel_order
								if channel_order == "bgr":
											lowerCamelCase__	:	int				=jnp.flip(lowerCamelCase,						axis=1  )

								# 1. time
								if not isinstance(lowerCamelCase,						jnp.ndarray  ):
											lowerCamelCase__	:	Any				=jnp.array([timesteps],						dtype=jnp.intaa  )
								elif isinstance(lowerCamelCase,						jnp.ndarray  ) and len(timesteps.shape  ) == 0:
											lowerCamelCase__	:	List[str]				=timesteps.astype(dtype=jnp.floataa  )
											lowerCamelCase__	:	int				=jnp.expand_dims(lowerCamelCase,						0  )

								lowerCamelCase__	:	Optional[Any]				=self.time_proj(lowerCamelCase  )
								lowerCamelCase__	:	Optional[Any]				=self.time_embedding(lowerCamelCase  )

								# 2. pre-process
								lowerCamelCase__	:	Optional[int]				=jnp.transpose(lowerCamelCase,						(0, 2, 3, 1)  )
								lowerCamelCase__	:	Dict				=self.conv_in(lowerCamelCase  )

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

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

								# 4. mid
								lowerCamelCase__	:	Optional[int]				=self.mid_block(lowerCamelCase,						lowerCamelCase,						lowerCamelCase,						deterministic=not train  )

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

								lowerCamelCase__	:	List[str]				=controlnet_down_block_res_samples

								lowerCamelCase__	:	List[str]				=self.controlnet_mid_block(lowerCamelCase  )

								# 6. scaling
								lowerCamelCase__	:	Union[str, Any]				=[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=lowerCamelCase,						mid_block_res_sample=lowerCamelCase  )




