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  UpperCamelCase				( __lowercase ):

     '''simple docstring'''

     A_										=  42





@flax_register_to_config
class  UpperCamelCase				( nn.Module   ,  __lowercase   ,  __lowercase ):

     '''simple docstring'''

     A_										=  32
     A_										=  4
     A_										=  4
     A_										=  (
         "CrossAttnDownBlock2D",
         "CrossAttnDownBlock2D",
         "CrossAttnDownBlock2D",
         "DownBlock2D",
     )
     A_										=  ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
     A_										=  False
     A_										=  (320, 640, 1_280, 1_280)
     A_										=  2
     A_										=  8
     A_										=  None
     A_										=  1_280
     A_										=  0.0
     A_										=  False
     A_										=  jnp.floataa
     A_										=  True
     A_										=  0
     A_										=  False



     def 		UpperCamelCase_						(					self				,				A_						) ->				FrozenDict:






           """simple docstring"""
           # init input tensors
           _lowerCamelCase			    =				(1, self.in_channels, self.sample_size, self.sample_size)
           _lowerCamelCase			    =				jnp.zeros(A_				,				dtype=jnp.floataa						)
           _lowerCamelCase			    =				jnp.ones((1,)				,				dtype=jnp.intaa						)
           _lowerCamelCase			    =				jnp.zeros((1, 1, self.cross_attention_dim)				,				dtype=jnp.floataa						)

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

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



     def 		UpperCamelCase_						(					self						) ->				Tuple:






           """simple docstring"""
           _lowerCamelCase			    =				self.block_out_channels
           _lowerCamelCase			    =				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.
           _lowerCamelCase			    =				self.num_attention_heads or self.attention_head_dim

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

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

           _lowerCamelCase			    =				self.only_cross_attention
           if isinstance(A_				,				A_						):
                 _lowerCamelCase			    =				(only_cross_attention,) * len(self.down_block_types						)

           if isinstance(A_				,				A_						):
                 _lowerCamelCase			    =				(num_attention_heads,) * len(self.down_block_types						)

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

                 if down_block_type == "CrossAttnDownBlock2D":
                       _lowerCamelCase			    =				FlaxCrossAttnDownBlockaD(
                           in_channels=A_				,				out_channels=A_				,				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:
                       _lowerCamelCase			    =				FlaxDownBlockaD(
                           in_channels=A_				,				out_channels=A_				,				dropout=self.dropout				,				num_layers=self.layers_per_block				,				add_downsample=not is_final_block				,				dtype=self.dtype				,				)

                 down_blocks.append(A_						)
           _lowerCamelCase			    =				down_blocks

           # mid
           _lowerCamelCase			    =				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
           _lowerCamelCase			    =				[]
           _lowerCamelCase			    =				list(reversed(A_						)						)
           _lowerCamelCase			    =				list(reversed(A_						)						)
           _lowerCamelCase			    =				list(reversed(A_						)						)
           _lowerCamelCase			    =				reversed_block_out_channels[0]
           for i, up_block_type in enumerate(self.up_block_types						):
                 _lowerCamelCase			    =				output_channel
                 _lowerCamelCase			    =				reversed_block_out_channels[i]
                 _lowerCamelCase			    =				reversed_block_out_channels[min(i + 1				,				len(A_						) - 1						)]

                 _lowerCamelCase			    =				i == len(A_						) - 1

                 if up_block_type == "CrossAttnUpBlock2D":
                       _lowerCamelCase			    =				FlaxCrossAttnUpBlockaD(
                           in_channels=A_				,				out_channels=A_				,				prev_output_channel=A_				,				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:
                       _lowerCamelCase			    =				FlaxUpBlockaD(
                           in_channels=A_				,				out_channels=A_				,				prev_output_channel=A_				,				num_layers=self.layers_per_block + 1				,				add_upsample=not is_final_block				,				dropout=self.dropout				,				dtype=self.dtype				,				)

                 up_blocks.append(A_						)
                 _lowerCamelCase			    =				output_channel
           _lowerCamelCase			    =				up_blocks

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





     def __call__(					self				,				A_				,				A_				,				A_				,				A_=None				,				A_=None				,				A_ = True				,				A_ = False				,				) ->				Union[FlaxUNetaDConditionOutput, Tuple]:






           """simple docstring"""
           # 1. time
           if not isinstance(A_				,				jnp.ndarray						):
                 _lowerCamelCase			    =				jnp.array([timesteps]				,				dtype=jnp.intaa						)
           elif isinstance(A_				,				jnp.ndarray						) and len(timesteps.shape						) == 0:
                 _lowerCamelCase			    =				timesteps.astype(dtype=jnp.floataa						)
                 _lowerCamelCase			    =				jnp.expand_dims(A_				,				0						)

           _lowerCamelCase			    =				self.time_proj(A_						)
           _lowerCamelCase			    =				self.time_embedding(A_						)

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

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

           if down_block_additional_residuals is not None:
                 _lowerCamelCase			    =				()

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

                 _lowerCamelCase			    =				new_down_block_res_samples

           # 4. mid
           _lowerCamelCase			    =				self.mid_block(A_				,				A_				,				A_				,				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:
                 _lowerCamelCase			    =				down_block_res_samples[-(self.layers_per_block + 1) :]
                 _lowerCamelCase			    =				down_block_res_samples[: -(self.layers_per_block + 1)]
                 if isinstance(A_				,				A_						):
                       _lowerCamelCase			    =				up_block(
                           A_				,				temb=A_				,				encoder_hidden_states=A_				,				res_hidden_states_tuple=A_				,				deterministic=not train				,				)
                 else:
                       _lowerCamelCase			    =				up_block(A_				,				temb=A_				,				res_hidden_states_tuple=A_				,				deterministic=not train						)

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

           if not return_dict:
                 return (sample,)

           return FlaxUNetaDConditionOutput(sample=A_						)
from math import acos, sin
from typing import List, Tuple, Union

import numpy as np
import torch
from PIL import Image

from ...models import AutoencoderKL, UNetaDConditionModel
from ...schedulers import DDIMScheduler, DDPMScheduler
from ...utils import randn_tensor
from ..pipeline_utils import AudioPipelineOutput, BaseOutput, DiffusionPipeline, ImagePipelineOutput
from .mel import Mel



class  UpperCamelCase				( __lowercase ):

     '''simple docstring'''

     A_										=  ['vqvae']



     def __init__(					self				,				A_				,				A_				,				A_				,				A_				,				) ->				Tuple:






           """simple docstring"""
           super().__init__()
           self.register_modules(unet=A_				,				scheduler=A_				,				mel=A_				,				vqvae=A_						)



     def 		UpperCamelCase_						(					self						) ->				int:






           """simple docstring"""
           return 50 if isinstance(self.scheduler				,				A_						) else 10_00



     @torch.no_grad()
     def __call__(					self				,				A_ = 1				,				A_ = None				,				A_ = None				,				A_ = 0				,				A_ = 0				,				A_ = None				,				A_ = None				,				A_ = 0				,				A_ = 0				,				A_ = None				,				A_ = 0				,				A_ = None				,				A_ = None				,				A_=True				,				) ->				Union[
         Union[AudioPipelineOutput, ImagePipelineOutput],
         Tuple[List[Image.Image], Tuple[int, List[np.ndarray]]],
     ]:






           """simple docstring"""
           _lowerCamelCase			    =				steps or self.get_default_steps()
           self.scheduler.set_timesteps(A_						)
           _lowerCamelCase			    =				step_generator or generator
           # For backwards compatibility
           if type(self.unet.config.sample_size						) == int:
                 _lowerCamelCase			    =				(self.unet.config.sample_size, self.unet.config.sample_size)
           if noise is None:
                 _lowerCamelCase			    =				randn_tensor(
                     (
                         batch_size,
                         self.unet.config.in_channels,
                         self.unet.config.sample_size[0],
                         self.unet.config.sample_size[1],
                     )				,				generator=A_				,				device=self.device				,				)
           _lowerCamelCase			    =				noise
           _lowerCamelCase			    =				None

           if audio_file is not None or raw_audio is not None:
                 self.mel.load_audio(A_				,				A_						)
                 _lowerCamelCase			    =				self.mel.audio_slice_to_image(A_						)
                 _lowerCamelCase			    =				np.frombuffer(input_image.tobytes()				,				dtype='''uint8'''						).reshape(
                     (input_image.height, input_image.width)						)
                 _lowerCamelCase			    =				(input_image / 2_55) * 2 - 1
                 _lowerCamelCase			    =				torch.tensor(input_image[np.newaxis, :, :]				,				dtype=torch.float						).to(self.device						)

                 if self.vqvae is not None:
                       _lowerCamelCase			    =				self.vqvae.encode(torch.unsqueeze(A_				,				0						)						).latent_dist.sample(
                           generator=A_						)[0]
                       _lowerCamelCase			    =				self.vqvae.config.scaling_factor * input_images

                 if start_step > 0:
                       _lowerCamelCase			    =				self.scheduler.add_noise(A_				,				A_				,				self.scheduler.timesteps[start_step - 1]						)

                 _lowerCamelCase			    =				(
                     self.unet.config.sample_size[1] * self.mel.get_sample_rate() / self.mel.x_res / self.mel.hop_length
                 )
                 _lowerCamelCase			    =				int(mask_start_secs * pixels_per_second						)
                 _lowerCamelCase			    =				int(mask_end_secs * pixels_per_second						)
                 _lowerCamelCase			    =				self.scheduler.add_noise(A_				,				A_				,				torch.tensor(self.scheduler.timesteps[start_step:]						)						)

           for step, t in enumerate(self.progress_bar(self.scheduler.timesteps[start_step:]						)						):
                 if isinstance(self.unet				,				A_						):
                       _lowerCamelCase			    =				self.unet(A_				,				A_				,				A_						)['''sample''']
                 else:
                       _lowerCamelCase			    =				self.unet(A_				,				A_						)['''sample''']

                 if isinstance(self.scheduler				,				A_						):
                       _lowerCamelCase			    =				self.scheduler.step(
                           model_output=A_				,				timestep=A_				,				sample=A_				,				eta=A_				,				generator=A_				,				)['''prev_sample''']
                 else:
                       _lowerCamelCase			    =				self.scheduler.step(
                           model_output=A_				,				timestep=A_				,				sample=A_				,				generator=A_				,				)['''prev_sample''']

                 if mask is not None:
                       if mask_start > 0:
                             _lowerCamelCase			    =				mask[:, step, :, :mask_start]
                       if mask_end > 0:
                             _lowerCamelCase			    =				mask[:, step, :, -mask_end:]

           if self.vqvae is not None:
                 # 0.18215 was scaling factor used in training to ensure unit variance
                 _lowerCamelCase			    =				1 / self.vqvae.config.scaling_factor * images
                 _lowerCamelCase			    =				self.vqvae.decode(A_						)['''sample''']

           _lowerCamelCase			    =				(images / 2 + 0.5).clamp(0				,				1						)
           _lowerCamelCase			    =				images.cpu().permute(0				,				2				,				3				,				1						).numpy()
           _lowerCamelCase			    =				(images * 2_55).round().astype('''uint8'''						)
           _lowerCamelCase			    =				list(
               (Image.fromarray(_[:, :, 0]						) for _ in images)
               if images.shape[3] == 1
               else (Image.fromarray(A_				,				mode='''RGB'''						).convert('''L'''						) for _ in images)						)

           _lowerCamelCase			    =				[self.mel.image_to_audio(A_						) for _ in images]
           if not return_dict:
                 return images, (self.mel.get_sample_rate(), audios)

           return BaseOutput(**AudioPipelineOutput(np.array(A_						)[:, np.newaxis, :]						)				,				**ImagePipelineOutput(A_						)						)



     @torch.no_grad()
     def 		UpperCamelCase_						(					self				,				A_				,				A_ = 50						) ->				np.ndarray:






           """simple docstring"""
           assert isinstance(self.scheduler				,				A_						)
           self.scheduler.set_timesteps(A_						)
           _lowerCamelCase			    =				np.array(
               [np.frombuffer(image.tobytes()				,				dtype='''uint8'''						).reshape((1, image.height, image.width)						) for image in images]						)
           _lowerCamelCase			    =				(sample / 2_55) * 2 - 1
           _lowerCamelCase			    =				torch.Tensor(A_						).to(self.device						)

           for t in self.progress_bar(torch.flip(self.scheduler.timesteps				,				(0,)						)						):
                 _lowerCamelCase			    =				t - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
                 _lowerCamelCase			    =				self.scheduler.alphas_cumprod[t]
                 _lowerCamelCase			    =				(
                     self.scheduler.alphas_cumprod[prev_timestep]
                     if prev_timestep >= 0
                     else self.scheduler.final_alpha_cumprod
                 )
                 _lowerCamelCase			    =				1 - alpha_prod_t
                 _lowerCamelCase			    =				self.unet(A_				,				A_						)['''sample''']
                 _lowerCamelCase			    =				(1 - alpha_prod_t_prev) ** 0.5 * model_output
                 _lowerCamelCase			    =				(sample - pred_sample_direction) * alpha_prod_t_prev ** (-0.5)
                 _lowerCamelCase			    =				sample * alpha_prod_t ** 0.5 + beta_prod_t ** 0.5 * model_output

           return sample





     @staticmethod
     def 		UpperCamelCase_						(					A_				,				A_				,				A_						) ->				torch.Tensor:






           """simple docstring"""
           _lowerCamelCase			    =				acos(torch.dot(torch.flatten(A_						)				,				torch.flatten(A_						)						) / torch.norm(A_						) / torch.norm(A_						)						)
           return sin((1 - alpha) * theta						) * xa / sin(A_						) + sin(alpha * theta						) * xa / sin(A_						)