


"""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      __lowerCamelCase	(       a__					):

     '''simple docstring'''
     A_							:						torch.FloatTensor







class      __lowerCamelCase	(       a__	,       a__					):

     '''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 = (224, 448, 672, 896)     ,					__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     ,					)							->				Union[str, Any]:
          super().__init__()

          _a     						=					sample_size
          _a     						=					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
          _a     						=					nn.Convad(__UpperCAmelCase     ,					block_out_channels[0]     ,					kernel_size=3     ,					padding=(1, 1)				)

          # time
          if time_embedding_type == "fourier":
               _a     						=					GaussianFourierProjection(embedding_size=block_out_channels[0]     ,					scale=16				)
               _a     						=					2 * block_out_channels[0]
          elif time_embedding_type == "positional":
               _a     						=					Timesteps(block_out_channels[0]     ,					__UpperCAmelCase     ,					__UpperCAmelCase				)
               _a     						=					block_out_channels[0]

          _a     						=					TimestepEmbedding(__UpperCAmelCase     ,					__UpperCAmelCase				)

          # class embedding
          if class_embed_type is None and num_class_embeds is not None:
               _a     						=					nn.Embedding(__UpperCAmelCase     ,					__UpperCAmelCase				)
          elif class_embed_type == "timestep":
               _a     						=					TimestepEmbedding(__UpperCAmelCase     ,					__UpperCAmelCase				)
          elif class_embed_type == "identity":
               _a     						=					nn.Identity(__UpperCAmelCase     ,					__UpperCAmelCase				)
          else:
               _a     						=					None

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

          # down
          _a     						=					block_out_channels[0]
          for i, down_block_type in enumerate(__UpperCAmelCase				):
               _a     						=					output_channel
               _a     						=					block_out_channels[i]
               _a     						=					i == len(__UpperCAmelCase				) - 1

               _a     						=					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
          _a     						=					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
          _a     						=					list(reversed(__UpperCAmelCase				)				)
          _a     						=					reversed_block_out_channels[0]
          for i, up_block_type in enumerate(__UpperCAmelCase				):
               _a     						=					output_channel
               _a     						=					reversed_block_out_channels[i]
               _a     						=					reversed_block_out_channels[min(i + 1     ,					len(__UpperCAmelCase				) - 1				)]

               _a     						=					i == len(__UpperCAmelCase				) - 1

               _a     						=					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				)
               _a     						=					output_channel

          # out
          _a     						=					norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4     ,					32				)
          _a     						=					nn.GroupNorm(num_channels=block_out_channels[0]     ,					num_groups=__UpperCAmelCase     ,					eps=__UpperCAmelCase				)
          _a     						=					nn.SiLU()
          _a     						=					nn.Convad(block_out_channels[0]     ,					__UpperCAmelCase     ,					kernel_size=3     ,					padding=1				)





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

          # 1. time
          _a     						=					timestep
          if not torch.is_tensor(__UpperCAmelCase				):
               _a     						=					torch.tensor([timesteps]     ,					dtype=torch.long     ,					device=sample.device				)
          elif torch.is_tensor(__UpperCAmelCase				) and len(timesteps.shape				) == 0:
               _a     						=					timesteps[None].to(sample.device				)

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

          _a     						=					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.
          _a     						=					t_emb.to(dtype=self.dtype				)
          _a     						=					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":
                    _a     						=					self.time_proj(__UpperCAmelCase				)

               _a     						=					self.class_embedding(__UpperCAmelCase				).to(dtype=self.dtype				)
               _a     						=					emb + class_emb

          # 2. pre-process
          _a     						=					sample
          _a     						=					self.conv_in(__UpperCAmelCase				)

          # 3. down
          _a     						=					(sample,)
          for downsample_block in self.down_blocks:
               if hasattr(__UpperCAmelCase     ,					'''skip_conv'''				):
                    _a						,   _a						,   _a     						=					downsample_block(
                        hidden_states=__UpperCAmelCase     ,					temb=__UpperCAmelCase     ,					skip_sample=__UpperCAmelCase				)
               else:
                    _a						,   _a     						=					downsample_block(hidden_states=__UpperCAmelCase     ,					temb=__UpperCAmelCase				)

               down_block_res_samples += res_samples

          # 4. mid
          _a     						=					self.mid_block(__UpperCAmelCase     ,					__UpperCAmelCase				)

          # 5. up
          _a     						=					None
          for upsample_block in self.up_blocks:
               _a     						=					down_block_res_samples[-len(upsample_block.resnets				) :]
               _a     						=					down_block_res_samples[: -len(upsample_block.resnets				)]

               if hasattr(__UpperCAmelCase     ,					'''skip_conv'''				):
                    _a						,   _a     						=					upsample_block(__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase				)
               else:
                    _a     						=					upsample_block(__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase				)

        # 6. post-process
          _a     						=					self.conv_norm_out(__UpperCAmelCase				)
          _a     						=					self.conv_act(__UpperCAmelCase				)
          _a     						=					self.conv_out(__UpperCAmelCase				)

          if skip_sample is not None:
               sample += skip_sample

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

          if not return_dict:
               return (sample,)

          return UNetaDOutput(sample=__UpperCAmelCase				)



"""simple docstring"""






import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionDualGuidedPipeline
from diffusers.utils.testing_utils import load_image, nightly, require_torch_gpu, torch_device


__snake_case												=							False


@nightly
@require_torch_gpu
class      __lowerCamelCase	(       unittest.TestCase					):

     '''simple docstring'''


     def 				_UpperCAmelCase   (   self				)							->				Tuple:
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()


     def 				_UpperCAmelCase   (   self				)							->				List[str]:
          _a     						=					VersatileDiffusionDualGuidedPipeline.from_pretrained('''shi-labs/versatile-diffusion'''				)
          # remove text_unet
          pipe.remove_unused_weights()
          pipe.to(__UpperCAmelCase				)
          pipe.set_progress_bar_config(disable=__UpperCAmelCase				)

          _a     						=					load_image(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'''				)

          _a     						=					torch.manual_seed(0				)
          _a     						=					pipe(
              prompt='''first prompt'''     ,					image=__UpperCAmelCase     ,					text_to_image_strength=0.75     ,					generator=__UpperCAmelCase     ,					guidance_scale=7.5     ,					num_inference_steps=2     ,					output_type='''numpy'''     ,					).images

          with tempfile.TemporaryDirectory() as tmpdirname:
               pipe.save_pretrained(__UpperCAmelCase				)
               _a     						=					VersatileDiffusionDualGuidedPipeline.from_pretrained(__UpperCAmelCase				)

          pipe.to(__UpperCAmelCase				)
          pipe.set_progress_bar_config(disable=__UpperCAmelCase				)

          _a     						=					generator.manual_seed(0				)
          _a     						=					pipe(
              prompt='''first prompt'''     ,					image=__UpperCAmelCase     ,					text_to_image_strength=0.75     ,					generator=__UpperCAmelCase     ,					guidance_scale=7.5     ,					num_inference_steps=2     ,					output_type='''numpy'''     ,					).images

          assert np.abs(image - new_image				).sum() < 1e-5, "Models don't have the same forward pass"





     def 				_UpperCAmelCase   (   self				)							->				Dict:
          _a     						=					VersatileDiffusionDualGuidedPipeline.from_pretrained('''shi-labs/versatile-diffusion'''				)
          pipe.remove_unused_weights()
          pipe.to(__UpperCAmelCase				)
          pipe.set_progress_bar_config(disable=__UpperCAmelCase				)

          _a     						=					'''cyberpunk 2077'''
          _a     						=					load_image(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'''				)
          _a     						=					torch.manual_seed(0				)
          _a     						=					pipe(
              prompt=__UpperCAmelCase     ,					image=__UpperCAmelCase     ,					text_to_image_strength=0.75     ,					generator=__UpperCAmelCase     ,					guidance_scale=7.5     ,					num_inference_steps=50     ,					output_type='''numpy'''     ,					).images

          _a     						=					image[0, 253:256, 253:256, -1]

          assert image.shape == (1, 512, 512, 3)
          _a     						=					np.array([0.0787, 0.0849, 0.0826, 0.0812, 0.0807, 0.0795, 0.0818, 0.0798, 0.0779]				)

          assert np.abs(image_slice.flatten() - expected_slice				).max() < 1e-2