import gc
import random
import unittest

import numpy as np
import torch

from diffusers import (
    DDIMScheduler,
    KandinskyVaaControlnetPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()



class        _lowerCAmelCase    (			a__							,    unittest.TestCase							):


  '''simple docstring'''
  a_							:       Union[str, Any]      					=KandinskyVaaControlnetPipeline
  a_							:       str      					=["""image_embeds""", """negative_image_embeds""", """hint"""]
  a_							:       Tuple      					=["""image_embeds""", """negative_image_embeds""", """hint"""]
  a_							:       List[str]      					=[
      """generator""",
      """height""",
      """width""",
      """latents""",
      """guidance_scale""",
      """num_inference_steps""",
      """return_dict""",
      """guidance_scale""",
      """num_images_per_prompt""",
      """output_type""",
      """return_dict""",
  ]
  a_							:       List[str]      					=False







  @property
  def UpperCamelCase_ (      self      :					Optional[int]							):


    '''simple docstring'''





    return 32







  @property
  def UpperCamelCase_ (      self      :					Optional[Any]							):


    '''simple docstring'''





    return 32







  @property
  def UpperCamelCase_ (      self      :					List[Any]							):


    '''simple docstring'''





    return self.time_input_dim







  @property
  def UpperCamelCase_ (      self      :					Dict							):


    '''simple docstring'''





    return self.time_input_dim * 4







  @property
  def UpperCamelCase_ (      self      :					List[Any]							):


    '''simple docstring'''





    return 1_00







  @property
  def UpperCamelCase_ (      self      :					Any							):


    '''simple docstring'''





    torch.manual_seed(0							)

    _snake_case   :				Union[str, Any]														=	{
        'in_channels': 8,
        # Out channels is double in channels because predicts mean and variance
        'out_channels': 8,
        'addition_embed_type': 'image_hint',
        'down_block_types': ('ResnetDownsampleBlock2D', 'SimpleCrossAttnDownBlock2D'),
        'up_block_types': ('SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D'),
        'mid_block_type': 'UNetMidBlock2DSimpleCrossAttn',
        'block_out_channels': (self.block_out_channels_a, self.block_out_channels_a * 2),
        'layers_per_block': 1,
        'encoder_hid_dim': self.text_embedder_hidden_size,
        'encoder_hid_dim_type': 'image_proj',
        'cross_attention_dim': self.cross_attention_dim,
        'attention_head_dim': 4,
        'resnet_time_scale_shift': 'scale_shift',
        'class_embed_type': None,
    }

    _snake_case   :				int														=	UNetaDConditionModel(**_A							)
    return model







  @property
  def UpperCamelCase_ (      self      :					Optional[int]							):


    '''simple docstring'''





    return {
        "block_out_channels": [32, 32, 64, 64],
        "down_block_types": [
            "DownEncoderBlock2D",
            "DownEncoderBlock2D",
            "DownEncoderBlock2D",
            "AttnDownEncoderBlock2D",
        ],
        "in_channels": 3,
        "latent_channels": 4,
        "layers_per_block": 1,
        "norm_num_groups": 8,
        "norm_type": "spatial",
        "num_vq_embeddings": 12,
        "out_channels": 3,
        "up_block_types": ["AttnUpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"],
        "vq_embed_dim": 4,
    }







  @property
  def UpperCamelCase_ (      self      :					Tuple							):


    '''simple docstring'''





    torch.manual_seed(0							)
    _snake_case   :				List[Any]														=	VQModel(**self.dummy_movq_kwargs							)
    return model







  def UpperCamelCase_ (      self      :					Union[str, Any]							):


    '''simple docstring'''





    _snake_case   :				str														=	self.dummy_unet
    _snake_case   :				int														=	self.dummy_movq

    _snake_case   :				str														=	DDIMScheduler(
        num_train_timesteps=10_00      ,				beta_schedule='linear'      ,				beta_start=0.0_00_85      ,				beta_end=0.0_12      ,				clip_sample=_A      ,				set_alpha_to_one=_A      ,				steps_offset=1      ,				prediction_type='epsilon'      ,				thresholding=_A      ,				)

    _snake_case   :				Any														=	{
        'unet': unet,
        'scheduler': scheduler,
        'movq': movq,
    }
    return components







  def UpperCamelCase_ (      self      :					Any      ,				UpperCamelCase      :					str      ,				UpperCamelCase      :					Dict=0							):


    '''simple docstring'''





    _snake_case   :				List[str]														=	floats_tensor((1, self.text_embedder_hidden_size)      ,				rng=random.Random(_A							)							).to(_A							)
    _snake_case   :				Any														=	floats_tensor((1, self.text_embedder_hidden_size)      ,				rng=random.Random(seed + 1							)							).to(
        _A							)

    # create hint
    _snake_case   :				Union[str, Any]														=	floats_tensor((1, 3, 64, 64)      ,				rng=random.Random(_A							)							).to(_A							)

    if str(_A							).startswith('mps'							):
      _snake_case   :				List[Any]														=	torch.manual_seed(_A							)
    else:
      _snake_case   :				str														=	torch.Generator(device=_A							).manual_seed(_A							)
    _snake_case   :				str														=	{
        'image_embeds': image_embeds,
        'negative_image_embeds': negative_image_embeds,
        'hint': hint,
        'generator': generator,
        'height': 64,
        'width': 64,
        'guidance_scale': 4.0,
        'num_inference_steps': 2,
        'output_type': 'np',
    }
    return inputs







  def UpperCamelCase_ (      self      :					Optional[Any]							):


    '''simple docstring'''





    _snake_case   :				List[Any]														=	'cpu'

    _snake_case   :				Tuple														=	self.get_dummy_components()

    _snake_case   :				List[Any]														=	self.pipeline_class(**_A							)
    _snake_case   :				Union[str, Any]														=	pipe.to(_A							)

    pipe.set_progress_bar_config(disable=_A							)

    _snake_case   :				int														=	pipe(**self.get_dummy_inputs(_A							)							)
    _snake_case   :				Any														=	output.images

    _snake_case   :				List[str]														=	pipe(
        **self.get_dummy_inputs(_A							)      ,				return_dict=_A      ,				)[0]

    _snake_case   :				List[Any]														=	image[0, -3:, -3:, -1]
    _snake_case   :				List[Any]														=	image_from_tuple[0, -3:, -3:, -1]

    assert image.shape == (1, 64, 64, 3)

    _snake_case   :				int														=	np.array(
        [0.6_95_98_26, 0.86_82_79, 0.7_55_80_92, 0.68_76_94_67, 0.85_80_58_04, 0.65_97_74_96, 0.44_88_53_02, 0.5_95_91_11, 0.4_25_15_95]							)

    assert (
        np.abs(image_slice.flatten() - expected_slice							).max() < 1e-2
    ), f""" expected_slice {expected_slice}, but got {image_slice.flatten()}"""

    assert (
        np.abs(image_from_tuple_slice.flatten() - expected_slice							).max() < 1e-2
    ), f""" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"""




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


  '''simple docstring'''







  def UpperCamelCase_ (      self      :					Dict							):


    '''simple docstring'''





    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()







  def UpperCamelCase_ (      self      :					Optional[Any]							):


    '''simple docstring'''





    _snake_case   :				Optional[Any]														=	load_numpy(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
        '/kandinskyv22/kandinskyv22_controlnet_robotcat_fp16.npy'							)

    _snake_case   :				List[str]														=	load_image(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
        '/kandinskyv22/hint_image_cat.png'							)
    _snake_case   :				Optional[int]														=	torch.from_numpy(np.array(_A							)							).float() / 2_55.0
    _snake_case   :				Tuple														=	hint.permute(2      ,				0      ,				1							).unsqueeze(0							)

    _snake_case   :				List[str]														=	KandinskyVaaPriorPipeline.from_pretrained(
        'kandinsky-community/kandinsky-2-2-prior'      ,				torch_dtype=torch.floataa							)
    pipe_prior.to(_A							)

    _snake_case   :				Any														=	KandinskyVaaControlnetPipeline.from_pretrained(
        'kandinsky-community/kandinsky-2-2-controlnet-depth'      ,				torch_dtype=torch.floataa							)
    _snake_case   :				int														=	pipeline.to(_A							)
    pipeline.set_progress_bar_config(disable=_A							)

    _snake_case   :				List[str]														=	'A robot, 4k photo'

    _snake_case   :				Optional[Any]														=	torch.Generator(device='cuda'							).manual_seed(0							)
    _snake_case   :				Tuple														=	pipe_prior(
        _A      ,				generator=_A      ,				num_inference_steps=5      ,				negative_prompt=''      ,				).to_tuple()

    _snake_case   :				Any														=	torch.Generator(device='cuda'							).manual_seed(0							)
    _snake_case   :				List[Any]														=	pipeline(
        image_embeds=_A      ,				negative_image_embeds=_A      ,				hint=_A      ,				generator=_A      ,				num_inference_steps=1_00      ,				output_type='np'      ,				)

    _snake_case   :				Dict														=	output.images[0]

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

    assert_mean_pixel_difference(_A      ,				_A							)











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

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .attention_processor import AttentionProcessor, AttnProcessor
from .embeddings import TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin



@dataclass
class        _lowerCAmelCase    (			UpperCAmelCase_							):


  '''simple docstring'''
  a_							:       torch.FloatTensor




class        _lowerCAmelCase    (			UpperCAmelCase_							,    UpperCAmelCase_							):


  '''simple docstring'''







  @register_to_config
  def __init__(      self      :					str      ,				UpperCamelCase      :					int = 32      ,				UpperCamelCase      :					int = 64      ,				UpperCamelCase      :					int = 20      ,				UpperCamelCase      :					int = 7_68      ,				UpperCamelCase      :					Optional[int]=77      ,				UpperCamelCase      :					int=4      ,				UpperCamelCase      :					float = 0.0      ,				UpperCamelCase      :					str = "silu"      ,				UpperCamelCase      :					Optional[str] = None      ,				UpperCamelCase      :					Optional[str] = None      ,				UpperCamelCase      :					Optional[str] = "linear"      ,				UpperCamelCase      :					Optional[str] = "prd"      ,				UpperCamelCase      :					Optional[int] = None      ,				UpperCamelCase      :					Optional[int] = None      ,				UpperCamelCase      :					Optional[int] = None      ,				):


    '''simple docstring'''





    super().__init__()
    _snake_case   :				str														=	num_attention_heads
    _snake_case   :				Optional[int]														=	attention_head_dim
    _snake_case   :				Any														=	num_attention_heads * attention_head_dim
    _snake_case   :				List[Any]														=	additional_embeddings

    _snake_case   :				List[str]														=	time_embed_dim or inner_dim
    _snake_case   :				int														=	embedding_proj_dim or embedding_dim
    _snake_case   :				List[Any]														=	clip_embed_dim or embedding_dim

    _snake_case   :				Optional[Any]														=	Timesteps(UpperCamelCase      ,				UpperCamelCase      ,				0							)
    _snake_case   :				List[Any]														=	TimestepEmbedding(UpperCamelCase      ,				UpperCamelCase      ,				out_dim=UpperCamelCase      ,				act_fn=UpperCamelCase							)

    _snake_case   :				Optional[int]														=	nn.Linear(UpperCamelCase      ,				UpperCamelCase							)

    if embedding_proj_norm_type is None:
      _snake_case   :				str														=	None
    elif embedding_proj_norm_type == "layer":
      _snake_case   :				List[Any]														=	nn.LayerNorm(UpperCamelCase							)
    else:
      raise ValueError(f"""unsupported embedding_proj_norm_type: {embedding_proj_norm_type}"""							)

    _snake_case   :				str														=	nn.Linear(UpperCamelCase      ,				UpperCamelCase							)

    if encoder_hid_proj_type is None:
      _snake_case   :				Any														=	None
    elif encoder_hid_proj_type == "linear":
      _snake_case   :				Optional[int]														=	nn.Linear(UpperCamelCase      ,				UpperCamelCase							)
    else:
      raise ValueError(f"""unsupported encoder_hid_proj_type: {encoder_hid_proj_type}"""							)

    _snake_case   :				List[str]														=	nn.Parameter(torch.zeros(1      ,				num_embeddings + additional_embeddings      ,				UpperCamelCase							)							)

    if added_emb_type == "prd":
      _snake_case   :				str														=	nn.Parameter(torch.zeros(1      ,				1      ,				UpperCamelCase							)							)
    elif added_emb_type is None:
      _snake_case   :				Dict														=	None
    else:
      raise ValueError(
          f"""`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `'prd'` or `None`."""							)

    _snake_case   :				Optional[int]														=	nn.ModuleList(
        [
            BasicTransformerBlock(
                UpperCamelCase      ,				UpperCamelCase      ,				UpperCamelCase      ,				dropout=UpperCamelCase      ,				activation_fn='gelu'      ,				attention_bias=UpperCamelCase      ,				)
            for d in range(UpperCamelCase							)
        ]							)

    if norm_in_type == "layer":
      _snake_case   :				Optional[int]														=	nn.LayerNorm(UpperCamelCase							)
    elif norm_in_type is None:
      _snake_case   :				Optional[Any]														=	None
    else:
      raise ValueError(f"""Unsupported norm_in_type: {norm_in_type}."""							)

    _snake_case   :				Optional[Any]														=	nn.LayerNorm(UpperCamelCase							)

    _snake_case   :				Union[str, Any]														=	nn.Linear(UpperCamelCase      ,				UpperCamelCase							)

    _snake_case   :				List[Any]														=	torch.full(
        [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]      ,				-1_00_00.0							)
    causal_attention_mask.triu_(1							)
    _snake_case   :				Optional[Any]														=	causal_attention_mask[None, ...]
    self.register_buffer('causal_attention_mask'      ,				UpperCamelCase      ,				persistent=UpperCamelCase							)

    _snake_case   :				str														=	nn.Parameter(torch.zeros(1      ,				UpperCamelCase							)							)
    _snake_case   :				List[str]														=	nn.Parameter(torch.zeros(1      ,				UpperCamelCase							)							)







  @property
  # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
  def UpperCamelCase_ (      self      :					Optional[Any]							):


    '''simple docstring'''





    _snake_case   :				Optional[Any]														=	{}

    def fn_recursive_add_processors(UpperCamelCase      :					str      ,				UpperCamelCase      :					torch.nn.Module      ,				UpperCamelCase      :					Dict[str, AttentionProcessor]							):
      if hasattr(UpperCamelCase      ,				'set_processor'							):
        _snake_case   :				Tuple														=	module.processor

      for sub_name, child in module.named_children():
        fn_recursive_add_processors(f"""{name}.{sub_name}"""      ,				UpperCamelCase      ,				UpperCamelCase							)

      return processors

    for name, module in self.named_children():
      fn_recursive_add_processors(UpperCamelCase      ,				UpperCamelCase      ,				UpperCamelCase							)

    return processors







  def UpperCamelCase_ (      self      :					List[Any]      ,				UpperCamelCase      :					Union[AttentionProcessor, Dict[str, AttentionProcessor]]							):


    '''simple docstring'''





    _snake_case   :				Optional[int]														=	len(self.attn_processors.keys()							)

    if isinstance(UpperCamelCase      ,				UpperCamelCase							) and len(UpperCamelCase							) != count:
      raise ValueError(
          f"""A dict of processors was passed, but the number of processors {len(UpperCamelCase							)} does not match the"""
          f""" number of attention layers: {count}. Please make sure to pass {count} processor classes."""							)

    def fn_recursive_attn_processor(UpperCamelCase      :					str      ,				UpperCamelCase      :					torch.nn.Module      ,				UpperCamelCase      :					Union[str, Any]							):
      if hasattr(UpperCamelCase      ,				'set_processor'							):
        if not isinstance(UpperCamelCase      ,				UpperCamelCase							):
          module.set_processor(UpperCamelCase							)
        else:
          module.set_processor(processor.pop(f"""{name}.processor"""							)							)

      for sub_name, child in module.named_children():
        fn_recursive_attn_processor(f"""{name}.{sub_name}"""      ,				UpperCamelCase      ,				UpperCamelCase							)

    for name, module in self.named_children():
      fn_recursive_attn_processor(UpperCamelCase      ,				UpperCamelCase      ,				UpperCamelCase							)







  def UpperCamelCase_ (      self      :					Optional[Any]							):


    '''simple docstring'''





    self.set_attn_processor(AttnProcessor()							)







  def UpperCamelCase_ (      self      :					Optional[Any]      ,				UpperCamelCase      :					Any      ,				UpperCamelCase      :					Union[torch.Tensor, float, int]      ,				UpperCamelCase      :					torch.FloatTensor      ,				UpperCamelCase      :					Optional[torch.FloatTensor] = None      ,				UpperCamelCase      :					Optional[torch.BoolTensor] = None      ,				UpperCamelCase      :					bool = True      ,				):


    '''simple docstring'''





    _snake_case   :				Dict														=	hidden_states.shape[0]

    _snake_case   :				str														=	timestep
    if not torch.is_tensor(UpperCamelCase							):
      _snake_case   :				Dict														=	torch.tensor([timesteps]      ,				dtype=torch.long      ,				device=hidden_states.device							)
    elif torch.is_tensor(UpperCamelCase							) and len(timesteps.shape							) == 0:
      _snake_case   :				Tuple														=	timesteps[None].to(hidden_states.device							)

    # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
    _snake_case   :				Optional[int]														=	timesteps * torch.ones(UpperCamelCase      ,				dtype=timesteps.dtype      ,				device=timesteps.device							)

    _snake_case   :				Union[str, Any]														=	self.time_proj(UpperCamelCase							)

    # timesteps does not contain any weights and will always return f32 tensors
    # but time_embedding might be fp16, so we need to cast here.
    _snake_case   :				Tuple														=	timesteps_projected.to(dtype=self.dtype							)
    _snake_case   :				List[Any]														=	self.time_embedding(UpperCamelCase							)

    if self.embedding_proj_norm is not None:
      _snake_case   :				Optional[Any]														=	self.embedding_proj_norm(UpperCamelCase							)

    _snake_case   :				Union[str, Any]														=	self.embedding_proj(UpperCamelCase							)
    if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
      _snake_case   :				Dict														=	self.encoder_hidden_states_proj(UpperCamelCase							)
    elif self.encoder_hidden_states_proj is not None and encoder_hidden_states is None:
      raise ValueError('`encoder_hidden_states_proj` requires `encoder_hidden_states` to be set'							)

    _snake_case   :				str														=	self.proj_in(UpperCamelCase							)

    _snake_case   :				int														=	self.positional_embedding.to(hidden_states.dtype							)

    _snake_case   :				Optional[int]														=	[]
    _snake_case   :				List[Any]														=	0

    if encoder_hidden_states is not None:
      additional_embeds.append(UpperCamelCase							)
      additional_embeddings_len += encoder_hidden_states.shape[1]

    if len(proj_embeddings.shape							) == 2:
      _snake_case   :				str														=	proj_embeddings[:, None, :]

    if len(hidden_states.shape							) == 2:
      _snake_case   :				str														=	hidden_states[:, None, :]

    _snake_case   :				str														=	additional_embeds + [
        proj_embeddings,
        time_embeddings[:, None, :],
        hidden_states,
    ]

    if self.prd_embedding is not None:
      _snake_case   :				int														=	self.prd_embedding.to(hidden_states.dtype							).expand(UpperCamelCase      ,				-1      ,				-1							)
      additional_embeds.append(UpperCamelCase							)

    _snake_case   :				Optional[int]														=	torch.cat(
        UpperCamelCase      ,				dim=1      ,				)

    # Allow positional_embedding to not include the `addtional_embeddings` and instead pad it with zeros for these additional tokens
    _snake_case   :				Tuple														=	additional_embeddings_len + proj_embeddings.shape[1] + 1
    if positional_embeddings.shape[1] < hidden_states.shape[1]:
      _snake_case   :				Optional[Any]														=	F.pad(
          UpperCamelCase      ,				(
              0,
              0,
              additional_embeddings_len,
              self.prd_embedding.shape[1] if self.prd_embedding is not None else 0,
          )      ,				value=0.0      ,				)

    _snake_case   :				Optional[Any]														=	hidden_states + positional_embeddings

    if attention_mask is not None:
      _snake_case   :				Any														=	(1 - attention_mask.to(hidden_states.dtype							)) * -1_00_00.0
      _snake_case   :				Tuple														=	F.pad(UpperCamelCase      ,				(0, self.additional_embeddings)      ,				value=0.0							)
      _snake_case   :				int														=	(attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype							)
      _snake_case   :				str														=	attention_mask.repeat_interleave(self.config.num_attention_heads      ,				dim=0							)

    if self.norm_in is not None:
      _snake_case   :				Tuple														=	self.norm_in(UpperCamelCase							)

    for block in self.transformer_blocks:
      _snake_case   :				Any														=	block(UpperCamelCase      ,				attention_mask=UpperCamelCase							)

    _snake_case   :				Dict														=	self.norm_out(UpperCamelCase							)

    if self.prd_embedding is not None:
      _snake_case   :				str														=	hidden_states[:, -1]
    else:
      _snake_case   :				Any														=	hidden_states[:, additional_embeddings_len:]

    _snake_case   :				List[Any]														=	self.proj_to_clip_embeddings(UpperCamelCase							)

    if not return_dict:
      return (predicted_image_embedding,)

    return PriorTransformerOutput(predicted_image_embedding=UpperCamelCase							)







  def UpperCamelCase_ (      self      :					Tuple      ,				UpperCamelCase      :					Union[str, Any]							):


    '''simple docstring'''





    _snake_case   :				List[Any]														=	(prior_latents * self.clip_std) + self.clip_mean
    return prior_latents




