






'''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,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)






@flax.struct.dataclass
class 				a_						(			lowerCamelCase				):
  lowercase     		=      42






@flax_register_to_config
class 				a_						(			nn.Module	,       lowerCamelCase	,       lowerCamelCase				):
  lowercase     		=      32
  lowercase     		=      4
  lowercase     		=      4
  lowercase     		=      (
      """CrossAttnDownBlock2D""",
      """CrossAttnDownBlock2D""",
      """CrossAttnDownBlock2D""",
      """DownBlock2D""",
  )
  lowercase     		=      ("""UpBlock2D""", """CrossAttnUpBlock2D""", """CrossAttnUpBlock2D""", """CrossAttnUpBlock2D""")
  lowercase     		=      False
  lowercase     		=      (3_20, 6_40, 12_80, 12_80)
  lowercase     		=      2
  lowercase     		=      8
  lowercase     		=      None
  lowercase     		=      12_80
  lowercase     		=      0.0
  lowercase     		=      False
  lowercase     		=      jnp.floataa
  lowercase     		=      True
  lowercase     		=      0
  lowercase     		=      False

  def   A__						(      self			,		_SCREAMING_SNAKE_CASE      )  ->      str:



        """simple docstring"""






        UpperCamelCase				  =					(1, self.in_channels, self.sample_size, self.sample_size)
        UpperCamelCase				  =					jnp.zeros(_SCREAMING_SNAKE_CASE			,		dtype=jnp.floataa      )
        UpperCamelCase				  =					jnp.ones((1,)			,		dtype=jnp.intaa      )
        UpperCamelCase				  =					jnp.zeros((1, 1, self.cross_attention_dim)			,		dtype=jnp.floataa      )

        UpperCamelCase				  =					jax.random.split(_SCREAMING_SNAKE_CASE      )
        UpperCamelCase				  =					{'params': params_rng, 'dropout': dropout_rng}

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

  def   A__						(      self      )  ->      Tuple:



        """simple docstring"""






        UpperCamelCase				  =					self.block_out_channels
        UpperCamelCase				  =					block_out_channels[0] * 4

        if self.num_attention_heads is not None:
              raise ValueError(
                  """At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19."""      )

        # 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				  =					self.num_attention_heads or self.attention_head_dim

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

        # time
        UpperCamelCase				  =					FlaxTimesteps(
            block_out_channels[0]			,		flip_sin_to_cos=self.flip_sin_to_cos			,		freq_shift=self.config.freq_shift      )
        UpperCamelCase				  =					FlaxTimestepEmbedding(_SCREAMING_SNAKE_CASE			,		dtype=self.dtype      )

        UpperCamelCase				  =					self.only_cross_attention
        if isinstance(_SCREAMING_SNAKE_CASE			,		_SCREAMING_SNAKE_CASE      ):
              UpperCamelCase				  =					(only_cross_attention,) * len(self.down_block_types      )

        if isinstance(_SCREAMING_SNAKE_CASE			,		_SCREAMING_SNAKE_CASE      ):
              UpperCamelCase				  =					(num_attention_heads,) * len(self.down_block_types      )

        # down
        UpperCamelCase				  =					[]
        UpperCamelCase				  =					block_out_channels[0]
        for i, down_block_type in enumerate(self.down_block_types      ):
              UpperCamelCase				  =					output_channel
              UpperCamelCase				  =					block_out_channels[i]
              UpperCamelCase				  =					i == len(_SCREAMING_SNAKE_CASE      ) - 1

              if down_block_type == "CrossAttnDownBlock2D":
                    UpperCamelCase				  =					FlaxCrossAttnDownBlockaD(
                        in_channels=_SCREAMING_SNAKE_CASE			,		out_channels=_SCREAMING_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]			,		use_memory_efficient_attention=self.use_memory_efficient_attention			,		dtype=self.dtype			,		)
              else:
                    UpperCamelCase				  =					FlaxDownBlockaD(
                        in_channels=_SCREAMING_SNAKE_CASE			,		out_channels=_SCREAMING_SNAKE_CASE			,		dropout=self.dropout			,		num_layers=self.layers_per_block			,		add_downsample=not is_final_block			,		dtype=self.dtype			,		)

              down_blocks.append(_SCREAMING_SNAKE_CASE      )
        UpperCamelCase				  =					down_blocks

        # mid
        UpperCamelCase				  =					FlaxUNetMidBlockaDCrossAttn(
            in_channels=block_out_channels[-1]			,		dropout=self.dropout			,		num_attention_heads=num_attention_heads[-1]			,		use_linear_projection=self.use_linear_projection			,		use_memory_efficient_attention=self.use_memory_efficient_attention			,		dtype=self.dtype			,		)

        # up
        UpperCamelCase				  =					[]
        UpperCamelCase				  =					list(reversed(_SCREAMING_SNAKE_CASE      )      )
        UpperCamelCase				  =					list(reversed(_SCREAMING_SNAKE_CASE      )      )
        UpperCamelCase				  =					list(reversed(_SCREAMING_SNAKE_CASE      )      )
        UpperCamelCase				  =					reversed_block_out_channels[0]
        for i, up_block_type in enumerate(self.up_block_types      ):
              UpperCamelCase				  =					output_channel
              UpperCamelCase				  =					reversed_block_out_channels[i]
              UpperCamelCase				  =					reversed_block_out_channels[min(i + 1			,		len(_SCREAMING_SNAKE_CASE      ) - 1      )]

              UpperCamelCase				  =					i == len(_SCREAMING_SNAKE_CASE      ) - 1

              if up_block_type == "CrossAttnUpBlock2D":
                    UpperCamelCase				  =					FlaxCrossAttnUpBlockaD(
                        in_channels=_SCREAMING_SNAKE_CASE			,		out_channels=_SCREAMING_SNAKE_CASE			,		prev_output_channel=_SCREAMING_SNAKE_CASE			,		num_layers=self.layers_per_block + 1			,		num_attention_heads=reversed_num_attention_heads[i]			,		add_upsample=not is_final_block			,		dropout=self.dropout			,		use_linear_projection=self.use_linear_projection			,		only_cross_attention=only_cross_attention[i]			,		use_memory_efficient_attention=self.use_memory_efficient_attention			,		dtype=self.dtype			,		)
              else:
                    UpperCamelCase				  =					FlaxUpBlockaD(
                        in_channels=_SCREAMING_SNAKE_CASE			,		out_channels=_SCREAMING_SNAKE_CASE			,		prev_output_channel=_SCREAMING_SNAKE_CASE			,		num_layers=self.layers_per_block + 1			,		add_upsample=not is_final_block			,		dropout=self.dropout			,		dtype=self.dtype			,		)

              up_blocks.append(_SCREAMING_SNAKE_CASE      )
              UpperCamelCase				  =					output_channel
        UpperCamelCase				  =					up_blocks

        # out
        UpperCamelCase				  =					nn.GroupNorm(num_groups=32			,		epsilon=1e-5      )
        UpperCamelCase				  =					nn.Conv(
            self.out_channels			,		kernel_size=(3, 3)			,		strides=(1, 1)			,		padding=((1, 1), (1, 1))			,		dtype=self.dtype			,		)


  def __call__(      self			,		_SCREAMING_SNAKE_CASE			,		_SCREAMING_SNAKE_CASE			,		_SCREAMING_SNAKE_CASE			,		_SCREAMING_SNAKE_CASE=None			,		_SCREAMING_SNAKE_CASE=None			,		_SCREAMING_SNAKE_CASE = True			,		_SCREAMING_SNAKE_CASE = False			,		)  ->      List[Any]:



        """simple docstring"""






        if not isinstance(_SCREAMING_SNAKE_CASE			,		jnp.ndarray      ):
              UpperCamelCase				  =					jnp.array([timesteps]			,		dtype=jnp.intaa      )
        elif isinstance(_SCREAMING_SNAKE_CASE			,		jnp.ndarray      ) and len(timesteps.shape      ) == 0:
              UpperCamelCase				  =					timesteps.astype(dtype=jnp.floataa      )
              UpperCamelCase				  =					jnp.expand_dims(_SCREAMING_SNAKE_CASE			,		0      )

        UpperCamelCase				  =					self.time_proj(_SCREAMING_SNAKE_CASE      )
        UpperCamelCase				  =					self.time_embedding(_SCREAMING_SNAKE_CASE      )

        # 2. pre-process
        UpperCamelCase				  =					jnp.transpose(_SCREAMING_SNAKE_CASE			,		(0, 2, 3, 1)      )
        UpperCamelCase				  =					self.conv_in(_SCREAMING_SNAKE_CASE      )

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

        if down_block_additional_residuals is not None:
              UpperCamelCase				  =					()

              for down_block_res_sample, down_block_additional_residual in zip(
                  _SCREAMING_SNAKE_CASE			,		_SCREAMING_SNAKE_CASE      ):
                    down_block_res_sample += down_block_additional_residual
                    new_down_block_res_samples += (down_block_res_sample,)

              UpperCamelCase				  =					new_down_block_res_samples

        # 4. mid
        UpperCamelCase				  =					self.mid_block(_SCREAMING_SNAKE_CASE			,		_SCREAMING_SNAKE_CASE			,		_SCREAMING_SNAKE_CASE			,		deterministic=not train      )

        if mid_block_additional_residual is not None:
              sample += mid_block_additional_residual

        # 5. up
        for up_block in self.up_blocks:
              UpperCamelCase				  =					down_block_res_samples[-(self.layers_per_block + 1) :]
              UpperCamelCase				  =					down_block_res_samples[: -(self.layers_per_block + 1)]
              if isinstance(_SCREAMING_SNAKE_CASE			,		_SCREAMING_SNAKE_CASE      ):
                    UpperCamelCase				  =					up_block(
                        _SCREAMING_SNAKE_CASE			,		temb=_SCREAMING_SNAKE_CASE			,		encoder_hidden_states=_SCREAMING_SNAKE_CASE			,		res_hidden_states_tuple=_SCREAMING_SNAKE_CASE			,		deterministic=not train			,		)
              else:
                    UpperCamelCase				  =					up_block(_SCREAMING_SNAKE_CASE			,		temb=_SCREAMING_SNAKE_CASE			,		res_hidden_states_tuple=_SCREAMING_SNAKE_CASE			,		deterministic=not train      )

        # 6. post-process
        UpperCamelCase				  =					self.conv_norm_out(_SCREAMING_SNAKE_CASE      )
        UpperCamelCase				  =					nn.silu(_SCREAMING_SNAKE_CASE      )
        UpperCamelCase				  =					self.conv_out(_SCREAMING_SNAKE_CASE      )
        UpperCamelCase				  =					jnp.transpose(_SCREAMING_SNAKE_CASE			,		(0, 3, 1, 2)      )

        if not return_dict:
              return (sample,)

        return FlaxUNetaDConditionOutput(sample=_SCREAMING_SNAKE_CASE      )






"""simple docstring"""






import datetime
import platform
import subprocess
from typing import Optional, Tuple, Union

import numpy as np







def      _snake_case			(				lowercase__							,    lowercase__		):
				_lowerCamelCase     :			Tuple    =    f'''{sampling_rate}'''
				_lowerCamelCase     :			str    =    '1'
				_lowerCamelCase     :			str    =    'f32le'
				_lowerCamelCase     :			Union[str, Any]    =    [
				    'ffmpeg',
				    '-i',
				    'pipe:0',
				    '-ac',
				    ac,
				    '-ar',
				    ar,
				    '-f',
				    format_for_conversion,
				    '-hide_banner',
				    '-loglevel',
				    'quiet',
				    'pipe:1',
				]

				try:
								with subprocess.Popen(lowercase__							,    stdin=subprocess.PIPE							,    stdout=subprocess.PIPE		) as ffmpeg_process:
												_lowerCamelCase     :			str    =    ffmpeg_process.communicate(lowercase__		)
				except FileNotFoundError as error:
								raise ValueError('ffmpeg was not found but is required to load audio files from filename'		) from error
				_lowerCamelCase     :			List[Any]    =    output_stream[0]
				_lowerCamelCase     :			Tuple    =    np.frombuffer(lowercase__							,    np.floataa		)
				if audio.shape[0] == 0:
								raise ValueError('Malformed soundfile'		)
				return audio







def      _snake_case			(				lowercase__							,    lowercase__							,    lowercase__ = "f32le"							,    ):
				_lowerCamelCase     :			Optional[Any]    =    f'''{sampling_rate}'''
				_lowerCamelCase     :			List[str]    =    '1'
				if format_for_conversion == "s16le":
								_lowerCamelCase     :			List[str]    =    2
				elif format_for_conversion == "f32le":
								_lowerCamelCase     :			List[Any]    =    4
				else:
								raise ValueError(f'''Unhandled format `{format_for_conversion}`. Please use `s16le` or `f32le`'''		)

				_lowerCamelCase     :			Dict    =    platform.system()
				if system == "Linux":
								_lowerCamelCase     :			Optional[int]    =    'alsa'
								_lowerCamelCase     :			Optional[Any]    =    'default'
				elif system == "Darwin":
								_lowerCamelCase     :			Optional[int]    =    'avfoundation'
								_lowerCamelCase     :			Any    =    ':0'
				elif system == "Windows":
								_lowerCamelCase     :			Tuple    =    'dshow'
								_lowerCamelCase     :			Tuple    =    'default'

				_lowerCamelCase     :			Optional[int]    =    [
				    'ffmpeg',
				    '-f',
				    format_,
				    '-i',
				    input_,
				    '-ac',
				    ac,
				    '-ar',
				    ar,
				    '-f',
				    format_for_conversion,
				    '-fflags',
				    'nobuffer',
				    '-hide_banner',
				    '-loglevel',
				    'quiet',
				    'pipe:1',
				]
				_lowerCamelCase     :			Tuple    =    int(round(sampling_rate * chunk_length_s		)		) * size_of_sample
				_lowerCamelCase     :			List[Any]    =    _ffmpeg_stream(lowercase__							,    lowercase__		)
				for item in iterator:
								yield item







def      _snake_case			(				lowercase__							,    lowercase__							,    lowercase__ = None							,    lowercase__ = None							,    lowercase__ = "f32le"							,    ):
				if stream_chunk_s is not None:
								_lowerCamelCase     :			int    =    stream_chunk_s
				else:
								_lowerCamelCase     :			Optional[Any]    =    chunk_length_s

				_lowerCamelCase     :			Optional[Any]    =    ffmpeg_microphone(lowercase__							,    lowercase__							,    format_for_conversion=lowercase__		)
				if format_for_conversion == "s16le":
								_lowerCamelCase     :			List[str]    =    np.intaa
								_lowerCamelCase     :			str    =    2
				elif format_for_conversion == "f32le":
								_lowerCamelCase     :			Any    =    np.floataa
								_lowerCamelCase     :			List[Any]    =    4
				else:
								raise ValueError(f'''Unhandled format `{format_for_conversion}`. Please use `s16le` or `f32le`'''		)

				if stride_length_s is None:
								_lowerCamelCase     :			Union[str, Any]    =    chunk_length_s / 6
				_lowerCamelCase     :			Optional[int]    =    int(round(sampling_rate * chunk_length_s		)		) * size_of_sample
				if isinstance(lowercase__							,    (int, float)		):
								_lowerCamelCase     :			Any    =    [stride_length_s, stride_length_s]

				_lowerCamelCase     :			Tuple    =    int(round(sampling_rate * stride_length_s[0]		)		) * size_of_sample
				_lowerCamelCase     :			Optional[Any]    =    int(round(sampling_rate * stride_length_s[1]		)		) * size_of_sample
				_lowerCamelCase     :			List[Any]    =    datetime.datetime.now()
				_lowerCamelCase     :			Optional[int]    =    datetime.timedelta(seconds=lowercase__		)
				for item in chunk_bytes_iter(lowercase__							,    lowercase__							,    stride=(stride_left, stride_right)							,    stream=lowercase__		):
								# Put everything back in numpy scale
								_lowerCamelCase     :			List[Any]    =    np.frombuffer(item['raw']							,    dtype=lowercase__		)
								_lowerCamelCase     :			int    =    (
								    item['stride'][0] // size_of_sample,
								    item['stride'][1] // size_of_sample,
								)
								_lowerCamelCase     :			Optional[int]    =    sampling_rate
								audio_time += delta
								if datetime.datetime.now() > audio_time + 10 * delta:
												# We're late !! SKIP
												continue
								yield item







def      _snake_case			(				lowercase__							,    lowercase__							,    lowercase__							,    lowercase__ = False		):
				_lowerCamelCase     :			int    =    B''
				_lowerCamelCase, _lowerCamelCase     :			Dict    =    stride
				if stride_left + stride_right >= chunk_len:
								raise ValueError(
								    f'''Stride needs to be strictly smaller than chunk_len: ({stride_left}, {stride_right}) vs {chunk_len}'''		)
				_lowerCamelCase     :			str    =    0
				for raw in iterator:
								acc += raw
								if stream and len(lowercase__		) < chunk_len:
												_lowerCamelCase     :			Optional[int]    =    (_stride_left, 0)
												yield {"raw": acc[:chunk_len], "stride": stride, "partial": True}
								else:
												while len(lowercase__		) >= chunk_len:
																# We are flushing the accumulator
																_lowerCamelCase     :			str    =    (_stride_left, stride_right)
																_lowerCamelCase     :			str    =    {'raw': acc[:chunk_len], 'stride': stride}
																if stream:
																				_lowerCamelCase     :			List[Any]    =    False
																yield item
																_lowerCamelCase     :			Optional[Any]    =    stride_left
																_lowerCamelCase     :			str    =    acc[chunk_len - stride_left - stride_right :]
    # Last chunk
				if len(lowercase__		) > stride_left:
								_lowerCamelCase     :			Optional[Any]    =    {'raw': acc, 'stride': (_stride_left, 0)}
								if stream:
												_lowerCamelCase     :			Tuple    =    False
								yield item







def      _snake_case			(				lowercase__							,    lowercase__		):
				_lowerCamelCase     :			int    =    2**24  # 16Mo
				try:
								with subprocess.Popen(lowercase__							,    stdout=subprocess.PIPE							,    bufsize=lowercase__		) as ffmpeg_process:
												while True:
																_lowerCamelCase     :			Optional[Any]    =    ffmpeg_process.stdout.read(lowercase__		)
																if raw == b"":
																				break
																yield raw
				except FileNotFoundError as error:
								raise ValueError('ffmpeg was not found but is required to stream audio files from filename'		) from error