





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 						lowerCamelCase       (__lowerCamelCase      ):







							"""simple docstring"""




							UpperCAmelCase_      						=   42
							UpperCAmelCase_      						=   42


class 						lowerCamelCase       (nn.Module      ):







							"""simple docstring"""




							UpperCAmelCase_      						=   42
							UpperCAmelCase_      						=   (16, 32, 96, 256)
							UpperCAmelCase_      						=   jnp.floataa






							def  A_     ( self     :					Optional[Any]						)				->			List[str]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Dict     =						nn.Conv(
														    self.block_out_channels[0],			kernel_size=(3, 3),			padding=((1, 1), (1, 1)),			dtype=self.dtype,			)

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

														SCREAMING_SNAKE_CASE__       :      Dict     =						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[Any],			_UpperCAmelCase     :					Dict						)				->			List[Any]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      List[str]     =						self.conv_in(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						nn.silu(_UpperCAmelCase						)

														for block in self.blocks:
																					SCREAMING_SNAKE_CASE__       :      Dict     =						block(_UpperCAmelCase						)
																					SCREAMING_SNAKE_CASE__       :      List[Any]     =						nn.silu(_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      Dict     =						self.conv_out(_UpperCAmelCase						)

														return embedding


@flax_register_to_config
class 						lowerCamelCase       (nn.Module							,     __lowerCamelCase							,     __lowerCamelCase      ):







							"""simple docstring"""




							UpperCAmelCase_      						=   32
							UpperCAmelCase_      						=   4
							UpperCAmelCase_      						=   (
							    "CrossAttnDownBlock2D",
							    "CrossAttnDownBlock2D",
							    "CrossAttnDownBlock2D",
							    "DownBlock2D",
							)
							UpperCAmelCase_      						=   False
							UpperCAmelCase_      						=   (320, 640, 1280, 1280)
							UpperCAmelCase_      						=   2
							UpperCAmelCase_      						=   8
							UpperCAmelCase_      						=   None
							UpperCAmelCase_      						=   1280
							UpperCAmelCase_      						=   0.0
							UpperCAmelCase_      						=   False
							UpperCAmelCase_      						=   jnp.floataa
							UpperCAmelCase_      						=   True
							UpperCAmelCase_      						=   0
							UpperCAmelCase_      						=   "rgb"
							UpperCAmelCase_      						=   (16, 32, 96, 256)






							def  A_     ( self     :					str,			_UpperCAmelCase     :					jax.random.KeyArray						)				->			FrozenDict:





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

														SCREAMING_SNAKE_CASE__   ,SCREAMING_SNAKE_CASE__       :      Tuple     =						jax.random.split(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						{"params": params_rng, "dropout": dropout_rng}

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






							def  A_     ( self     :					List[str]						)				->			int:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						self.block_out_channels
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						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.
														SCREAMING_SNAKE_CASE__       :      int     =						self.num_attention_heads or self.attention_head_dim

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

														# time
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						FlaxTimesteps(
														    block_out_channels[0],			flip_sin_to_cos=self.flip_sin_to_cos,			freq_shift=self.config.freq_shift						)
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						FlaxTimestepEmbedding(_UpperCAmelCase,			dtype=self.dtype						)

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

														SCREAMING_SNAKE_CASE__       :      Dict     =						self.only_cross_attention
														if isinstance(_UpperCAmelCase,			_UpperCAmelCase						):
																					SCREAMING_SNAKE_CASE__       :      Tuple     =						(only_cross_attention,) * len(self.down_block_types						)

														if isinstance(_UpperCAmelCase,			_UpperCAmelCase						):
																					SCREAMING_SNAKE_CASE__       :      Optional[int]     =						(num_attention_heads,) * len(self.down_block_types						)

														# down
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						[]
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						[]

														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						block_out_channels[0]

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

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

																					if down_block_type == "CrossAttnDownBlock2D":
																												SCREAMING_SNAKE_CASE__       :      int     =						FlaxCrossAttnDownBlockaD(
																												    in_channels=_UpperCAmelCase,			out_channels=_UpperCAmelCase,			dropout=self.dropout,			num_layers=self.layers_per_block,			num_attention_heads=num_attention_heads[i],			add_downsample=not is_final_block,			use_linear_projection=self.use_linear_projection,			only_cross_attention=only_cross_attention[i],			dtype=self.dtype,			)
																					else:
																												SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						FlaxDownBlockaD(
																												    in_channels=_UpperCAmelCase,			out_channels=_UpperCAmelCase,			dropout=self.dropout,			num_layers=self.layers_per_block,			add_downsample=not is_final_block,			dtype=self.dtype,			)

																					down_blocks.append(_UpperCAmelCase						)

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

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

														SCREAMING_SNAKE_CASE__       :      int     =						down_blocks
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						controlnet_down_blocks

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

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






							def __call__( self     :					Dict,			_UpperCAmelCase     :					Optional[int],			_UpperCAmelCase     :					Any,			_UpperCAmelCase     :					Optional[Any],			_UpperCAmelCase     :					Tuple,			_UpperCAmelCase     :					float = 1.0,			_UpperCAmelCase     :					bool = True,			_UpperCAmelCase     :					bool = False,			)				->			Union[FlaxControlNetOutput, Tuple]:





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

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

														SCREAMING_SNAKE_CASE__       :      List[str]     =						self.time_proj(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						self.time_embedding(_UpperCAmelCase						)

														# 2. pre-process
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						jnp.transpose(_UpperCAmelCase,			(0, 2, 3, 1)						)
														SCREAMING_SNAKE_CASE__       :      Tuple     =						self.conv_in(_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      str     =						jnp.transpose(_UpperCAmelCase,			(0, 2, 3, 1)						)
														SCREAMING_SNAKE_CASE__       :      Tuple     =						self.controlnet_cond_embedding(_UpperCAmelCase						)
														sample += controlnet_cond

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

														# 4. mid
														SCREAMING_SNAKE_CASE__       :      Dict     =						self.mid_block(_UpperCAmelCase,			_UpperCAmelCase,			_UpperCAmelCase,			deterministic=not train						)

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

														SCREAMING_SNAKE_CASE__       :      List[Any]     =						controlnet_down_block_res_samples

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

														# 6. scaling
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						[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=_UpperCAmelCase,			mid_block_res_sample=_UpperCAmelCase						)













import qiskit




def _a (  SCREAMING_SNAKE_CASE__      :  int			,				SCREAMING_SNAKE_CASE__      :  int       )	->    qiskit.result.counts.Counts:

							'''simple docstring'''






							SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						qiskit.Aer.get_backend("aer_simulator"       )

							# Create a Quantum Circuit acting on the q register
							SCREAMING_SNAKE_CASE__       :      List[Any]     =						qiskit.QuantumCircuit(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__       )

							# Apply X (NOT) Gate to Qubits 0 & 1
							circuit.x(0       )
							circuit.x(1       )

							# Map the quantum measurement to the classical bits
							circuit.measure([0, 1]			,				[0, 1]       )

							# Execute the circuit on the qasm simulator
							SCREAMING_SNAKE_CASE__       :      int     =						qiskit.execute(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__			,				shots=10_00       )

							# Return the histogram data of the results of the experiment.
							return job.result().get_counts(SCREAMING_SNAKE_CASE__       )


if __name__ == "__main__":
		_lowerCamelCase :    List[str]  							= single_qubit_measure(2, 2)
		print(f"Total count for various states are: {counts}")






