


from typing import List, Optional, Union

import numpy as np
import tensorflow as tf

from .utils import logging


_SCREAMING_SNAKE_CASE      				=   logging.get_logger(__name__)
def       UpperCAmelCase__						(UpperCamelCase_    ):



     """simple docstring"""






     if isinstance(_lowercase							,np.ndarray    ):
          return list(tensor.shape    )

     snake_case					=    tf.shape(_lowercase    )

     if tensor.shape == tf.TensorShape(_lowercase    ):
          return dynamic

     snake_case					=    tensor.shape.as_list()

     return [dynamic[i] if s is None else s for i, s in enumerate(_lowercase    )]
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_ = None							,UpperCamelCase_ = None    ):



     """simple docstring"""






     return tf.nn.softmax(logits=logits + 1e-9							,axis=_lowercase							,name=_lowercase    )
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_=1e-5							,UpperCamelCase_=-1    ):



     """simple docstring"""






     if weight.shape.rank != 1 or bias.shape.rank != 1 or not isinstance(_lowercase							,_lowercase    ):
          raise NotImplementedError('''Only 1D weight and bias tensors are supported for now, with only a single axis.'''    )

     # Get mean and variance on the axis to be normalized
     snake_case					=    tf.nn.moments(_lowercase							,axes=[axis]							,keepdims=_lowercase    )

     if axis != -1:
          # Reshape scale and weight to have the same rank as inputs, but with 1 dimensions
          # on every dimension except axis
          snake_case					=    [1] * inputs.shape.rank
          snake_case					=    shape_list(_lowercase    )[axis]
          snake_case					=    tf.reshape(_lowercase							,_lowercase    )
          snake_case					=    tf.reshape(_lowercase							,_lowercase    )

     # Compute layer normalization using the batch_normalization
     # function.
     snake_case					=    tf.nn.batch_normalization(
         _lowercase							,_lowercase							,_lowercase							,offset=_lowercase							,scale=_lowercase							,variance_epsilon=_lowercase							,)
     return outputs
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_=0							,UpperCamelCase_=-1    ):



     """simple docstring"""






     if end_dim < 0:
          end_dim += input.shape.rank
     if start_dim < 0:
          start_dim += input.shape.rank

     if start_dim == end_dim:
          return input

     snake_case					=    tf.shape(_lowercase    )
     snake_case					=    tf.math.reduce_prod(in_shape[start_dim : end_dim + 1]    )
     snake_case					=    tf.concat([in_shape[:start_dim], [flattened_dim], in_shape[end_dim + 1 :]]							,axis=0    )
     return tf.reshape(_lowercase							,_lowercase    )
def       UpperCAmelCase__						(UpperCamelCase_    ):



     """simple docstring"""






     if not isinstance(_lowercase							,tf.Tensor    ):
          snake_case					=    tf.convert_to_tensor(_lowercase    )  # Catches stray NumPy inputs
     if encoder_attention_mask.shape.rank == 3:
          snake_case					=    encoder_attention_mask[:, None, :, :]
     if encoder_attention_mask.shape.rank == 2:
          snake_case					=    encoder_attention_mask[:, None, None, :]
     # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
     # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow
     # /transformer/transformer_layers.py#L270
     # encoder_extended_attention_mask = (encoder_extended_attention_mask ==
     # encoder_extended_attention_mask.transpose(-1, -2))
     snake_case					=    (
         tf.cast(1							,encoder_attention_mask.dtype    ) - encoder_extended_attention_mask
     ) * encoder_extended_attention_mask.dtype.min

     return encoder_extended_attention_mask
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_ = "input_ids"    ):



     """simple docstring"""






     tf.debugging.assert_less(
         _lowercase							,tf.cast(_lowercase							,dtype=tensor.dtype    )							,message=(
             F'''The maximum value of {tensor_name} ({tf.math.reduce_max(_lowercase    )}) must be smaller than the embedding '''
             F'''layer\'s input dimension ({embed_dim}). The likely cause is some problem at tokenization time.'''
         )							,)
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_    ):



     """simple docstring"""






     snake_case					=    6_45_12
     # Check that no item in `data` is larger than `HDF5_OBJECT_HEADER_LIMIT`
     # because in that case even chunking the array would not make the saving
     # possible.
     snake_case					=    [x for x in data if len(_lowercase    ) > HDF5_OBJECT_HEADER_LIMIT]

     # Expecting this to never be true.
     if bad_attributes:
          raise RuntimeError(
              '''The following attributes cannot be saved to HDF5 file because '''
              F'''they are larger than {HDF5_OBJECT_HEADER_LIMIT} '''
              F'''bytes: {bad_attributes}'''    )

     snake_case					=    np.asarray(_lowercase    )

     snake_case					=    1
     snake_case					=    np.array_split(_lowercase							,_lowercase    )

     # This will never loop forever thanks to the test above.
     while any(x.nbytes > HDF5_OBJECT_HEADER_LIMIT for x in chunked_data    ):
          num_chunks += 1
          snake_case					=    np.array_split(_lowercase							,_lowercase    )

     if num_chunks > 1:
          for chunk_id, chunk_data in enumerate(_lowercase    ):
               snake_case					=    chunk_data
     else:
          snake_case					=    data
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_    ):



     """simple docstring"""






     if name in group.attrs:
          snake_case					=    [n.decode('''utf8'''    ) if hasattr(_lowercase							,'''decode'''    ) else n for n in group.attrs[name]]
     else:
          snake_case					=    []
          snake_case					=    0
          while "%s%d" % (name, chunk_id) in group.attrs:
               data.extend(
                   [n.decode('''utf8'''    ) if hasattr(_lowercase							,'''decode'''    ) else n for n in group.attrs['''%s%d''' % (name, chunk_id)]]    )
               chunk_id += 1
     return data

def       UpperCAmelCase__						(UpperCamelCase_    ):



     """simple docstring"""






     def _expand_single_ad_tensor(UpperCamelCase_    ):
          if isinstance(_lowercase							,tf.Tensor    ) and t.shape.rank == 1:
               return tf.expand_dims(_lowercase							,axis=-1    )
          return t

     return tf.nest.map_structure(_expand_single_ad_tensor							,_lowercase    )








import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyImgaImgPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
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   A__ (   snake_case__	, unittest.TestCase     ):







    """simple docstring"""

    __magic_name__        =		KandinskyImgaImgPipeline
    __magic_name__        =		['prompt', 'image_embeds', 'negative_image_embeds', 'image']
    __magic_name__        =		[
        'prompt',
        'negative_prompt',
        'image_embeds',
        'negative_image_embeds',
        'image',
    ]
    __magic_name__        =		[
        'generator',
        'height',
        'width',
        'strength',
        'guidance_scale',
        'negative_prompt',
        'num_inference_steps',
        'return_dict',
        'guidance_scale',
        'num_images_per_prompt',
        'output_type',
        'return_dict',
    ]
    __magic_name__        =		False
    @property
    def 		a_      (	self   ):
         return 3_2
    @property
    def 		a_      (	self   ):
         return 3_2
    @property
    def 		a_      (	self   ):
         return self.time_input_dim
    @property
    def 		a_      (	self   ):
         return self.time_input_dim * 4
    @property
    def 		a_      (	self   ):
         return 1_0_0
    @property
    def 		a_      (	self   ):
         snake_case					=    XLMRobertaTokenizerFast.from_pretrained('''YiYiXu/tiny-random-mclip-base'''   )
         return tokenizer
    @property
    def 		a_      (	self   ):
         torch.manual_seed(0   )
         snake_case					=    MCLIPConfig(
             numDims=self.cross_attention_dim      , transformerDimensions=self.text_embedder_hidden_size      , hidden_size=self.text_embedder_hidden_size      , intermediate_size=3_7      , num_attention_heads=4      , num_hidden_layers=5      , vocab_size=1_0_0_5      , )

         snake_case					=    MultilingualCLIP(__snake_case   )
         snake_case					=    text_encoder.eval()

         return text_encoder
    @property
    def 		a_      (	self   ):
         torch.manual_seed(0   )

         snake_case					=    {
             '''in_channels''': 4,
             # Out channels is double in channels because predicts mean and variance
             '''out_channels''': 8,
             '''addition_embed_type''': '''text_image''',
             '''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''': '''text_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					=    UNetaDConditionModel(**__snake_case   )
         return model
    @property
    def 		a_      (	self   ):
         return {
             "block_out_channels": [3_2, 6_4],
             "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
             "in_channels": 3,
             "latent_channels": 4,
             "layers_per_block": 1,
             "norm_num_groups": 8,
             "norm_type": "spatial",
             "num_vq_embeddings": 1_2,
             "out_channels": 3,
             "up_block_types": [
                 "AttnUpDecoderBlock2D",
                 "UpDecoderBlock2D",
             ],
             "vq_embed_dim": 4,
         }
    @property
    def 		a_      (	self   ):
         torch.manual_seed(0   )
         snake_case					=    VQModel(**self.dummy_movq_kwargs   )
         return model
    def 		a_      (	self   ):
         snake_case					=    self.dummy_text_encoder
         snake_case					=    self.dummy_tokenizer
         snake_case					=    self.dummy_unet
         snake_case					=    self.dummy_movq

         snake_case					=    {
             '''num_train_timesteps''': 1_0_0_0,
             '''beta_schedule''': '''linear''',
             '''beta_start''': 0.0_0085,
             '''beta_end''': 0.012,
             '''clip_sample''': False,
             '''set_alpha_to_one''': False,
             '''steps_offset''': 0,
             '''prediction_type''': '''epsilon''',
             '''thresholding''': False,
         }

         snake_case					=    DDIMScheduler(**__snake_case   )

         snake_case					=    {
             '''text_encoder''': text_encoder,
             '''tokenizer''': tokenizer,
             '''unet''': unet,
             '''scheduler''': scheduler,
             '''movq''': movq,
         }

         return components
    def 		a_      (	self      , __snake_case      , __snake_case=0   ):
         snake_case					=    floats_tensor((1, self.cross_attention_dim)      , rng=random.Random(__snake_case   )   ).to(__snake_case   )
         snake_case					=    floats_tensor((1, self.cross_attention_dim)      , rng=random.Random(seed + 1   )   ).to(__snake_case   )
         # create init_image
         snake_case					=    floats_tensor((1, 3, 6_4, 6_4)      , rng=random.Random(__snake_case   )   ).to(__snake_case   )
         snake_case					=    image.cpu().permute(0      , 2      , 3      , 1   )[0]
         snake_case					=    Image.fromarray(np.uinta(__snake_case   )   ).convert('''RGB'''   ).resize((2_5_6, 2_5_6)   )

         if str(__snake_case   ).startswith('''mps'''   ):
              snake_case					=    torch.manual_seed(__snake_case   )
         else:
              snake_case					=    torch.Generator(device=__snake_case   ).manual_seed(__snake_case   )
         snake_case					=    {
             '''prompt''': '''horse''',
             '''image''': init_image,
             '''image_embeds''': image_embeds,
             '''negative_image_embeds''': negative_image_embeds,
             '''generator''': generator,
             '''height''': 6_4,
             '''width''': 6_4,
             '''num_inference_steps''': 1_0,
             '''guidance_scale''': 7.0,
             '''strength''': 0.2,
             '''output_type''': '''np''',
         }
         return inputs




    def 		a_      (	self   ):
         snake_case					=    '''cpu'''

         snake_case					=    self.get_dummy_components()

         snake_case					=    self.pipeline_class(**__snake_case   )
         snake_case					=    pipe.to(__snake_case   )

         pipe.set_progress_bar_config(disable=__snake_case   )

         snake_case					=    pipe(**self.get_dummy_inputs(__snake_case   )   )
         snake_case					=    output.images

         snake_case					=    pipe(
             **self.get_dummy_inputs(__snake_case   )      , return_dict=__snake_case      , )[0]

         snake_case					=    image[0, -3:, -3:, -1]
         snake_case					=    image_from_tuple[0, -3:, -3:, -1]

         assert image.shape == (1, 6_4, 6_4, 3)

         snake_case					=    np.array(
             [0.6147_4943, 0.607_3539, 0.4330_8544, 0.592_8269, 0.4749_3595, 0.4675_5973, 0.461_3838, 0.4536_8797, 0.5011_9233]   )
         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   A__ (   unittest.TestCase     ):







    """simple docstring"""

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




    def 		a_      (	self   ):
         snake_case					=    load_numpy(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/kandinsky/kandinsky_img2img_frog.npy'''   )

         snake_case					=    load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/kandinsky/cat.png'''   )
         snake_case					=    '''A red cartoon frog, 4k'''

         snake_case					=    KandinskyPriorPipeline.from_pretrained(
             '''kandinsky-community/kandinsky-2-1-prior'''      , torch_dtype=torch.floataa   )
         pipe_prior.to(__snake_case   )

         snake_case					=    KandinskyImgaImgPipeline.from_pretrained(
             '''kandinsky-community/kandinsky-2-1'''      , torch_dtype=torch.floataa   )
         snake_case					=    pipeline.to(__snake_case   )

         pipeline.set_progress_bar_config(disable=__snake_case   )

         snake_case					=    torch.Generator(device='''cpu'''   ).manual_seed(0   )
         snake_case			,  snake_case					=    pipe_prior(
             __snake_case      , generator=__snake_case      , num_inference_steps=5      , negative_prompt=''''''      , ).to_tuple()

         snake_case					=    pipeline(
             __snake_case      , image=__snake_case      , image_embeds=__snake_case      , negative_image_embeds=__snake_case      , generator=__snake_case      , num_inference_steps=1_0_0      , height=7_6_8      , width=7_6_8      , strength=0.2      , output_type='''np'''      , )

         snake_case					=    output.images[0]

         assert image.shape == (7_6_8, 7_6_8, 3)

         assert_mean_pixel_difference(__snake_case      , __snake_case   )




