






"""simple docstring"""







from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin
from .unet_ad_blocks import get_down_block, get_mid_block, get_out_block, get_up_block







@dataclass
class        _SCREAMING_SNAKE_CASE						(  __UpperCAmelCase							):





       """simple docstring"""





       _SCREAMING_SNAKE_CASE    =42







class        _SCREAMING_SNAKE_CASE						(  __UpperCAmelCase       , __UpperCAmelCase							):





       """simple docstring"""





       @register_to_config
       def __init__(   self:     Optional[int]				,		__A:     int = 65536				,		__A:     Optional[int] = None				,		__A:     int = 2				,		__A:     int = 2				,		__A:     int = 0				,		__A:     str = "fourier"				,		__A:     bool = True				,		__A:     bool = False				,		__A:     float = 0.0				,		__A:     Tuple[str] = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D")				,		__A:     Tuple[str] = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip")				,		__A:     Tuple[str] = "UNetMidBlock1D"				,		__A:     str = None				,		__A:     Tuple[int] = (32, 32, 64)				,		__A:     str = None				,		__A:     int = 8				,		__A:     int = 1				,		__A:     bool = False				,		):

            '''simple docstring'''

            super().__init__()
            a__          =  sample_size

            # time
            if time_embedding_type == "fourier":
                 a__          =  GaussianFourierProjection(
                     embedding_size=8				,		set_W_to_weight=lowerCAmelCase_				,		log=lowerCAmelCase_				,		flip_sin_to_cos=lowerCAmelCase_    )
                 a__          =  2 * block_out_channels[0]
            elif time_embedding_type == "positional":
                 a__          =  Timesteps(
                     block_out_channels[0]				,		flip_sin_to_cos=lowerCAmelCase_				,		downscale_freq_shift=lowerCAmelCase_    )
                 a__          =  block_out_channels[0]

            if use_timestep_embedding:
                 a__          =  block_out_channels[0] * 4
                 a__          =  TimestepEmbedding(
                     in_channels=lowerCAmelCase_				,		time_embed_dim=lowerCAmelCase_				,		act_fn=lowerCAmelCase_				,		out_dim=block_out_channels[0]				,		)

            a__          =  nn.ModuleList([]    )
            a__          =  None
            a__          =  nn.ModuleList([]    )
            a__          =  None

            # down
            a__          =  in_channels
            for i, down_block_type in enumerate(lowerCAmelCase_    ):
                 a__          =  output_channel
                 a__          =  block_out_channels[i]

                 if i == 0:
                      input_channel += extra_in_channels

                 a__          =  i == len(lowerCAmelCase_    ) - 1

                 a__          =  get_down_block(
                     lowerCAmelCase_				,		num_layers=lowerCAmelCase_				,		in_channels=lowerCAmelCase_				,		out_channels=lowerCAmelCase_				,		temb_channels=block_out_channels[0]				,		add_downsample=not is_final_block or downsample_each_block				,		)
                 self.down_blocks.append(lowerCAmelCase_    )

            # mid
            a__          =  get_mid_block(
                lowerCAmelCase_				,		in_channels=block_out_channels[-1]				,		mid_channels=block_out_channels[-1]				,		out_channels=block_out_channels[-1]				,		embed_dim=block_out_channels[0]				,		num_layers=lowerCAmelCase_				,		add_downsample=lowerCAmelCase_				,		)

            # up
            a__          =  list(reversed(lowerCAmelCase_    )    )
            a__          =  reversed_block_out_channels[0]
            if out_block_type is None:
                 a__          =  out_channels
            else:
                 a__          =  block_out_channels[0]

            for i, up_block_type in enumerate(lowerCAmelCase_    ):
                 a__          =  output_channel
                 a__          =  (
                     reversed_block_out_channels[i + 1] if i < len(lowerCAmelCase_    ) - 1 else final_upsample_channels
                 )

                 a__          =  i == len(lowerCAmelCase_    ) - 1

                 a__          =  get_up_block(
                     lowerCAmelCase_				,		num_layers=lowerCAmelCase_				,		in_channels=lowerCAmelCase_				,		out_channels=lowerCAmelCase_				,		temb_channels=block_out_channels[0]				,		add_upsample=not is_final_block				,		)
                 self.up_blocks.append(lowerCAmelCase_    )
                 a__          =  output_channel

            # out
            a__          =  norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4				,		32    )
            a__          =  get_out_block(
                out_block_type=lowerCAmelCase_				,		num_groups_out=lowerCAmelCase_				,		embed_dim=block_out_channels[0]				,		out_channels=lowerCAmelCase_				,		act_fn=lowerCAmelCase_				,		fc_dim=block_out_channels[-1] // 4				,		)





       def  lowercase		(   self:     Union[str, Any]				,		__A:     torch.FloatTensor				,		__A:     Union[torch.Tensor, float, int]				,		__A:     bool = True				,		):

            '''simple docstring'''

            a__          =  timestep
            if not torch.is_tensor(lowerCAmelCase_    ):
                 a__          =  torch.tensor([timesteps]				,		dtype=torch.long				,		device=sample.device    )
            elif torch.is_tensor(lowerCAmelCase_    ) and len(timesteps.shape    ) == 0:
                 a__          =  timesteps[None].to(sample.device    )

            a__          =  self.time_proj(lowerCAmelCase_    )
            if self.config.use_timestep_embedding:
                 a__          =  self.time_mlp(lowerCAmelCase_    )
            else:
                 a__          =  timestep_embed[..., None]
                 a__          =  timestep_embed.repeat([1, 1, sample.shape[2]]    ).to(sample.dtype    )
                 a__          =  timestep_embed.broadcast_to((sample.shape[:1] + timestep_embed.shape[1:])    )

            # 2. down
            a__          =  ()
            for downsample_block in self.down_blocks:
                 a__       ,a__          =  downsample_block(hidden_states=lowerCAmelCase_				,		temb=lowerCAmelCase_    )
                 down_block_res_samples += res_samples

            # 3. mid
            if self.mid_block:
                 a__          =  self.mid_block(lowerCAmelCase_				,		lowerCAmelCase_    )

            # 4. up
            for i, upsample_block in enumerate(self.up_blocks    ):
                 a__          =  down_block_res_samples[-1:]
                 a__          =  down_block_res_samples[:-1]
                 a__          =  upsample_block(lowerCAmelCase_				,		res_hidden_states_tuple=lowerCAmelCase_				,		temb=lowerCAmelCase_    )

            # 5. post-process
            if self.out_block:
                 a__          =  self.out_block(lowerCAmelCase_				,		lowerCAmelCase_    )

            if not return_dict:
                 return (sample,)

            return UNetaDOutput(sample=lowerCAmelCase_    )










"""simple docstring"""







__a       :	List[Any]             = frozenset(
    [
        'prompt',
        'height',
        'width',
        'guidance_scale',
        'negative_prompt',
        'prompt_embeds',
        'negative_prompt_embeds',
        'cross_attention_kwargs',
    ]
)

__a       :	Union[str, Any]             = frozenset(['prompt', 'negative_prompt'])

__a       :	Any             = frozenset([])

__a       :	Union[str, Any]             = frozenset(['image'])

__a       :	Dict             = frozenset(
    [
        'image',
        'height',
        'width',
        'guidance_scale',
    ]
)

__a       :	Dict             = frozenset(['image'])

__a       :	Dict             = frozenset(
    [
        'prompt',
        'image',
        'height',
        'width',
        'guidance_scale',
        'negative_prompt',
        'prompt_embeds',
        'negative_prompt_embeds',
    ]
)

__a       :	Optional[Any]             = frozenset(['prompt', 'image', 'negative_prompt'])

__a       :	List[Any]             = frozenset(
    [
        # Text guided image variation with an image mask
        'prompt',
        'image',
        'mask_image',
        'height',
        'width',
        'guidance_scale',
        'negative_prompt',
        'prompt_embeds',
        'negative_prompt_embeds',
    ]
)

__a       :	Union[str, Any]             = frozenset(['prompt', 'image', 'mask_image', 'negative_prompt'])

__a       :	Optional[Any]             = frozenset(
    [
        # image variation with an image mask
        'image',
        'mask_image',
        'height',
        'width',
        'guidance_scale',
    ]
)

__a       :	int             = frozenset(['image', 'mask_image'])

__a       :	Tuple             = frozenset(
    [
        'example_image',
        'image',
        'mask_image',
        'height',
        'width',
        'guidance_scale',
    ]
)

__a       :	Optional[Any]             = frozenset(['example_image', 'image', 'mask_image'])

__a       :	Optional[Any]             = frozenset(['class_labels'])

__a       :	Tuple             = frozenset(['class_labels'])

__a       :	int             = frozenset(['batch_size'])

__a       :	int             = frozenset([])

__a       :	Union[str, Any]             = frozenset(['batch_size'])

__a       :	Tuple             = frozenset([])

__a       :	Dict             = frozenset(
    [
        'prompt',
        'audio_length_in_s',
        'guidance_scale',
        'negative_prompt',
        'prompt_embeds',
        'negative_prompt_embeds',
        'cross_attention_kwargs',
    ]
)

__a       :	Dict             = frozenset(['prompt', 'negative_prompt'])
__a       :	Optional[int]             = frozenset(['input_tokens'])

__a       :	str             = frozenset(['input_tokens'])


