




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__								=(1_6, 3_2, 9_6, 2_5_6)
			UpperCamelCase__								=jnp.floataa

			def 							UpperCAmelCase__       ( self			:   Any  )						->   Union[str, Any]:
								__magic_name__							:	str    							=							nn.Conv(
								    self.block_out_channels[0]       ,						kernel_size=(3, 3)       ,						padding=((1, 1), (1, 1))       ,						dtype=self.dtype       ,						)

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

								__magic_name__							:	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			:   Dict       ,						lowerCamelCase_			:   int  )						->   int:
								__magic_name__							:	Any    							=							self.conv_in(lowerCamelCase_  )
								__magic_name__							:	Any    							=							nn.silu(lowerCamelCase_  )

								for block in self.blocks:
													__magic_name__							:	int    							=							block(lowerCamelCase_  )
													__magic_name__							:	Optional[Any]    							=							nn.silu(lowerCamelCase_  )

								__magic_name__							:	Any    							=							self.conv_out(lowerCamelCase_  )

								return embedding






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






			'''simple docstring'''





			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 							UpperCAmelCase__       ( self			:   Dict       ,						lowerCamelCase_			:   jax.random.KeyArray  )						->   FrozenDict:
								# init input tensors
								__magic_name__							:	Optional[int]    							=							(1, self.in_channels, self.sample_size, self.sample_size)
								__magic_name__							:	List[str]    							=							jnp.zeros(lowerCamelCase_       ,						dtype=jnp.floataa  )
								__magic_name__							:	List[str]    							=							jnp.ones((1,)       ,						dtype=jnp.intaa  )
								__magic_name__							:	Any    							=							jnp.zeros((1, 1, self.cross_attention_dim)       ,						dtype=jnp.floataa  )
								__magic_name__							:	List[str]    							=							(1, 3, self.sample_size * 8, self.sample_size * 8)
								__magic_name__							:	Optional[int]    							=							jnp.zeros(lowerCamelCase_       ,						dtype=jnp.floataa  )

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

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

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

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

								# time
								__magic_name__							:	Optional[int]    							=							FlaxTimesteps(
								    block_out_channels[0]       ,						flip_sin_to_cos=self.flip_sin_to_cos       ,						freq_shift=self.config.freq_shift  )
								__magic_name__							:	Dict    							=							FlaxTimestepEmbedding(lowerCamelCase_       ,						dtype=self.dtype  )

								__magic_name__							:	str    							=							FlaxControlNetConditioningEmbedding(
								    conditioning_embedding_channels=block_out_channels[0]       ,						block_out_channels=self.conditioning_embedding_out_channels       ,						)

								__magic_name__							:	Optional[Any]    							=							self.only_cross_attention
								if isinstance(lowerCamelCase_       ,						lowerCamelCase_  ):
													__magic_name__							:	List[str]    							=							(only_cross_attention,) * len(self.down_block_types  )

								if isinstance(lowerCamelCase_       ,						lowerCamelCase_  ):
													__magic_name__							:	Dict    							=							(num_attention_heads,) * len(self.down_block_types  )

								# down
								__magic_name__							:	int    							=							[]
								__magic_name__							:	Tuple    							=							[]

								__magic_name__							:	Optional[int]    							=							block_out_channels[0]

								__magic_name__							:	Optional[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  ):
													__magic_name__							:	Tuple    							=							output_channel
													__magic_name__							:	Optional[int]    							=							block_out_channels[i]
													__magic_name__							:	Optional[Any]    							=							i == len(lowerCamelCase_  ) - 1

													if down_block_type == "CrossAttnDownBlock2D":
																		__magic_name__							:	Dict    							=							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:
																		__magic_name__							:	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  ):
																		__magic_name__							:	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:
																		__magic_name__							:	Optional[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_  )

								__magic_name__							:	Optional[int]    							=							down_blocks
								__magic_name__							:	Optional[int]    							=							controlnet_down_blocks

								# mid
								__magic_name__							:	Any    							=							block_out_channels[-1]
								__magic_name__							:	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       ,						)

								__magic_name__							:	Dict    							=							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			:   Dict       ,						lowerCamelCase_			:   str       ,						lowerCamelCase_			:   List[Any]       ,						lowerCamelCase_			:   str       ,						lowerCamelCase_			:   str       ,						lowerCamelCase_			:   float = 1.0       ,						lowerCamelCase_			:   bool = True       ,						lowerCamelCase_			:   bool = False       ,						)						->   Union[FlaxControlNetOutput, Tuple]:
								__magic_name__							:	List[str]    							=							self.controlnet_conditioning_channel_order
								if channel_order == "bgr":
													__magic_name__							:	int    							=							jnp.flip(lowerCamelCase_       ,						axis=1  )

								# 1. time
								if not isinstance(lowerCamelCase_       ,						jnp.ndarray  ):
													__magic_name__							:	Tuple    							=							jnp.array([timesteps]       ,						dtype=jnp.intaa  )
								elif isinstance(lowerCamelCase_       ,						jnp.ndarray  ) and len(timesteps.shape  ) == 0:
													__magic_name__							:	Any    							=							timesteps.astype(dtype=jnp.floataa  )
													__magic_name__							:	Optional[Any]    							=							jnp.expand_dims(lowerCamelCase_       ,						0  )

								__magic_name__							:	Tuple    							=							self.time_proj(lowerCamelCase_  )
								__magic_name__							:	Union[str, Any]    							=							self.time_embedding(lowerCamelCase_  )

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

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

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

								# 4. mid
								__magic_name__							:	Optional[Any]    							=							self.mid_block(lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_       ,						deterministic=not train  )

								# 5. contronet blocks
								__magic_name__							:	Tuple    							=							()
								for down_block_res_sample, controlnet_block in zip(lowerCamelCase_       ,						self.controlnet_down_blocks  ):
													__magic_name__							:	Optional[Any]    							=							controlnet_block(lowerCamelCase_  )
													controlnet_down_block_res_samples += (down_block_res_sample,)

								__magic_name__							:	int    							=							controlnet_down_block_res_samples

								__magic_name__							:	str    							=							self.controlnet_mid_block(lowerCamelCase_  )

								# 6. scaling
								__magic_name__							:	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=lowerCamelCase_       ,						mid_block_res_sample=lowerCamelCase_  )









import inspect
import os
import unittest
from pathlib import Path

import torch

import accelerate
from accelerate.test_utils import execute_subprocess_async
from accelerate.test_utils.testing import run_command


class    lowerCamelCase    (		unittest.TestCase						):






			'''simple docstring'''





			UpperCamelCase__								=inspect.getfile(accelerate.test_utils						)
			UpperCamelCase__								=os.path.sep.join(mod_file.split(os.path.sep						)[:-1] + ['''scripts''', '''test_cli.py''']						)

			UpperCamelCase__								=['''accelerate''', '''launch''']
			UpperCamelCase__								=Path.home() / '''.cache/huggingface/accelerate'''
			UpperCamelCase__								='''default_config.yaml'''
			UpperCamelCase__								=config_folder / config_file
			UpperCamelCase__								=config_folder / '''_default_config.yaml'''

			UpperCamelCase__								=Path('''tests/test_configs'''						)

			@classmethod
			def 							UpperCAmelCase__       ( cls			:   Optional[int]  )						->   int:
								if cls.config_path.is_file():
													cls.config_path.rename(cls.changed_path  )

			@classmethod
			def 							UpperCAmelCase__       ( cls			:   Dict  )						->   List[str]:
								if cls.changed_path.is_file():
													cls.changed_path.rename(cls.config_path  )

			def 							UpperCAmelCase__       ( self			:   str  )						->   Optional[Any]:
								__magic_name__							:	Dict    							=							self.base_cmd
								if torch.cuda.is_available() and (torch.cuda.device_count() > 1):
													cmd += ["--multi_gpu"]
								execute_subprocess_async(cmd + [self.test_file_path]       ,						env=os.environ.copy()  )

			def 							UpperCAmelCase__       ( self			:   int  )						->   Optional[int]:
								for config in sorted(self.test_config_path.glob('''**/*.yaml'''  )  ):
													with self.subTest(config_file=lowerCamelCase_  ):
																		execute_subprocess_async(
																		    self.base_cmd + ['''--config_file''', str(lowerCamelCase_  ), self.test_file_path]       ,						env=os.environ.copy()  )


			def 							UpperCAmelCase__       ( self			:   int  )						->   Union[str, Any]:
								execute_subprocess_async(['''accelerate''', '''test''']       ,						env=os.environ.copy()  )






class    lowerCamelCase    (		unittest.TestCase						):






			'''simple docstring'''





			UpperCamelCase__								='''test-tpu'''
			UpperCamelCase__								='''us-central1-a'''
			UpperCamelCase__								='''ls'''
			UpperCamelCase__								=['''accelerate''', '''tpu-config''']
			UpperCamelCase__								='''cd /usr/share'''
			UpperCamelCase__								='''tests/test_samples/test_command_file.sh'''
			UpperCamelCase__								='''Running gcloud compute tpus tpu-vm ssh'''

			def 							UpperCAmelCase__       ( self			:   Optional[Any]  )						->   Optional[Any]:
								__magic_name__							:	str    							=							run_command(
								    self.cmd
								    + ['''--command''', self.command, '''--tpu_zone''', self.tpu_zone, '''--tpu_name''', self.tpu_name, '''--debug''']       ,						return_stdout=lowerCamelCase_       ,						)
								self.assertIn(
								    F'''{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all'''       ,						lowerCamelCase_       ,						)

			def 							UpperCAmelCase__       ( self			:   Union[str, Any]  )						->   int:
								__magic_name__							:	int    							=							run_command(
								    self.cmd
								    + [
								        '''--config_file''',
								        '''tests/test_configs/0_12_0.yaml''',
								        '''--command''',
								        self.command,
								        '''--tpu_zone''',
								        self.tpu_zone,
								        '''--tpu_name''',
								        self.tpu_name,
								        '''--debug''',
								    ]       ,						return_stdout=lowerCamelCase_       ,						)
								self.assertIn(
								    F'''{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all'''       ,						lowerCamelCase_       ,						)

			def 							UpperCAmelCase__       ( self			:   List[Any]  )						->   Tuple:
								__magic_name__							:	Any    							=							run_command(
								    self.cmd + ['''--config_file''', '''tests/test_configs/latest.yaml''', '''--debug''']       ,						return_stdout=lowerCamelCase_  )
								self.assertIn(
								    F'''{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo "hello world"; echo "this is a second command" --worker all'''       ,						lowerCamelCase_       ,						)

			def 							UpperCAmelCase__       ( self			:   Optional[int]  )						->   Tuple:
								__magic_name__							:	Tuple    							=							run_command(
								    self.cmd + ['''--config_file''', '''tests/test_configs/latest.yaml''', '''--command''', self.command, '''--debug''']       ,						return_stdout=lowerCamelCase_       ,						)
								self.assertIn(
								    F'''{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all'''       ,						lowerCamelCase_       ,						)

			def 							UpperCAmelCase__       ( self			:   List[Any]  )						->   Dict:
								__magic_name__							:	str    							=							run_command(
								    self.cmd
								    + [
								        '''--config_file''',
								        '''tests/test_configs/latest.yaml''',
								        '''--command''',
								        self.command,
								        '''--command''',
								        '''echo "Hello World"''',
								        '''--debug''',
								    ]       ,						return_stdout=lowerCamelCase_       ,						)
								self.assertIn(
								    F'''{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls; echo "Hello World" --worker all'''       ,						lowerCamelCase_       ,						)

			def 							UpperCAmelCase__       ( self			:   Optional[Any]  )						->   Optional[int]:
								__magic_name__							:	Union[str, Any]    							=							run_command(
								    self.cmd
								    + ['''--config_file''', '''tests/test_configs/latest.yaml''', '''--command_file''', self.command_file, '''--debug''']       ,						return_stdout=lowerCamelCase_       ,						)
								self.assertIn(
								    F'''{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo "hello world"; echo "this is a second command" --worker all'''       ,						lowerCamelCase_       ,						)

			def 							UpperCAmelCase__       ( self			:   int  )						->   List[Any]:
								__magic_name__							:	List[Any]    							=							run_command(
								    self.cmd
								    + [
								        '''--config_file''',
								        '''tests/test_configs/0_12_0.yaml''',
								        '''--command_file''',
								        self.command_file,
								        '''--tpu_zone''',
								        self.tpu_zone,
								        '''--tpu_name''',
								        self.tpu_name,
								        '''--debug''',
								    ]       ,						return_stdout=lowerCamelCase_       ,						)
								self.assertIn(
								    F'''{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo "hello world"; echo "this is a second command" --worker all'''       ,						lowerCamelCase_       ,						)

			def 							UpperCAmelCase__       ( self			:   str  )						->   Dict:
								__magic_name__							:	Tuple    							=							run_command(
								    self.cmd + ['''--config_file''', '''tests/test_configs/latest.yaml''', '''--install_accelerate''', '''--debug''']       ,						return_stdout=lowerCamelCase_       ,						)
								self.assertIn(
								    F'''{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; pip install accelerate -U; echo "hello world"; echo "this is a second command" --worker all'''       ,						lowerCamelCase_       ,						)


			def 							UpperCAmelCase__       ( self			:   Any  )						->   Tuple:
								__magic_name__							:	Tuple    							=							run_command(
								    self.cmd
								    + [
								        '''--config_file''',
								        '''tests/test_configs/latest.yaml''',
								        '''--install_accelerate''',
								        '''--accelerate_version''',
								        '''12.0.0''',
								        '''--debug''',
								    ]       ,						return_stdout=lowerCamelCase_       ,						)
								self.assertIn(
								    F'''{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; pip install accelerate==12.0.0; echo "hello world"; echo "this is a second command" --worker all'''       ,						lowerCamelCase_       ,						)



