
"""simple docstring"""




import argparse
import os
from functools import reduce

import fairseq
import torch
from datasets import load_dataset

from transformers import WavaVecaProcessor, logging
from transformers.models.dataavec.configuration_dataavec_audio import DataaVecAudioConfig

# Copied from https://github.com/pytorch/fairseq/blob/main/examples/data2vec/models/data2vec_audio.py
from transformers.models.dataavec.dataavec_audio import DataaVecAudioModel as Dummy  # noqa: F401
from transformers.models.dataavec.modeling_dataavec_audio import DataaVecAudioForCTC, DataaVecAudioModel


logging.set_verbosity_info()
A__   :						str 							=     logging.get_logger(__name__)

A__   :						List[Any] 							=     {
    """post_extract_proj""": """feature_projection.projection""",
    """models.0.layer_norm""": """feature_projection.layer_norm""",
    """self_attn.k_proj""": """encoder.layers.*.attention.k_proj""",
    """self_attn.v_proj""": """encoder.layers.*.attention.v_proj""",
    """self_attn.q_proj""": """encoder.layers.*.attention.q_proj""",
    """self_attn.out_proj""": """encoder.layers.*.attention.out_proj""",
    """self_attn_layer_norm""": """encoder.layers.*.layer_norm""",
    """fc1""": """encoder.layers.*.feed_forward.intermediate_dense""",
    """fc2""": """encoder.layers.*.feed_forward.output_dense""",
    """final_layer_norm""": """encoder.layers.*.final_layer_norm""",
    """encoder.layer_norm""": """encoder.layer_norm""",
    """w2v_model.layer_norm""": """feature_projection.layer_norm""",
    """w2v_encoder.proj""": """lm_head""",
    """mask_emb""": """masked_spec_embed""",
}
A__   :						Optional[Any] 							=     [
    """lm_head""",
]



def 				a__			(					lowerCAmelCase						:							str							,     lowerCAmelCase						:							Optional[int]							,     lowerCAmelCase						:							int							,     lowerCAmelCase						:							Optional[int]							,     lowerCAmelCase						:							str		):






     '''simple docstring'''
     for attribute in key.split("."		):
          UpperCAmelCase__				:			int							=    getattr(lowerCAmelCase							,     lowerCAmelCase		)

     if weight_type is not None:
          UpperCAmelCase__				:			Any							=    getattr(lowerCAmelCase							,     lowerCAmelCase		).shape
     else:
          UpperCAmelCase__				:			Dict							=    hf_pointer.shape

     if hf_shape != value.shape:
          raise ValueError(
              F"Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be"
              F" {value.shape} for {full_name}"		)

     if weight_type == "weight":
          UpperCAmelCase__				:			int							=    value
     elif weight_type == "weight_g":
          UpperCAmelCase__				:			Dict							=    value
     elif weight_type == "weight_v":
          UpperCAmelCase__				:			Dict							=    value
     elif weight_type == "bias":
          UpperCAmelCase__				:			Optional[int]							=    value
     else:
          UpperCAmelCase__				:			str							=    value

     logger.info(F"{key + '.' + weight_type if weight_type is not None else ''} was initialized from {full_name}."		)



def 				a__			(					lowerCAmelCase						:							Optional[Any]							,     lowerCAmelCase						:							Any							,     lowerCAmelCase						:							Dict		):






     '''simple docstring'''
     UpperCAmelCase__				:			List[Any]							=    []
     UpperCAmelCase__				:			Union[str, Any]							=    fairseq_model.state_dict()

     if not is_headless:
          UpperCAmelCase__				:			Any							=    hf_model.dataavec_audio.feature_extractor
          UpperCAmelCase__				:			int							=    hf_model.dataavec_audio.encoder.pos_conv_embed

     else:
          UpperCAmelCase__				:			Any							=    hf_model.feature_extractor
          UpperCAmelCase__				:			Dict							=    hf_model.encoder.pos_conv_embed

     for name, value in fairseq_dict.items():
          UpperCAmelCase__				:			Any							=    False
          if "conv_layers" in name:
               load_conv_layer(
                   lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     )
               UpperCAmelCase__				:			Dict							=    True
          elif "pos_conv" in name:
               load_pos_conv_layer(
                   lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     )
               UpperCAmelCase__				:			List[Any]							=    True
          else:
               for key, mapped_key in MAPPING.items():
                    if not is_headless:
                         UpperCAmelCase__				:			Dict							=    "data2vec_audio." + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key
                    if key in name or key.split("w2v_model."		)[-1] == name.split("."		)[0]:
                         UpperCAmelCase__				:			int							=    True
                         if "*" in mapped_key:
                              UpperCAmelCase__				:			Optional[int]							=    name.split(lowerCAmelCase		)[0].split("."		)[-2]
                              UpperCAmelCase__				:			Dict							=    mapped_key.replace("*"							,     lowerCAmelCase		)
                         if "weight_g" in name:
                              UpperCAmelCase__				:			Any							=    "weight_g"
                         elif "weight_v" in name:
                              UpperCAmelCase__				:			Any							=    "weight_v"
                         elif "bias" in name:
                              UpperCAmelCase__				:			Optional[Any]							=    "bias"
                         elif "weight" in name:
                              # TODO: don't match quantizer.weight_proj
                              UpperCAmelCase__				:			Tuple							=    "weight"
                         else:
                              UpperCAmelCase__				:			Tuple							=    None
                         set_recursively(lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase		)
                    continue
          if not is_used:
               unused_weights.append(lowerCAmelCase		)

     logger.warning(F"Unused weights: {unused_weights}"		)



def 				a__			(					lowerCAmelCase						:							Optional[Any]							,     lowerCAmelCase						:							int		):






     '''simple docstring'''
     UpperCAmelCase__				:			Optional[int]							=    path.split("."		)
     return reduce(lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase		)



def 				a__			(					lowerCAmelCase						:							Tuple							,     lowerCAmelCase						:							List[Any]							,     lowerCAmelCase						:							List[Any]							,     lowerCAmelCase						:							str		):






     '''simple docstring'''
     UpperCAmelCase__				:			List[Any]							=    access_by_string(lowerCAmelCase							,     lowerCAmelCase		)
     UpperCAmelCase__				:			int							=    hf_weight.data

     if fsq_value.shape != hf_value.shape:
          raise ValueError(F"{full_name} has size {fsq_value.shape}, but {hf_value.shape} was found."		)
     UpperCAmelCase__				:			List[str]							=    fsq_value
     logger.info(F"{full_name} was correctly initialized from {hf_weight_path}."		)



def 				a__			(					lowerCAmelCase						:							List[str]							,     lowerCAmelCase						:							str							,     lowerCAmelCase						:							Tuple							,     lowerCAmelCase						:							Any		):






     '''simple docstring'''
     UpperCAmelCase__				:			List[Any]							=    full_name.split("conv_layers."		)[-1]
     UpperCAmelCase__				:			List[str]							=    name.split("."		)
     UpperCAmelCase__				:			Optional[Any]							=    int(items[0]		)
     UpperCAmelCase__				:			List[str]							=    int(items[1]		)

     UpperCAmelCase__				:			str							=    name.split("."		)[-1]
     if type_id == 0:
          UpperCAmelCase__				:			Any							=    "conv"
     elif type_id == 2:
          UpperCAmelCase__				:			Optional[Any]							=    "layer_norm"
     else:
          unused_weights.append(lowerCAmelCase		)
          return

     set_weights(lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     F"conv_layers.{layer_id}.{layer_type}.{weight_type}"		)



def 				a__			(					lowerCAmelCase						:							Any							,     lowerCAmelCase						:							List[str]							,     lowerCAmelCase						:							Optional[int]							,     lowerCAmelCase						:							Union[str, Any]		):






     '''simple docstring'''
     UpperCAmelCase__				:			int							=    full_name.split("pos_conv."		)[-1]
     UpperCAmelCase__				:			int							=    name.split("."		)
     UpperCAmelCase__				:			Dict							=    int(items[0]		)
     UpperCAmelCase__				:			Optional[Any]							=    int(items[1]		)

     UpperCAmelCase__				:			List[Any]							=    name.split("."		)[-1]
     if type_id != 0:
          unused_weights.append(lowerCAmelCase		)
          return
     else:
          UpperCAmelCase__				:			List[Any]							=    "conv"

     set_weights(lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     F"layers.{layer_id}.{layer_type}.{weight_type}"		)



@torch.no_grad()
def 				a__			(					lowerCAmelCase						:							Any							,     lowerCAmelCase						:							Union[str, Any]							,     lowerCAmelCase						:							Any=None							,     lowerCAmelCase						:							Dict=None							,     lowerCAmelCase						:							Any=True		):






     '''simple docstring'''
     if config_path is not None:
          UpperCAmelCase__				:			str							=    DataaVecAudioConfig.from_pretrained(lowerCAmelCase		)
     else:
          UpperCAmelCase__				:			List[Any]							=    DataaVecAudioConfig()

     if not is_finetuned:
          # Modify final_proj layer name
          UpperCAmelCase__				:			Dict							=    DataaVecAudioModel(lowerCAmelCase		)
          UpperCAmelCase__				:			Dict							=    os.path.dirname(lowerCAmelCase		)

          UpperCAmelCase__				:			List[str]							=    torch.load(lowerCAmelCase		)
          UpperCAmelCase__				:			Any							=    state_dict["model"].pop("final_proj.0.weight"		)
          UpperCAmelCase__				:			List[str]							=    state_dict["model"].pop("final_proj.0.bias"		)
          UpperCAmelCase__				:			Optional[int]							=    os.path.join(lowerCAmelCase							,     "converted.pt"		)
          torch.save(lowerCAmelCase							,     lowerCAmelCase		)
     else:
          UpperCAmelCase__				:			List[Any]							=    DataaVecAudioForCTC(lowerCAmelCase		)
          UpperCAmelCase__				:			Union[str, Any]							=    checkpoint_path

     def load_dataavec(lowerCAmelCase						:							Optional[int]		):
          UpperCAmelCase__				,						UpperCAmelCase__				,						UpperCAmelCase__				:			Any							=    fairseq.checkpoint_utils.load_model_ensemble_and_task([path]		)
          return model[0].eval()

     UpperCAmelCase__				:			Tuple							=    load_dataavec(lowerCAmelCase		)

     recursively_load_weights(lowerCAmelCase							,     lowerCAmelCase							,     not is_finetuned		)

     UpperCAmelCase__				:			int							=    WavaVecaProcessor.from_pretrained("facebook/wav2vec2-large-lv60"		)

     UpperCAmelCase__				:			Any							=    load_dataset("patrickvonplaten/librispeech_asr_dummy"							,     "clean"							,     split="validation"		)
     UpperCAmelCase__				:			Optional[Any]							=    [x["array"] for x in ds[:4]["audio"]]

     UpperCAmelCase__				:			Optional[Any]							=    processor(lowerCAmelCase							,     return_tensors="pt"							,     padding=lowerCAmelCase		)

     UpperCAmelCase__				:			str							=    inputs.input_values
     UpperCAmelCase__				:			Union[str, Any]							=    inputs.attention_mask
     #    input_values = inputs.input_values[:, :-1]
     #    attention_mask = inputs.attention_mask[:, :-1]

     hf_wavavec.eval()
     model.eval()
     if is_finetuned:
          UpperCAmelCase__				:			Optional[Any]							=    model(source=lowerCAmelCase							,     padding_mask=(1 - attention_mask)							,     mask=lowerCAmelCase							,     features_only=lowerCAmelCase		)[
              "encoder_out"
          ].transpose(0							,     1		)
          UpperCAmelCase__				:			List[Any]							=    hf_wavavec(lowerCAmelCase							,     attention_mask=lowerCAmelCase		)["logits"]

          UpperCAmelCase__				:			Dict							=    torch.argmax(lowerCAmelCase							,     dim=-1		)
          UpperCAmelCase__				:			str							=    processor.batch_decode(lowerCAmelCase		)

          print(F"Expected Output: {ds[:4]['text']}, Pred: {output_string}"		)
     else:
          UpperCAmelCase__				:			List[Any]							=    model(source=lowerCAmelCase							,     padding_mask=(1 - attention_mask)							,     mask=lowerCAmelCase							,     features_only=lowerCAmelCase		)[
              "layer_results"
          ][-1][0].transpose(0							,     1		)
          UpperCAmelCase__				:			Union[str, Any]							=    hf_wavavec(lowerCAmelCase							,     attention_mask=lowerCAmelCase		)["last_hidden_state"]

     print(our_output.shape							,     their_output.shape		)
     UpperCAmelCase__				:			int							=    torch.max(torch.abs(our_output - their_output		)		).item()
     print(F"max_absolute_diff = {max_absolute_diff}"		)  # ~ 1e-7
     UpperCAmelCase__				:			str							=    torch.allclose(lowerCAmelCase							,     lowerCAmelCase							,     atol=1E-3		)
     print("Do both models output the same tensors?"							,     "🔥" if success else "💩"		)
     if not success:
          raise Exception("Something went wRoNg"		)

     hf_wavavec.save_pretrained(lowerCAmelCase		)

     if is_finetuned:
          processor.save_pretrained(lowerCAmelCase		)
     else:
          processor.feature_extractor.save_pretrained(lowerCAmelCase		)


if __name__ == "__main__":
    A__   :						Any 							=     argparse.ArgumentParser()
    parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
    parser.add_argument("""--checkpoint_path""", default=None, type=str, help="""Path to fairseq checkpoint""")
    parser.add_argument("""--dict_path""", default=None, type=str, help="""Path to dict of fine-tuned model""")
    parser.add_argument("""--config_path""", default=None, type=str, help="""Path to hf config.json of model to convert""")
    parser.add_argument(
        """--not_finetuned""", action="""store_true""", help="""Whether the model to convert is a fine-tuned model or not"""
    )
    A__   :						Tuple 							=     parser.parse_args()
    convert_wavaveca_checkpoint(
        args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
    )





"""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 UNetMidBlockaD, get_down_block, get_up_block







@dataclass
class 	_lowercase    (							lowerCAmelCase_					):


     '''simple docstring'''





     _A		    =				42







class 	_lowercase    (							lowerCAmelCase_						, lowerCAmelCase_					):


     '''simple docstring'''





     @register_to_config
     def __init__(   self			,      __UpperCamelCase = None			,      __UpperCamelCase = 3			,      __UpperCamelCase = 3			,      __UpperCamelCase = False			,      __UpperCamelCase = "positional"			,      __UpperCamelCase = 0			,      __UpperCamelCase = True			,      __UpperCamelCase = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")			,      __UpperCamelCase = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")			,      __UpperCamelCase = (2_24, 4_48, 6_72, 8_96)			,      __UpperCamelCase = 2			,      __UpperCamelCase = 1			,      __UpperCamelCase = 1			,      __UpperCamelCase = "conv"			,      __UpperCamelCase = "conv"			,      __UpperCamelCase = "silu"			,      __UpperCamelCase = 8			,      __UpperCamelCase = 32			,      __UpperCamelCase = 1E-5			,      __UpperCamelCase = "default"			,      __UpperCamelCase = True			,      __UpperCamelCase = None			,      __UpperCamelCase = None			,      )->       Any:
          super().__init__()

          UpperCAmelCase__				:			Optional[int]							=    sample_size
          UpperCAmelCase__				:			Any							=    block_out_channels[0] * 4

          # Check inputs
          if len(__UpperCamelCase	) != len(__UpperCamelCase	):
               raise ValueError(
                   F"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."	)

          if len(__UpperCamelCase	) != len(__UpperCamelCase	):
               raise ValueError(
                   F"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."	)

          # input
          UpperCAmelCase__				:			List[Any]							=    nn.Convad(__UpperCamelCase			,      block_out_channels[0]			,      kernel_size=3			,      padding=(1, 1)	)

          # time
          if time_embedding_type == "fourier":
               UpperCAmelCase__				:			int							=    GaussianFourierProjection(embedding_size=block_out_channels[0]			,      scale=16	)
               UpperCAmelCase__				:			Dict							=    2 * block_out_channels[0]
          elif time_embedding_type == "positional":
               UpperCAmelCase__				:			List[str]							=    Timesteps(block_out_channels[0]			,      __UpperCamelCase			,      __UpperCamelCase	)
               UpperCAmelCase__				:			Any							=    block_out_channels[0]

          UpperCAmelCase__				:			List[str]							=    TimestepEmbedding(__UpperCamelCase			,      __UpperCamelCase	)

          # class embedding
          if class_embed_type is None and num_class_embeds is not None:
               UpperCAmelCase__				:			Optional[Any]							=    nn.Embedding(__UpperCamelCase			,      __UpperCamelCase	)
          elif class_embed_type == "timestep":
               UpperCAmelCase__				:			str							=    TimestepEmbedding(__UpperCamelCase			,      __UpperCamelCase	)
          elif class_embed_type == "identity":
               UpperCAmelCase__				:			Union[str, Any]							=    nn.Identity(__UpperCamelCase			,      __UpperCamelCase	)
          else:
               UpperCAmelCase__				:			Optional[int]							=    None

          UpperCAmelCase__				:			str							=    nn.ModuleList([]	)
          UpperCAmelCase__				:			Any							=    None
          UpperCAmelCase__				:			Tuple							=    nn.ModuleList([]	)

          # down
          UpperCAmelCase__				:			Any							=    block_out_channels[0]
          for i, down_block_type in enumerate(__UpperCamelCase	):
               UpperCAmelCase__				:			List[Any]							=    output_channel
               UpperCAmelCase__				:			Optional[int]							=    block_out_channels[i]
               UpperCAmelCase__				:			Any							=    i == len(__UpperCamelCase	) - 1

               UpperCAmelCase__				:			List[Any]							=    get_down_block(
                   __UpperCamelCase			,      num_layers=__UpperCamelCase			,      in_channels=__UpperCamelCase			,      out_channels=__UpperCamelCase			,      temb_channels=__UpperCamelCase			,      add_downsample=not is_final_block			,      resnet_eps=__UpperCamelCase			,      resnet_act_fn=__UpperCamelCase			,      resnet_groups=__UpperCamelCase			,      attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel			,      downsample_padding=__UpperCamelCase			,      resnet_time_scale_shift=__UpperCamelCase			,      downsample_type=__UpperCamelCase			,      )
               self.down_blocks.append(__UpperCamelCase	)

          # mid
          UpperCAmelCase__				:			List[str]							=    UNetMidBlockaD(
              in_channels=block_out_channels[-1]			,      temb_channels=__UpperCamelCase			,      resnet_eps=__UpperCamelCase			,      resnet_act_fn=__UpperCamelCase			,      output_scale_factor=__UpperCamelCase			,      resnet_time_scale_shift=__UpperCamelCase			,      attention_head_dim=attention_head_dim if attention_head_dim is not None else block_out_channels[-1]			,      resnet_groups=__UpperCamelCase			,      add_attention=__UpperCamelCase			,      )

          # up
          UpperCAmelCase__				:			Tuple							=    list(reversed(__UpperCamelCase	)	)
          UpperCAmelCase__				:			Optional[Any]							=    reversed_block_out_channels[0]
          for i, up_block_type in enumerate(__UpperCamelCase	):
               UpperCAmelCase__				:			int							=    output_channel
               UpperCAmelCase__				:			List[Any]							=    reversed_block_out_channels[i]
               UpperCAmelCase__				:			Dict							=    reversed_block_out_channels[min(i + 1			,      len(__UpperCamelCase	) - 1	)]

               UpperCAmelCase__				:			int							=    i == len(__UpperCamelCase	) - 1

               UpperCAmelCase__				:			Dict							=    get_up_block(
                   __UpperCamelCase			,      num_layers=layers_per_block + 1			,      in_channels=__UpperCamelCase			,      out_channels=__UpperCamelCase			,      prev_output_channel=__UpperCamelCase			,      temb_channels=__UpperCamelCase			,      add_upsample=not is_final_block			,      resnet_eps=__UpperCamelCase			,      resnet_act_fn=__UpperCamelCase			,      resnet_groups=__UpperCamelCase			,      attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel			,      resnet_time_scale_shift=__UpperCamelCase			,      upsample_type=__UpperCamelCase			,      )
               self.up_blocks.append(__UpperCamelCase	)
               UpperCAmelCase__				:			Dict							=    output_channel

          # out
          UpperCAmelCase__				:			str							=    norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4			,      32	)
          UpperCAmelCase__				:			Tuple							=    nn.GroupNorm(num_channels=block_out_channels[0]			,      num_groups=__UpperCamelCase			,      eps=__UpperCamelCase	)
          UpperCAmelCase__				:			Any							=    nn.SiLU()
          UpperCAmelCase__				:			Optional[int]							=    nn.Convad(block_out_channels[0]			,      __UpperCamelCase			,      kernel_size=3			,      padding=1	)



     def 			lowerCAmelCase__							(   self			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase = None			,      __UpperCamelCase = True			,      )->       Union[UNetaDOutput, Tuple]:
          # 0. center input if necessary
          if self.config.center_input_sample:
               UpperCAmelCase__				:			Any							=    2 * sample - 1.0

          # 1. time
          UpperCAmelCase__				:			Tuple							=    timestep
          if not torch.is_tensor(__UpperCamelCase	):
               UpperCAmelCase__				:			Optional[int]							=    torch.tensor([timesteps]			,      dtype=torch.long			,      device=sample.device	)
          elif torch.is_tensor(__UpperCamelCase	) and len(timesteps.shape	) == 0:
               UpperCAmelCase__				:			List[Any]							=    timesteps[None].to(sample.device	)

          # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
          UpperCAmelCase__				:			Dict							=    timesteps * torch.ones(sample.shape[0]			,      dtype=timesteps.dtype			,      device=timesteps.device	)

          UpperCAmelCase__				:			List[str]							=    self.time_proj(__UpperCamelCase	)

          # timesteps does not contain any weights and will always return f32 tensors
          # but time_embedding might actually be running in fp16. so we need to cast here.
          # there might be better ways to encapsulate this.
          UpperCAmelCase__				:			Dict							=    t_emb.to(dtype=self.dtype	)
          UpperCAmelCase__				:			Optional[int]							=    self.time_embedding(__UpperCamelCase	)

          if self.class_embedding is not None:
               if class_labels is None:
                    raise ValueError("class_labels should be provided when doing class conditioning"	)

               if self.config.class_embed_type == "timestep":
                    UpperCAmelCase__				:			List[Any]							=    self.time_proj(__UpperCamelCase	)

               UpperCAmelCase__				:			List[str]							=    self.class_embedding(__UpperCamelCase	).to(dtype=self.dtype	)
               UpperCAmelCase__				:			Optional[int]							=    emb + class_emb

          # 2. pre-process
          UpperCAmelCase__				:			List[str]							=    sample
          UpperCAmelCase__				:			str							=    self.conv_in(__UpperCamelCase	)

          # 3. down
          UpperCAmelCase__				:			List[str]							=    (sample,)
          for downsample_block in self.down_blocks:
               if hasattr(__UpperCamelCase			,      "skip_conv"	):
                    UpperCAmelCase__				,						UpperCAmelCase__				,						UpperCAmelCase__				:			Optional[Any]							=    downsample_block(
                        hidden_states=__UpperCamelCase			,      temb=__UpperCamelCase			,      skip_sample=__UpperCamelCase	)
               else:
                    UpperCAmelCase__				,						UpperCAmelCase__				:			Optional[Any]							=    downsample_block(hidden_states=__UpperCamelCase			,      temb=__UpperCamelCase	)

               down_block_res_samples += res_samples

          # 4. mid
          UpperCAmelCase__				:			Union[str, Any]							=    self.mid_block(__UpperCamelCase			,      __UpperCamelCase	)

          # 5. up
          UpperCAmelCase__				:			Dict							=    None
          for upsample_block in self.up_blocks:
               UpperCAmelCase__				:			Union[str, Any]							=    down_block_res_samples[-len(upsample_block.resnets	) :]
               UpperCAmelCase__				:			Optional[Any]							=    down_block_res_samples[: -len(upsample_block.resnets	)]

               if hasattr(__UpperCamelCase			,      "skip_conv"	):
                    UpperCAmelCase__				,						UpperCAmelCase__				:			str							=    upsample_block(__UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase	)
               else:
                    UpperCAmelCase__				:			Tuple							=    upsample_block(__UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase	)

        # 6. post-process
          UpperCAmelCase__				:			List[Any]							=    self.conv_norm_out(__UpperCamelCase	)
          UpperCAmelCase__				:			Any							=    self.conv_act(__UpperCamelCase	)
          UpperCAmelCase__				:			Optional[Any]							=    self.conv_out(__UpperCamelCase	)

          if skip_sample is not None:
               sample += skip_sample

          if self.config.time_embedding_type == "fourier":
               UpperCAmelCase__				:			str							=    timesteps.reshape((sample.shape[0], *([1] * len(sample.shape[1:]	)))	)
               UpperCAmelCase__				:			int							=    sample / timesteps

          if not return_dict:
               return (sample,)

          return UNetaDOutput(sample=__UpperCamelCase	)



