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 			__magic_name__	(		__a   ,						unittest.TestCase     ):





   """simple docstring"""




   lowerCAmelCase   :     Optional[int]	      = KandinskyVaaControlnetPipeline
   lowerCAmelCase   :     int	      = ['''image_embeds''', '''negative_image_embeds''', '''hint''']
   lowerCAmelCase   :     List[str]	      = ['''image_embeds''', '''negative_image_embeds''', '''hint''']
   lowerCAmelCase   :     str	      = [
       '''generator''',
       '''height''',
       '''width''',
       '''latents''',
       '''guidance_scale''',
       '''num_inference_steps''',
       '''return_dict''',
       '''guidance_scale''',
       '''num_images_per_prompt''',
       '''output_type''',
       '''return_dict''',
   ]
   lowerCAmelCase   :     List[str]	      = False



   @property
   def 	lowerCAmelCase    (						self       :	int				):



    """simple docstring"""



    return 32



   @property
   def 	lowerCAmelCase    (						self       :	Any				):



    """simple docstring"""



    return 32



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



    """simple docstring"""



    return self.time_input_dim



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



    """simple docstring"""



    return self.time_input_dim * 4



   @property
   def 	lowerCAmelCase    (						self       :	str				):



    """simple docstring"""



    return 100



   @property
   def 	lowerCAmelCase    (						self       :	Tuple				):



    """simple docstring"""



    torch.manual_seed(0				)

    _UpperCamelCase:  int    						=			{
        '''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,
    }

    _UpperCamelCase:  str    						=			UNetaDConditionModel(**_lowercase				)
    return model



   @property
   def 	lowerCAmelCase    (						self       :	List[str]				):



    """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 	lowerCAmelCase    (						self       :	Any				):



    """simple docstring"""



    torch.manual_seed(0				)
    _UpperCamelCase:  Optional[int]    						=			VQModel(**self.dummy_movq_kwargs				)
    return model



   def 	lowerCAmelCase    (						self       :	Any				):



    """simple docstring"""



    _UpperCamelCase:  int    						=			self.dummy_unet
    _UpperCamelCase:  List[Any]    						=			self.dummy_movq

    _UpperCamelCase:  Dict    						=			DDIMScheduler(
        num_train_timesteps=1_000		,	beta_schedule='''linear'''		,	beta_start=0.00085		,	beta_end=0.012		,	clip_sample=_lowercase		,	set_alpha_to_one=_lowercase		,	steps_offset=1		,	prediction_type='''epsilon'''		,	thresholding=_lowercase		,	)

    _UpperCamelCase:  List[Any]    						=			{
        '''unet''': unet,
        '''scheduler''': scheduler,
        '''movq''': movq,
    }
    return components



   def 	lowerCAmelCase    (						self       :	Any		,	_lowercase       :	Any		,	_lowercase       :	Optional[Any]=0				):



    """simple docstring"""



    _UpperCamelCase:  int    						=			floats_tensor((1, self.text_embedder_hidden_size)		,	rng=random.Random(_lowercase				)				).to(_lowercase				)
    _UpperCamelCase:  Tuple    						=			floats_tensor((1, self.text_embedder_hidden_size)		,	rng=random.Random(seed + 1				)				).to(
        _lowercase				)

    # create hint
    _UpperCamelCase:  str    						=			floats_tensor((1, 3, 64, 64)		,	rng=random.Random(_lowercase				)				).to(_lowercase				)

    if str(_lowercase				).startswith('''mps'''				):
     _UpperCamelCase:  Union[str, Any]    						=			torch.manual_seed(_lowercase				)
    else:
     _UpperCamelCase:  str    						=			torch.Generator(device=_lowercase				).manual_seed(_lowercase				)
    _UpperCamelCase:  Optional[int]    						=			{
        '''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 	lowerCAmelCase    (						self       :	List[Any]				):



    """simple docstring"""



    _UpperCamelCase:  Tuple    						=			'''cpu'''

    _UpperCamelCase:  int    						=			self.get_dummy_components()

    _UpperCamelCase:  Dict    						=			self.pipeline_class(**_lowercase				)
    _UpperCamelCase:  Dict    						=			pipe.to(_lowercase				)

    pipe.set_progress_bar_config(disable=_lowercase				)

    _UpperCamelCase:  Union[str, Any]    						=			pipe(**self.get_dummy_inputs(_lowercase				)				)
    _UpperCamelCase:  Dict    						=			output.images

    _UpperCamelCase:  Optional[Any]    						=			pipe(
        **self.get_dummy_inputs(_lowercase				)		,	return_dict=_lowercase		,	)[0]

    _UpperCamelCase:  int    						=			image[0, -3:, -3:, -1]
    _UpperCamelCase:  int    						=			image_from_tuple[0, -3:, -3:, -1]

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

    _UpperCamelCase:  List[Any]    						=			np.array(
        [0.6959826, 0.868279, 0.7558092, 0.68769467, 0.85805804, 0.65977496, 0.44885302, 0.5959111, 0.4251595]				)

    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 			__magic_name__	(		unittest.TestCase     ):





   """simple docstring"""




   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



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



   def 	lowerCAmelCase    (						self       :	Optional[Any]				):



    """simple docstring"""



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

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

    _UpperCamelCase:  str    						=			KandinskyVaaPriorPipeline.from_pretrained(
        '''kandinsky-community/kandinsky-2-2-prior'''		,	torch_dtype=torch.floataa				)
    pipe_prior.to(_lowercase				)

    _UpperCamelCase:  List[str]    						=			KandinskyVaaControlnetPipeline.from_pretrained(
        '''kandinsky-community/kandinsky-2-2-controlnet-depth'''		,	torch_dtype=torch.floataa				)
    _UpperCamelCase:  Dict    						=			pipeline.to(_lowercase				)
    pipeline.set_progress_bar_config(disable=_lowercase				)

    _UpperCamelCase:  Optional[int]    						=			'''A robot, 4k photo'''

    _UpperCamelCase:  str    						=			torch.Generator(device='''cuda'''				).manual_seed(0				)
    _UpperCamelCase						,				_UpperCamelCase:  Dict    						=			pipe_prior(
        _lowercase		,	generator=_lowercase		,	num_inference_steps=5		,	negative_prompt=''''''		,	).to_tuple()

    _UpperCamelCase:  Any    						=			torch.Generator(device='''cuda'''				).manual_seed(0				)
    _UpperCamelCase:  Union[str, Any]    						=			pipeline(
        image_embeds=_lowercase		,	negative_image_embeds=_lowercase		,	hint=_lowercase		,	generator=_lowercase		,	num_inference_steps=100		,	output_type='''np'''		,	)

    _UpperCamelCase:  Optional[Any]    						=			output.images[0]

    assert image.shape == (512, 512, 3)

    assert_mean_pixel_difference(_lowercase		,	_lowercase				)
import unittest

import numpy as np

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
  from transformers.models.bert.modeling_flax_bert import (
      FlaxBertForMaskedLM,
      FlaxBertForMultipleChoice,
      FlaxBertForNextSentencePrediction,
      FlaxBertForPreTraining,
      FlaxBertForQuestionAnswering,
      FlaxBertForSequenceClassification,
      FlaxBertForTokenClassification,
      FlaxBertModel,
  )




class 			__magic_name__	(		unittest.TestCase     ):





   """simple docstring"""




   def __init__(						self       :	Tuple		,	_lowercase       :	List[str]		,	_lowercase       :	Dict=13		,	_lowercase       :	Any=7		,	_lowercase       :	List[str]=True		,	_lowercase       :	Any=True		,	_lowercase       :	Optional[int]=True		,	_lowercase       :	int=True		,	_lowercase       :	str=99		,	_lowercase       :	Union[str, Any]=32		,	_lowercase       :	Any=5		,	_lowercase       :	Dict=4		,	_lowercase       :	Tuple=37		,	_lowercase       :	Optional[Any]="gelu"		,	_lowercase       :	Optional[int]=0.1		,	_lowercase       :	Any=0.1		,	_lowercase       :	Union[str, Any]=512		,	_lowercase       :	Dict=16		,	_lowercase       :	int=2		,	_lowercase       :	Union[str, Any]=0.02		,	_lowercase       :	Dict=4		,	):



    """simple docstring"""



    _UpperCamelCase:  List[str]    						=			parent
    _UpperCamelCase:  Optional[int]    						=			batch_size
    _UpperCamelCase:  Any    						=			seq_length
    _UpperCamelCase:  List[str]    						=			is_training
    _UpperCamelCase:  Union[str, Any]    						=			use_attention_mask
    _UpperCamelCase:  Union[str, Any]    						=			use_token_type_ids
    _UpperCamelCase:  str    						=			use_labels
    _UpperCamelCase:  List[str]    						=			vocab_size
    _UpperCamelCase:  Dict    						=			hidden_size
    _UpperCamelCase:  List[Any]    						=			num_hidden_layers
    _UpperCamelCase:  List[Any]    						=			num_attention_heads
    _UpperCamelCase:  Tuple    						=			intermediate_size
    _UpperCamelCase:  Optional[Any]    						=			hidden_act
    _UpperCamelCase:  Any    						=			hidden_dropout_prob
    _UpperCamelCase:  int    						=			attention_probs_dropout_prob
    _UpperCamelCase:  Optional[int]    						=			max_position_embeddings
    _UpperCamelCase:  List[Any]    						=			type_vocab_size
    _UpperCamelCase:  List[Any]    						=			type_sequence_label_size
    _UpperCamelCase:  List[Any]    						=			initializer_range
    _UpperCamelCase:  Optional[int]    						=			num_choices



   def 	lowerCAmelCase    (						self       :	Optional[Any]				):



    """simple docstring"""



    _UpperCamelCase:  int    						=			ids_tensor([self.batch_size, self.seq_length]		,	self.vocab_size				)

    _UpperCamelCase:  Optional[int]    						=			None
    if self.use_attention_mask:
     _UpperCamelCase:  List[str]    						=			random_attention_mask([self.batch_size, self.seq_length]				)

    _UpperCamelCase:  Optional[int]    						=			None
    if self.use_token_type_ids:
     _UpperCamelCase:  Tuple    						=			ids_tensor([self.batch_size, self.seq_length]		,	self.type_vocab_size				)

    _UpperCamelCase:  str    						=			BertConfig(
        vocab_size=self.vocab_size		,	hidden_size=self.hidden_size		,	num_hidden_layers=self.num_hidden_layers		,	num_attention_heads=self.num_attention_heads		,	intermediate_size=self.intermediate_size		,	hidden_act=self.hidden_act		,	hidden_dropout_prob=self.hidden_dropout_prob		,	attention_probs_dropout_prob=self.attention_probs_dropout_prob		,	max_position_embeddings=self.max_position_embeddings		,	type_vocab_size=self.type_vocab_size		,	is_decoder=_lowercase		,	initializer_range=self.initializer_range		,	)

    return config, input_ids, token_type_ids, attention_mask



   def 	lowerCAmelCase    (						self       :	Optional[Any]				):



    """simple docstring"""



    _UpperCamelCase:  Optional[Any]    						=			self.prepare_config_and_inputs()
    _UpperCamelCase						,				_UpperCamelCase						,				_UpperCamelCase						,				_UpperCamelCase:  Any    						=			config_and_inputs
    _UpperCamelCase:  List[Any]    						=			{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
    return config, inputs_dict



   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    _UpperCamelCase:  Any    						=			self.prepare_config_and_inputs()
    _UpperCamelCase						,				_UpperCamelCase						,				_UpperCamelCase						,				_UpperCamelCase:  str    						=			config_and_inputs

    _UpperCamelCase:  int    						=			True
    _UpperCamelCase:  Optional[Any]    						=			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]				)
    _UpperCamelCase:  Optional[Any]    						=			ids_tensor([self.batch_size, self.seq_length]		,	vocab_size=2				)

    return (
        config,
        input_ids,
        attention_mask,
        encoder_hidden_states,
        encoder_attention_mask,
    )




@require_flax
class 			__magic_name__	(		__a   ,						unittest.TestCase     ):





   """simple docstring"""




   lowerCAmelCase   :     Dict	      = True

   lowerCAmelCase   :     List[Any]	      = (
       (
           FlaxBertModel,
           FlaxBertForPreTraining,
           FlaxBertForMaskedLM,
           FlaxBertForMultipleChoice,
           FlaxBertForQuestionAnswering,
           FlaxBertForNextSentencePrediction,
           FlaxBertForSequenceClassification,
           FlaxBertForTokenClassification,
           FlaxBertForQuestionAnswering,
       )
       if is_flax_available()
       else ()
   )



   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    _UpperCamelCase:  Dict    						=			FlaxBertModelTester(self				)



   @slow
   def 	lowerCAmelCase    (						self       :	Tuple				):



    """simple docstring"""



    _UpperCamelCase:  Dict    						=			FlaxBertModel.from_pretrained('''bert-base-cased'''				)
    _UpperCamelCase:  int    						=			model(np.ones((1, 1)				)				)
    self.assertIsNotNone(_lowercase				)