





# This model implementation is heavily inspired by https://github.com/haofanwang/ControlNet-for-Diffusers/

import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    DDIMScheduler,
    StableDiffusionControlNetImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel
from diffusers.utils import floats_tensor, load_image, load_numpy, randn_tensor, slow, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
)


enable_full_determinism()



class     _UpperCAmelCase   (   _UpperCamelCase		,   _UpperCamelCase		,   _UpperCamelCase		,   unittest.TestCase   ):

  """simple docstring"""


  a_      =			StableDiffusionControlNetImgaImgPipeline
  a_      =			TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
  a_      =			TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
  a_      =			IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"""control_image"""}   )
  a_      =			IMAGE_TO_IMAGE_IMAGE_PARAMS






  def     lowercase		(    self			:		List[str]					)							->						List[str]:
     torch.manual_seed(0					)
     __lowerCAmelCase							     =     UNetaDConditionModel(
         block_out_channels=(3_2, 6_4) ,  layers_per_block=2 ,  sample_size=3_2 ,  in_channels=4 ,  out_channels=4 ,  down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D') ,  up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D') ,  cross_attention_dim=3_2 ,  )
     torch.manual_seed(0					)
     __lowerCAmelCase							     =     ControlNetModel(
         block_out_channels=(3_2, 6_4) ,  layers_per_block=2 ,  in_channels=4 ,  down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D') ,  cross_attention_dim=3_2 ,  conditioning_embedding_out_channels=(1_6, 3_2) ,  )
     torch.manual_seed(0					)
     __lowerCAmelCase							     =     DDIMScheduler(
         beta_start=0.0_00_85 ,  beta_end=0.0_12 ,  beta_schedule='scaled_linear' ,  clip_sample=lowerCAmelCase_ ,  set_alpha_to_one=lowerCAmelCase_ ,  )
     torch.manual_seed(0					)
     __lowerCAmelCase							     =     AutoencoderKL(
         block_out_channels=[3_2, 6_4] ,  in_channels=3 ,  out_channels=3 ,  down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D'] ,  up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D'] ,  latent_channels=4 ,  )
     torch.manual_seed(0					)
     __lowerCAmelCase							     =     CLIPTextConfig(
         bos_token_id=0 ,  eos_token_id=2 ,  hidden_size=3_2 ,  intermediate_size=3_7 ,  layer_norm_eps=1e-05 ,  num_attention_heads=4 ,  num_hidden_layers=5 ,  pad_token_id=1 ,  vocab_size=1_0_0_0 ,  )
     __lowerCAmelCase							     =     CLIPTextModel(lowerCAmelCase_					)
     __lowerCAmelCase							     =     CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'					)

     __lowerCAmelCase							     =     {
         'unet': unet,
         'controlnet': controlnet,
         'scheduler': scheduler,
         'vae': vae,
         'text_encoder': text_encoder,
         'tokenizer': tokenizer,
         'safety_checker': None,
         'feature_extractor': None,
     }
     return components






  def     lowercase		(    self			:		List[str] ,  lowerCAmelCase_			:		Union[str, Any] ,  lowerCAmelCase_			:		str=0					)							->						List[str]:
     if str(lowerCAmelCase_					).startswith('mps'					):
        __lowerCAmelCase							     =     torch.manual_seed(lowerCAmelCase_					)
     else:
        __lowerCAmelCase							     =     torch.Generator(device=lowerCAmelCase_					).manual_seed(lowerCAmelCase_					)

     __lowerCAmelCase							     =     2
     __lowerCAmelCase							     =     randn_tensor(
         (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor) ,  generator=lowerCAmelCase_ ,  device=torch.device(lowerCAmelCase_					) ,  )
     __lowerCAmelCase							     =     floats_tensor(control_image.shape ,  rng=random.Random(lowerCAmelCase_					)					).to(lowerCAmelCase_					)
     __lowerCAmelCase							     =     image.cpu().permute(0 ,  2 ,  3 ,  1					)[0]
     __lowerCAmelCase							     =     Image.fromarray(np.uinta(lowerCAmelCase_					)					).convert('RGB'					).resize((6_4, 6_4)					)
     __lowerCAmelCase							     =     {
         'prompt': 'A painting of a squirrel eating a burger',
         'generator': generator,
         'num_inference_steps': 2,
         'guidance_scale': 6.0,
         'output_type': 'numpy',
         'image': image,
         'control_image': control_image,
     }

     return inputs






  def     lowercase		(    self			:		List[str]					)							->						List[Any]:
     return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3					)






  @unittest.skipIf(
      torch_device != 'cuda' or not is_xformers_available() ,  reason='XFormers attention is only available with CUDA and `xformers` installed' ,  )
  def     lowercase		(    self			:		Any					)							->						int:
     self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3					)






  def     lowercase		(    self			:		str					)							->						Tuple:
     self._test_inference_batch_single_identical(expected_max_diff=2e-3					)





class     _UpperCAmelCase   (   _UpperCamelCase		,   _UpperCamelCase		,   unittest.TestCase   ):

  """simple docstring"""


  a_      =			StableDiffusionControlNetImgaImgPipeline
  a_      =			TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
  a_      =			TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
  a_      =			frozenset([]   )  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess






  def     lowercase		(    self			:		List[Any]					)							->						Tuple:
     torch.manual_seed(0					)
     __lowerCAmelCase							     =     UNetaDConditionModel(
         block_out_channels=(3_2, 6_4) ,  layers_per_block=2 ,  sample_size=3_2 ,  in_channels=4 ,  out_channels=4 ,  down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D') ,  up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D') ,  cross_attention_dim=3_2 ,  )
     torch.manual_seed(0					)

     def init_weights(lowerCAmelCase_			:		List[str]					):
        if isinstance(lowerCAmelCase_ ,  torch.nn.Convad					):
           torch.nn.init.normal(m.weight					)
           m.bias.data.fill_(1.0					)

     __lowerCAmelCase							     =     ControlNetModel(
         block_out_channels=(3_2, 6_4) ,  layers_per_block=2 ,  in_channels=4 ,  down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D') ,  cross_attention_dim=3_2 ,  conditioning_embedding_out_channels=(1_6, 3_2) ,  )
     controlneta.controlnet_down_blocks.apply(lowerCAmelCase_					)

     torch.manual_seed(0					)
     __lowerCAmelCase							     =     ControlNetModel(
         block_out_channels=(3_2, 6_4) ,  layers_per_block=2 ,  in_channels=4 ,  down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D') ,  cross_attention_dim=3_2 ,  conditioning_embedding_out_channels=(1_6, 3_2) ,  )
     controlneta.controlnet_down_blocks.apply(lowerCAmelCase_					)

     torch.manual_seed(0					)
     __lowerCAmelCase							     =     DDIMScheduler(
         beta_start=0.0_00_85 ,  beta_end=0.0_12 ,  beta_schedule='scaled_linear' ,  clip_sample=lowerCAmelCase_ ,  set_alpha_to_one=lowerCAmelCase_ ,  )
     torch.manual_seed(0					)
     __lowerCAmelCase							     =     AutoencoderKL(
         block_out_channels=[3_2, 6_4] ,  in_channels=3 ,  out_channels=3 ,  down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D'] ,  up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D'] ,  latent_channels=4 ,  )
     torch.manual_seed(0					)
     __lowerCAmelCase							     =     CLIPTextConfig(
         bos_token_id=0 ,  eos_token_id=2 ,  hidden_size=3_2 ,  intermediate_size=3_7 ,  layer_norm_eps=1e-05 ,  num_attention_heads=4 ,  num_hidden_layers=5 ,  pad_token_id=1 ,  vocab_size=1_0_0_0 ,  )
     __lowerCAmelCase							     =     CLIPTextModel(lowerCAmelCase_					)
     __lowerCAmelCase							     =     CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'					)

     __lowerCAmelCase							     =     MultiControlNetModel([controlneta, controlneta]					)

     __lowerCAmelCase							     =     {
         'unet': unet,
         'controlnet': controlnet,
         'scheduler': scheduler,
         'vae': vae,
         'text_encoder': text_encoder,
         'tokenizer': tokenizer,
         'safety_checker': None,
         'feature_extractor': None,
     }
     return components






  def     lowercase		(    self			:		List[Any] ,  lowerCAmelCase_			:		str ,  lowerCAmelCase_			:		Dict=0					)							->						List[Any]:
     if str(lowerCAmelCase_					).startswith('mps'					):
        __lowerCAmelCase							     =     torch.manual_seed(lowerCAmelCase_					)
     else:
        __lowerCAmelCase							     =     torch.Generator(device=lowerCAmelCase_					).manual_seed(lowerCAmelCase_					)

     __lowerCAmelCase							     =     2

     __lowerCAmelCase							     =     [
         randn_tensor(
             (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor) ,  generator=lowerCAmelCase_ ,  device=torch.device(lowerCAmelCase_					) ,  ),
         randn_tensor(
             (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor) ,  generator=lowerCAmelCase_ ,  device=torch.device(lowerCAmelCase_					) ,  ),
     ]

     __lowerCAmelCase							     =     floats_tensor(control_image[0].shape ,  rng=random.Random(lowerCAmelCase_					)					).to(lowerCAmelCase_					)
     __lowerCAmelCase							     =     image.cpu().permute(0 ,  2 ,  3 ,  1					)[0]
     __lowerCAmelCase							     =     Image.fromarray(np.uinta(lowerCAmelCase_					)					).convert('RGB'					).resize((6_4, 6_4)					)
     __lowerCAmelCase							     =     {
         'prompt': 'A painting of a squirrel eating a burger',
         'generator': generator,
         'num_inference_steps': 2,
         'guidance_scale': 6.0,
         'output_type': 'numpy',
         'image': image,
         'control_image': control_image,
     }

     return inputs






  def     lowercase		(    self			:		List[str]					)							->						Dict:
     __lowerCAmelCase							     =     self.get_dummy_components()
     __lowerCAmelCase							     =     self.pipeline_class(**lowerCAmelCase_					)
     pipe.to(lowerCAmelCase_					)

     __lowerCAmelCase							     =     10.0
     __lowerCAmelCase							     =     4

     __lowerCAmelCase							     =     self.get_dummy_inputs(lowerCAmelCase_					)
     __lowerCAmelCase							     =     steps
     __lowerCAmelCase							     =     scale
     __lowerCAmelCase							     =     pipe(**lowerCAmelCase_					)[0]

     __lowerCAmelCase							     =     self.get_dummy_inputs(lowerCAmelCase_					)
     __lowerCAmelCase							     =     steps
     __lowerCAmelCase							     =     scale
     __lowerCAmelCase							     =     pipe(**lowerCAmelCase_ ,  control_guidance_start=0.1 ,  control_guidance_end=0.2					)[0]

     __lowerCAmelCase							     =     self.get_dummy_inputs(lowerCAmelCase_					)
     __lowerCAmelCase							     =     steps
     __lowerCAmelCase							     =     scale
     __lowerCAmelCase							     =     pipe(**lowerCAmelCase_ ,  control_guidance_start=[0.1, 0.3] ,  control_guidance_end=[0.2, 0.7]					)[0]

     __lowerCAmelCase							     =     self.get_dummy_inputs(lowerCAmelCase_					)
     __lowerCAmelCase							     =     steps
     __lowerCAmelCase							     =     scale
     __lowerCAmelCase							     =     pipe(**lowerCAmelCase_ ,  control_guidance_start=0.4 ,  control_guidance_end=[0.5, 0.8]					)[0]

     # make sure that all outputs are different
     assert np.sum(np.abs(output_a - output_a					)					) > 1e-3
     assert np.sum(np.abs(output_a - output_a					)					) > 1e-3
     assert np.sum(np.abs(output_a - output_a					)					) > 1e-3






  def     lowercase		(    self			:		Optional[int]					)							->						List[str]:
     return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3					)






  @unittest.skipIf(
      torch_device != 'cuda' or not is_xformers_available() ,  reason='XFormers attention is only available with CUDA and `xformers` installed' ,  )
  def     lowercase		(    self			:		Optional[Any]					)							->						str:
     self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3					)






  def     lowercase		(    self			:		List[str]					)							->						int:
     self._test_inference_batch_single_identical(expected_max_diff=2e-3					)






  def     lowercase		(    self			:		Tuple					)							->						Dict:
     __lowerCAmelCase							     =     self.get_dummy_components()
     __lowerCAmelCase							     =     self.pipeline_class(**lowerCAmelCase_					)
     pipe.to(lowerCAmelCase_					)
     pipe.set_progress_bar_config(disable=lowerCAmelCase_					)
     with tempfile.TemporaryDirectory() as tmpdir:
        try:
           # save_pretrained is not implemented for Multi-ControlNet
           pipe.save_pretrained(lowerCAmelCase_					)
        except NotImplementedError:
           pass





@slow
@require_torch_gpu
class     _UpperCAmelCase   (   unittest.TestCase   ):

  """simple docstring"""


  def     lowercase		(    self			:		Dict					)							->						List[str]:
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()






  def     lowercase		(    self			:		Optional[int]					)							->						List[str]:
     __lowerCAmelCase							     =     ControlNetModel.from_pretrained('lllyasviel/sd-controlnet-canny'					)

     __lowerCAmelCase							     =     StableDiffusionControlNetImgaImgPipeline.from_pretrained(
         'runwayml/stable-diffusion-v1-5' ,  safety_checker=lowerCAmelCase_ ,  controlnet=lowerCAmelCase_					)
     pipe.enable_model_cpu_offload()
     pipe.set_progress_bar_config(disable=lowerCAmelCase_					)

     __lowerCAmelCase							     =     torch.Generator(device='cpu'					).manual_seed(0					)
     __lowerCAmelCase							     =     'evil space-punk bird'
     __lowerCAmelCase							     =     load_image(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png'					).resize((5_1_2, 5_1_2)					)
     __lowerCAmelCase							     =     load_image(
         'https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png'					).resize((5_1_2, 5_1_2)					)

     __lowerCAmelCase							     =     pipe(
         lowerCAmelCase_ ,  lowerCAmelCase_ ,  control_image=lowerCAmelCase_ ,  generator=lowerCAmelCase_ ,  output_type='np' ,  num_inference_steps=5_0 ,  strength=0.6 ,  )

     __lowerCAmelCase							     =     output.images[0]

     assert image.shape == (5_1_2, 5_1_2, 3)

     __lowerCAmelCase							     =     load_numpy(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/img2img.npy'					)

     assert np.abs(expected_image - image					).max() < 9e-2








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     _UpperCAmelCase   (   _UpperCamelCase   ):

  """simple docstring"""


  a_      =			42





class     _UpperCAmelCase   (   _UpperCamelCase		,   _UpperCamelCase   ):

  """simple docstring"""


  @register_to_config
  def __init__(    self			:		List[Any] ,  lowerCAmelCase_			:		int = 6_5_5_3_6 ,  lowerCAmelCase_			:		Optional[int] = None ,  lowerCAmelCase_			:		int = 2 ,  lowerCAmelCase_			:		int = 2 ,  lowerCAmelCase_			:		int = 0 ,  lowerCAmelCase_			:		str = "fourier" ,  lowerCAmelCase_			:		bool = True ,  lowerCAmelCase_			:		bool = False ,  lowerCAmelCase_			:		float = 0.0 ,  lowerCAmelCase_			:		Tuple[str] = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D") ,  lowerCAmelCase_			:		Tuple[str] = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip") ,  lowerCAmelCase_			:		Tuple[str] = "UNetMidBlock1D" ,  lowerCAmelCase_			:		str = None ,  lowerCAmelCase_			:		Tuple[int] = (3_2, 3_2, 6_4) ,  lowerCAmelCase_			:		str = None ,  lowerCAmelCase_			:		int = 8 ,  lowerCAmelCase_			:		int = 1 ,  lowerCAmelCase_			:		bool = False ,  )							->						Optional[int]:
     super().__init__()
     __lowerCAmelCase							     =     sample_size

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

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

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

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

        if i == 0:
           input_channel += extra_in_channels

        __lowerCAmelCase							     =     i == len(lowerCAmelCase_					) - 1

        __lowerCAmelCase							     =     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
     __lowerCAmelCase							     =     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
     __lowerCAmelCase							     =     list(reversed(lowerCAmelCase_					)					)
     __lowerCAmelCase							     =     reversed_block_out_channels[0]
     if out_block_type is None:
        __lowerCAmelCase							     =     out_channels
     else:
        __lowerCAmelCase							     =     block_out_channels[0]

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

        __lowerCAmelCase							     =     i == len(lowerCAmelCase_					) - 1

        __lowerCAmelCase							     =     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_					)
        __lowerCAmelCase							     =     output_channel

     # out
     __lowerCAmelCase							     =     norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4 ,  3_2					)
     __lowerCAmelCase							     =     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			:		Optional[Any] ,  lowerCAmelCase_			:		torch.FloatTensor ,  lowerCAmelCase_			:		Union[torch.Tensor, float, int] ,  lowerCAmelCase_			:		bool = True ,  )							->						Union[UNetaDOutput, Tuple]:
     __lowerCAmelCase							     =     timestep
     if not torch.is_tensor(lowerCAmelCase_					):
        __lowerCAmelCase							     =     torch.tensor([timesteps] ,  dtype=torch.long ,  device=sample.device					)
     elif torch.is_tensor(lowerCAmelCase_					) and len(timesteps.shape					) == 0:
        __lowerCAmelCase							     =     timesteps[None].to(sample.device					)

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

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

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

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

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

     if not return_dict:
        return (sample,)

     return UNetaDOutput(sample=lowerCAmelCase_					)

