




import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotConfig, is_flax_available
from transformers.testing_utils import jax_device, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
   import os

   # The slow tests are often failing with OOM error on GPU
   # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
   # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
   lowerCamelCase__					:  List[str]      							=				'platform'

   import jax
   import jax.numpy as jnp

   from transformers import BlenderbotTokenizer
   from transformers.models.blenderbot.modeling_flax_blenderbot import (
       FlaxBlenderbotForConditionalGeneration,
       FlaxBlenderbotModel,
       shift_tokens_right,
   )

def    UpperCAmelCase_				( __UpperCAmelCase		:      List[Any]     ,     __UpperCAmelCase		:      Optional[int]     ,     __UpperCAmelCase		:      str=None     ,     __UpperCAmelCase		:      str=None     ,     __UpperCAmelCase		:      str=None     ,     __UpperCAmelCase		:      Union[str, Any]=None     ,     __UpperCAmelCase		:      Optional[int]=None     ,     __UpperCAmelCase		:      Optional[Any]=None     ,     )   ->					Tuple:
      if attention_mask is None:
            SCREAMING_SNAKE_CASE_      		=     np.where(input_ids != config.pad_token_id     ,     1     ,     0				)
      if decoder_attention_mask is None:
            SCREAMING_SNAKE_CASE_      		=     np.where(decoder_input_ids != config.pad_token_id     ,     1     ,     0				)
      if head_mask is None:
            SCREAMING_SNAKE_CASE_      		=     np.ones((config.encoder_layers, config.encoder_attention_heads)				)
      if decoder_head_mask is None:
            SCREAMING_SNAKE_CASE_      		=     np.ones((config.decoder_layers, config.decoder_attention_heads)				)
      if cross_attn_head_mask is None:
            SCREAMING_SNAKE_CASE_      		=     np.ones((config.decoder_layers, config.decoder_attention_heads)				)
      return {
          "input_ids": input_ids,
          "decoder_input_ids": decoder_input_ids,
          "attention_mask": attention_mask,
          "decoder_attention_mask": attention_mask,
      }





class 			lowerCamelCase_				:


 '''simple docstring'''







 def __init__(				self    :						List[str]							,							_lowerCAmelCase    :						Any							,							_lowerCAmelCase    :						int=13							,							_lowerCAmelCase    :						str=7							,							_lowerCAmelCase    :						Optional[Any]=True							,							_lowerCAmelCase    :						List[str]=False							,							_lowerCAmelCase    :						Dict=99							,							_lowerCAmelCase    :						Tuple=16							,							_lowerCAmelCase    :						Optional[int]=2							,							_lowerCAmelCase    :						int=4							,							_lowerCAmelCase    :						int=4							,							_lowerCAmelCase    :						str="gelu"							,							_lowerCAmelCase    :						List[Any]=0.1							,							_lowerCAmelCase    :						Tuple=0.1							,							_lowerCAmelCase    :						Tuple=32							,							_lowerCAmelCase    :						List[Any]=2							,							_lowerCAmelCase    :						Optional[int]=1							,							_lowerCAmelCase    :						Dict=0							,							_lowerCAmelCase    :						Any=0.02							,							):
       SCREAMING_SNAKE_CASE_      		=     parent
       SCREAMING_SNAKE_CASE_      		=     batch_size
       SCREAMING_SNAKE_CASE_      		=     seq_length
       SCREAMING_SNAKE_CASE_      		=     is_training
       SCREAMING_SNAKE_CASE_      		=     use_labels
       SCREAMING_SNAKE_CASE_      		=     vocab_size
       SCREAMING_SNAKE_CASE_      		=     hidden_size
       SCREAMING_SNAKE_CASE_      		=     num_hidden_layers
       SCREAMING_SNAKE_CASE_      		=     num_attention_heads
       SCREAMING_SNAKE_CASE_      		=     intermediate_size
       SCREAMING_SNAKE_CASE_      		=     hidden_act
       SCREAMING_SNAKE_CASE_      		=     hidden_dropout_prob
       SCREAMING_SNAKE_CASE_      		=     attention_probs_dropout_prob
       SCREAMING_SNAKE_CASE_      		=     max_position_embeddings
       SCREAMING_SNAKE_CASE_      		=     eos_token_id
       SCREAMING_SNAKE_CASE_      		=     pad_token_id
       SCREAMING_SNAKE_CASE_      		=     bos_token_id
       SCREAMING_SNAKE_CASE_      		=     initializer_range



 def  lowerCAmelCase_    (				self    :						str		):
       SCREAMING_SNAKE_CASE_      		=     np.clip(ids_tensor([self.batch_size, self.seq_length - 1]							,							self.vocab_size		)							,							3							,							self.vocab_size		)
       SCREAMING_SNAKE_CASE_      		=     np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)							,							dtype=np.intaa		))							,							-1		)

       SCREAMING_SNAKE_CASE_      		=     shift_tokens_right(_lowerCAmelCase							,							1							,							2		)

       SCREAMING_SNAKE_CASE_      		=     BlenderbotConfig(
           vocab_size=self.vocab_size							,							d_model=self.hidden_size							,							encoder_layers=self.num_hidden_layers							,							decoder_layers=self.num_hidden_layers							,							encoder_attention_heads=self.num_attention_heads							,							decoder_attention_heads=self.num_attention_heads							,							encoder_ffn_dim=self.intermediate_size							,							decoder_ffn_dim=self.intermediate_size							,							dropout=self.hidden_dropout_prob							,							attention_dropout=self.attention_probs_dropout_prob							,							max_position_embeddings=self.max_position_embeddings							,							eos_token_id=self.eos_token_id							,							bos_token_id=self.bos_token_id							,							pad_token_id=self.pad_token_id							,							initializer_range=self.initializer_range							,							use_cache=_lowerCAmelCase							,							)
       SCREAMING_SNAKE_CASE_      		=     prepare_blenderbot_inputs_dict(_lowerCAmelCase							,							_lowerCAmelCase							,							_lowerCAmelCase		)
       return config, inputs_dict



 def  lowerCAmelCase_    (				self    :						Optional[Any]		):
       SCREAMING_SNAKE_CASE_		,					SCREAMING_SNAKE_CASE_      		=     self.prepare_config_and_inputs()
       return config, inputs_dict



 def  lowerCAmelCase_    (				self    :						List[str]							,							_lowerCAmelCase    :						Dict							,							_lowerCAmelCase    :						List[str]							,							_lowerCAmelCase    :						List[Any]		):
       SCREAMING_SNAKE_CASE_      		=     20
       SCREAMING_SNAKE_CASE_      		=     model_class_name(_lowerCAmelCase		)

       SCREAMING_SNAKE_CASE_      		=     model.encode(inputs_dict['input_ids']		)

       SCREAMING_SNAKE_CASE_		,					SCREAMING_SNAKE_CASE_      		=     (
           inputs_dict['decoder_input_ids'],
           inputs_dict['decoder_attention_mask'],
       )

       SCREAMING_SNAKE_CASE_      		=     model.init_cache(decoder_input_ids.shape[0]							,							_lowerCAmelCase							,							_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     jnp.ones((decoder_input_ids.shape[0], max_decoder_length)							,							dtype='i4'		)

       SCREAMING_SNAKE_CASE_      		=     jnp.broadcast_to(
           jnp.arange(decoder_input_ids.shape[-1] - 1		)[None, :]							,							(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)							,							)
       SCREAMING_SNAKE_CASE_      		=     model.decode(
           decoder_input_ids[:, :-1]							,							_lowerCAmelCase							,							decoder_attention_mask=_lowerCAmelCase							,							past_key_values=_lowerCAmelCase							,							decoder_position_ids=_lowerCAmelCase							,							)

       SCREAMING_SNAKE_CASE_      		=     jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]							,							dtype='i4'		)
       SCREAMING_SNAKE_CASE_      		=     model.decode(
           decoder_input_ids[:, -1:]							,							_lowerCAmelCase							,							decoder_attention_mask=_lowerCAmelCase							,							past_key_values=outputs_cache.past_key_values							,							decoder_position_ids=_lowerCAmelCase							,							)

       SCREAMING_SNAKE_CASE_      		=     model.decode(_lowerCAmelCase							,							_lowerCAmelCase		)

       SCREAMING_SNAKE_CASE_      		=     np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])		)		)
       self.parent.assertTrue(diff < 1E-3							,							msg=F"Max diff is {diff}"		)




 def  lowerCAmelCase_    (				self    :						int							,							_lowerCAmelCase    :						List[str]							,							_lowerCAmelCase    :						Dict							,							_lowerCAmelCase    :						int		):
       SCREAMING_SNAKE_CASE_      		=     20
       SCREAMING_SNAKE_CASE_      		=     model_class_name(_lowerCAmelCase		)

       SCREAMING_SNAKE_CASE_      		=     model.encode(inputs_dict['input_ids']		)

       SCREAMING_SNAKE_CASE_		,					SCREAMING_SNAKE_CASE_      		=     (
           inputs_dict['decoder_input_ids'],
           inputs_dict['decoder_attention_mask'],
       )

       SCREAMING_SNAKE_CASE_      		=     jnp.concatenate(
           [
               decoder_attention_mask,
               jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])		),
           ]							,							axis=-1							,							)

       SCREAMING_SNAKE_CASE_      		=     model.init_cache(decoder_input_ids.shape[0]							,							_lowerCAmelCase							,							_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     jnp.broadcast_to(
           jnp.arange(decoder_input_ids.shape[-1] - 1		)[None, :]							,							(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)							,							)

       SCREAMING_SNAKE_CASE_      		=     model.decode(
           decoder_input_ids[:, :-1]							,							_lowerCAmelCase							,							decoder_attention_mask=_lowerCAmelCase							,							past_key_values=_lowerCAmelCase							,							decoder_position_ids=_lowerCAmelCase							,							)
       SCREAMING_SNAKE_CASE_      		=     jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]							,							dtype='i4'		)
       SCREAMING_SNAKE_CASE_      		=     model.decode(
           decoder_input_ids[:, -1:]							,							_lowerCAmelCase							,							past_key_values=outputs_cache.past_key_values							,							decoder_attention_mask=_lowerCAmelCase							,							decoder_position_ids=_lowerCAmelCase							,							)

       SCREAMING_SNAKE_CASE_      		=     model.decode(_lowerCAmelCase							,							_lowerCAmelCase							,							decoder_attention_mask=_lowerCAmelCase		)

       SCREAMING_SNAKE_CASE_      		=     np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])		)		)
       self.parent.assertTrue(diff < 1E-3							,							msg=F"Max diff is {diff}"		)





@require_flax
class 			lowerCamelCase_				( unittest.TestCase   ):


 '''simple docstring'''







 lowercase_							       =       99



 def  lowerCAmelCase_    (				self    :						Optional[Any]		):
       SCREAMING_SNAKE_CASE_      		=     np.array(
           [
               [71, 82, 18, 33, 46, 91, 2],
               [68, 34, 26, 58, 30, 82, 2],
               [5, 97, 17, 39, 94, 40, 2],
               [76, 83, 94, 25, 70, 78, 2],
               [87, 59, 41, 35, 48, 66, 2],
               [55, 13, 16, 58, 5, 2, 1],  # note padding
               [64, 27, 31, 51, 12, 75, 2],
               [52, 64, 86, 17, 83, 39, 2],
               [48, 61, 9, 24, 71, 82, 2],
               [26, 1, 60, 48, 22, 13, 2],
               [21, 5, 62, 28, 14, 76, 2],
               [45, 98, 37, 86, 59, 48, 2],
               [70, 70, 50, 9, 28, 0, 2],
           ]							,							dtype=np.intaa							,							)

       SCREAMING_SNAKE_CASE_      		=     input_ids.shape[0]
       SCREAMING_SNAKE_CASE_      		=     BlenderbotConfig(
           vocab_size=self.vocab_size							,							d_model=24							,							encoder_layers=2							,							decoder_layers=2							,							encoder_attention_heads=2							,							decoder_attention_heads=2							,							encoder_ffn_dim=32							,							decoder_ffn_dim=32							,							max_position_embeddings=48							,							eos_token_id=2							,							pad_token_id=1							,							bos_token_id=0							,							)
       return config, input_ids, batch_size



 def  lowerCAmelCase_    (				self    :						Optional[Any]		):
       SCREAMING_SNAKE_CASE_		,					SCREAMING_SNAKE_CASE_		,					SCREAMING_SNAKE_CASE_      		=     self._get_config_and_data()
       SCREAMING_SNAKE_CASE_      		=     FlaxBlenderbotForConditionalGeneration(_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     lm_model(input_ids=_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     (batch_size, input_ids.shape[1], config.vocab_size)
       self.assertEqual(outputs['logits'].shape							,							_lowerCAmelCase		)



 def  lowerCAmelCase_    (				self    :						List[Any]		):
       SCREAMING_SNAKE_CASE_      		=     BlenderbotConfig(
           vocab_size=self.vocab_size							,							d_model=14							,							encoder_layers=2							,							decoder_layers=2							,							encoder_attention_heads=2							,							decoder_attention_heads=2							,							encoder_ffn_dim=8							,							decoder_ffn_dim=8							,							max_position_embeddings=48							,							)
       SCREAMING_SNAKE_CASE_      		=     FlaxBlenderbotForConditionalGeneration(_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]							,							dtype=np.intaa		)
       SCREAMING_SNAKE_CASE_      		=     np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]							,							dtype=np.intaa		)
       SCREAMING_SNAKE_CASE_      		=     lm_model(input_ids=_lowerCAmelCase							,							decoder_input_ids=_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     (*summary.shape, config.vocab_size)
       self.assertEqual(outputs['logits'].shape							,							_lowerCAmelCase		)




 def  lowerCAmelCase_    (				self    :						Union[str, Any]		):
       SCREAMING_SNAKE_CASE_      		=     np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]							,							dtype=np.intaa		)
       SCREAMING_SNAKE_CASE_      		=     shift_tokens_right(_lowerCAmelCase							,							1							,							2		)
       SCREAMING_SNAKE_CASE_      		=     np.equal(_lowerCAmelCase							,							1		).astype(np.floataa		).sum()
       SCREAMING_SNAKE_CASE_      		=     np.equal(_lowerCAmelCase							,							1		).astype(np.floataa		).sum()
       self.assertEqual(shifted.shape							,							input_ids.shape		)
       self.assertEqual(_lowerCAmelCase							,							n_pad_before - 1		)
       self.assertTrue(np.equal(shifted[:, 0]							,							2		).all()		)





@require_flax
class 			lowerCamelCase_				( _SCREAMING_SNAKE_CASE			,  unittest.TestCase			,  _SCREAMING_SNAKE_CASE   ):


 '''simple docstring'''







 lowercase_							       =       True
 lowercase_							       =       (
     (
         FlaxBlenderbotModel,
         FlaxBlenderbotForConditionalGeneration,
     )
     if is_flax_available()
     else ()
 )
 lowercase_							       =       (FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()



 def  lowerCAmelCase_    (				self    :						Union[str, Any]		):
       SCREAMING_SNAKE_CASE_      		=     FlaxBlenderbotModelTester(self		)



 def  lowerCAmelCase_    (				self    :						Optional[Any]		):
       SCREAMING_SNAKE_CASE_		,					SCREAMING_SNAKE_CASE_      		=     self.model_tester.prepare_config_and_inputs()
       for model_class in self.all_model_classes:
             self.model_tester.check_use_cache_forward(_lowerCAmelCase							,							_lowerCAmelCase							,							_lowerCAmelCase		)



 def  lowerCAmelCase_    (				self    :						str		):
       SCREAMING_SNAKE_CASE_		,					SCREAMING_SNAKE_CASE_      		=     self.model_tester.prepare_config_and_inputs()
       for model_class in self.all_model_classes:
             self.model_tester.check_use_cache_forward_with_attn_mask(_lowerCAmelCase							,							_lowerCAmelCase							,							_lowerCAmelCase		)



 def  lowerCAmelCase_    (				self    :						Tuple		):
       SCREAMING_SNAKE_CASE_		,					SCREAMING_SNAKE_CASE_      		=     self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             with self.subTest(model_class.__name__		):
                   SCREAMING_SNAKE_CASE_      		=     self._prepare_for_class(_lowerCAmelCase							,							_lowerCAmelCase		)
                   SCREAMING_SNAKE_CASE_      		=     model_class(_lowerCAmelCase		)

                   @jax.jit
                   def encode_jitted(_lowerCAmelCase    :						str							,							_lowerCAmelCase    :						Optional[Any]=None							,							**_lowerCAmelCase    :						str		):
                         return model.encode(input_ids=_lowerCAmelCase							,							attention_mask=_lowerCAmelCase		)

                   with self.subTest('JIT Enabled'		):
                         SCREAMING_SNAKE_CASE_      		=     encode_jitted(**_lowerCAmelCase		).to_tuple()

                   with self.subTest('JIT Disabled'		):
                         with jax.disable_jit():
                               SCREAMING_SNAKE_CASE_      		=     encode_jitted(**_lowerCAmelCase		).to_tuple()

                   self.assertEqual(len(_lowerCAmelCase		)							,							len(_lowerCAmelCase		)		)
                   for jitted_output, output in zip(_lowerCAmelCase							,							_lowerCAmelCase		):
                         self.assertEqual(jitted_output.shape							,							output.shape		)



 def  lowerCAmelCase_    (				self    :						Union[str, Any]		):
       SCREAMING_SNAKE_CASE_		,					SCREAMING_SNAKE_CASE_      		=     self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             with self.subTest(model_class.__name__		):
                   SCREAMING_SNAKE_CASE_      		=     model_class(_lowerCAmelCase		)
                   SCREAMING_SNAKE_CASE_      		=     model.encode(inputs_dict['input_ids']							,							inputs_dict['attention_mask']		)

                   SCREAMING_SNAKE_CASE_      		=     {
                       'decoder_input_ids': inputs_dict['decoder_input_ids'],
                       'decoder_attention_mask': inputs_dict['decoder_attention_mask'],
                       'encoder_outputs': encoder_outputs,
                   }

                   @jax.jit
                   def decode_jitted(_lowerCAmelCase    :						int							,							_lowerCAmelCase    :						List[Any]							,							_lowerCAmelCase    :						Dict		):
                         return model.decode(
                             decoder_input_ids=_lowerCAmelCase							,							decoder_attention_mask=_lowerCAmelCase							,							encoder_outputs=_lowerCAmelCase							,							)

                   with self.subTest('JIT Enabled'		):
                         SCREAMING_SNAKE_CASE_      		=     decode_jitted(**_lowerCAmelCase		).to_tuple()

                   with self.subTest('JIT Disabled'		):
                         with jax.disable_jit():
                               SCREAMING_SNAKE_CASE_      		=     decode_jitted(**_lowerCAmelCase		).to_tuple()

                   self.assertEqual(len(_lowerCAmelCase		)							,							len(_lowerCAmelCase		)		)
                   for jitted_output, output in zip(_lowerCAmelCase							,							_lowerCAmelCase		):
                         self.assertEqual(jitted_output.shape							,							output.shape		)



 @slow
 def  lowerCAmelCase_    (				self    :						Union[str, Any]		):
       for model_class_name in self.all_model_classes:
             SCREAMING_SNAKE_CASE_      		=     model_class_name.from_pretrained('facebook/blenderbot-400M-distill'		)
             # FlaxBlenderbotForSequenceClassification expects eos token in input_ids
             SCREAMING_SNAKE_CASE_      		=     np.ones((1, 1)		) * model.config.eos_token_id
             SCREAMING_SNAKE_CASE_      		=     model(_lowerCAmelCase		)
             self.assertIsNotNone(_lowerCAmelCase		)




 @unittest.skipUnless(jax_device != 'cpu'							,							'3B test too slow on CPU.'		)
 @slow
 def  lowerCAmelCase_    (				self    :						Optional[Any]		):
       SCREAMING_SNAKE_CASE_      		=     {'num_beams': 1, 'early_stopping': True, 'min_length': 15, 'max_length': 25}
       SCREAMING_SNAKE_CASE_      		=     {'skip_special_tokens': True, 'clean_up_tokenization_spaces': True}

       SCREAMING_SNAKE_CASE_      		=     FlaxBlenderbotForConditionalGeneration.from_pretrained('facebook/blenderbot-3B'							,							from_pt=_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     BlenderbotTokenizer.from_pretrained('facebook/blenderbot-3B'		)

       SCREAMING_SNAKE_CASE_      		=     ['Sam']
       SCREAMING_SNAKE_CASE_      		=     tokenizer(_lowerCAmelCase							,							return_tensors='jax'		)

       SCREAMING_SNAKE_CASE_      		=     model.generate(**_lowerCAmelCase							,							**_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     'Sam is a great name. It means "sun" in Gaelic.'

       SCREAMING_SNAKE_CASE_      		=     tokenizer.batch_decode(_lowerCAmelCase							,							**_lowerCAmelCase		)
       assert generated_txt[0].strip() == tgt_text





import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizer

from diffusers import (
    AltDiffusionImgaImgPipeline,
    AutoencoderKL,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
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


enable_full_determinism()





class 			lowerCamelCase_				( unittest.TestCase   ):


 '''simple docstring'''







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



 @property
 def  lowerCAmelCase_    (				self    :						Union[str, Any]		):
       SCREAMING_SNAKE_CASE_      		=     1
       SCREAMING_SNAKE_CASE_      		=     3
       SCREAMING_SNAKE_CASE_      		=     (32, 32)

       SCREAMING_SNAKE_CASE_      		=     floats_tensor((batch_size, num_channels) + sizes							,							rng=random.Random(0		)		).to(_lowerCAmelCase		)
       return image



 @property
 def  lowerCAmelCase_    (				self    :						Union[str, Any]		):
       torch.manual_seed(0		)
       SCREAMING_SNAKE_CASE_      		=     UNetaDConditionModel(
           block_out_channels=(32, 64)							,							layers_per_block=2							,							sample_size=32							,							in_channels=4							,							out_channels=4							,							down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')							,							up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')							,							cross_attention_dim=32							,							)
       return model



 @property
 def  lowerCAmelCase_    (				self    :						Tuple		):
       torch.manual_seed(0		)
       SCREAMING_SNAKE_CASE_      		=     AutoencoderKL(
           block_out_channels=[32, 64]							,							in_channels=3							,							out_channels=3							,							down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']							,							up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']							,							latent_channels=4							,							)
       return model



 @property
 def  lowerCAmelCase_    (				self    :						Optional[int]		):
       torch.manual_seed(0		)
       SCREAMING_SNAKE_CASE_      		=     RobertaSeriesConfig(
           hidden_size=32							,							project_dim=32							,							intermediate_size=37							,							layer_norm_eps=1E-05							,							num_attention_heads=4							,							num_hidden_layers=5							,							pad_token_id=1							,							vocab_size=5_006							,							)
       return RobertaSeriesModelWithTransformation(_lowerCAmelCase		)



 @property
 def  lowerCAmelCase_    (				self    :						List[Any]		):
       def extract(*_lowerCAmelCase    :						Optional[int]							,							**_lowerCAmelCase    :						str		):





             class 			lowerCamelCase_				:


              '''simple docstring'''







              def __init__(				self    :						str		):
                    SCREAMING_SNAKE_CASE_      		=     torch.ones([0]		)




              def  lowerCAmelCase_    (				self    :						Union[str, Any]							,							_lowerCAmelCase    :						int		):
                    self.pixel_values.to(_lowerCAmelCase		)
                    return self


             return Out()

       return extract



 def  lowerCAmelCase_    (				self    :						Optional[Any]		):
       SCREAMING_SNAKE_CASE_      		=     'cpu'  # ensure determinism for the device-dependent torch.Generator
       SCREAMING_SNAKE_CASE_      		=     self.dummy_cond_unet
       SCREAMING_SNAKE_CASE_      		=     PNDMScheduler(skip_prk_steps=_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     self.dummy_vae
       SCREAMING_SNAKE_CASE_      		=     self.dummy_text_encoder
       SCREAMING_SNAKE_CASE_      		=     XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'		)
       SCREAMING_SNAKE_CASE_      		=     77

       SCREAMING_SNAKE_CASE_      		=     self.dummy_image.to(_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     init_image / 2 + 0.5

       # make sure here that pndm scheduler skips prk
       SCREAMING_SNAKE_CASE_      		=     AltDiffusionImgaImgPipeline(
           unet=_lowerCAmelCase							,							scheduler=_lowerCAmelCase							,							vae=_lowerCAmelCase							,							text_encoder=_lowerCAmelCase							,							tokenizer=_lowerCAmelCase							,							safety_checker=_lowerCAmelCase							,							feature_extractor=self.dummy_extractor							,							)
       SCREAMING_SNAKE_CASE_      		=     VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor							,							do_normalize=_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     alt_pipe.to(_lowerCAmelCase		)
       alt_pipe.set_progress_bar_config(disable=_lowerCAmelCase		)

       SCREAMING_SNAKE_CASE_      		=     'A painting of a squirrel eating a burger'
       SCREAMING_SNAKE_CASE_      		=     torch.Generator(device=_lowerCAmelCase		).manual_seed(0		)
       SCREAMING_SNAKE_CASE_      		=     alt_pipe(
           [prompt]							,							generator=_lowerCAmelCase							,							guidance_scale=6.0							,							num_inference_steps=2							,							output_type='np'							,							image=_lowerCAmelCase							,							)

       SCREAMING_SNAKE_CASE_      		=     output.images

       SCREAMING_SNAKE_CASE_      		=     torch.Generator(device=_lowerCAmelCase		).manual_seed(0		)
       SCREAMING_SNAKE_CASE_      		=     alt_pipe(
           [prompt]							,							generator=_lowerCAmelCase							,							guidance_scale=6.0							,							num_inference_steps=2							,							output_type='np'							,							image=_lowerCAmelCase							,							return_dict=_lowerCAmelCase							,							)[0]

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

       assert image.shape == (1, 32, 32, 3)
       SCREAMING_SNAKE_CASE_      		=     np.array([0.4427, 0.3731, 0.4249, 0.4941, 0.4546, 0.4148, 0.4193, 0.4666, 0.4499]		)

       assert np.abs(image_slice.flatten() - expected_slice		).max() < 5E-3
       assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 5E-3



 @unittest.skipIf(torch_device != 'cuda'							,							'This test requires a GPU'		)
 def  lowerCAmelCase_    (				self    :						Tuple		):
       SCREAMING_SNAKE_CASE_      		=     self.dummy_cond_unet
       SCREAMING_SNAKE_CASE_      		=     PNDMScheduler(skip_prk_steps=_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     self.dummy_vae
       SCREAMING_SNAKE_CASE_      		=     self.dummy_text_encoder
       SCREAMING_SNAKE_CASE_      		=     XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'		)
       SCREAMING_SNAKE_CASE_      		=     77

       SCREAMING_SNAKE_CASE_      		=     self.dummy_image.to(_lowerCAmelCase		)

       # put models in fp16
       SCREAMING_SNAKE_CASE_      		=     unet.half()
       SCREAMING_SNAKE_CASE_      		=     vae.half()
       SCREAMING_SNAKE_CASE_      		=     bert.half()

       # make sure here that pndm scheduler skips prk
       SCREAMING_SNAKE_CASE_      		=     AltDiffusionImgaImgPipeline(
           unet=_lowerCAmelCase							,							scheduler=_lowerCAmelCase							,							vae=_lowerCAmelCase							,							text_encoder=_lowerCAmelCase							,							tokenizer=_lowerCAmelCase							,							safety_checker=_lowerCAmelCase							,							feature_extractor=self.dummy_extractor							,							)
       SCREAMING_SNAKE_CASE_      		=     VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor							,							do_normalize=_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     alt_pipe.to(_lowerCAmelCase		)
       alt_pipe.set_progress_bar_config(disable=_lowerCAmelCase		)

       SCREAMING_SNAKE_CASE_      		=     'A painting of a squirrel eating a burger'
       SCREAMING_SNAKE_CASE_      		=     torch.manual_seed(0		)
       SCREAMING_SNAKE_CASE_      		=     alt_pipe(
           [prompt]							,							generator=_lowerCAmelCase							,							num_inference_steps=2							,							output_type='np'							,							image=_lowerCAmelCase							,							).images

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




 @unittest.skipIf(torch_device != 'cuda'							,							'This test requires a GPU'		)
 def  lowerCAmelCase_    (				self    :						str		):
       SCREAMING_SNAKE_CASE_      		=     load_image(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
           '/img2img/sketch-mountains-input.jpg'		)
       # resize to resolution that is divisible by 8 but not 16 or 32
       SCREAMING_SNAKE_CASE_      		=     init_image.resize((760, 504)		)

       SCREAMING_SNAKE_CASE_      		=     'BAAI/AltDiffusion'
       SCREAMING_SNAKE_CASE_      		=     AltDiffusionImgaImgPipeline.from_pretrained(
           _lowerCAmelCase							,							safety_checker=_lowerCAmelCase							,							)
       pipe.to(_lowerCAmelCase		)
       pipe.set_progress_bar_config(disable=_lowerCAmelCase		)
       pipe.enable_attention_slicing()

       SCREAMING_SNAKE_CASE_      		=     'A fantasy landscape, trending on artstation'

       SCREAMING_SNAKE_CASE_      		=     torch.manual_seed(0		)
       SCREAMING_SNAKE_CASE_      		=     pipe(
           prompt=_lowerCAmelCase							,							image=_lowerCAmelCase							,							strength=0.75							,							guidance_scale=7.5							,							generator=_lowerCAmelCase							,							output_type='np'							,							)
       SCREAMING_SNAKE_CASE_      		=     output.images[0]

       SCREAMING_SNAKE_CASE_      		=     image[255:258, 383:386, -1]

       assert image.shape == (504, 760, 3)
       SCREAMING_SNAKE_CASE_      		=     np.array([0.9358, 0.9397, 0.9599, 0.9901, 1.0000, 1.0000, 0.9882, 1.0000, 1.0000]		)

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





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


 '''simple docstring'''







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




 def  lowerCAmelCase_    (				self    :						int		):
       SCREAMING_SNAKE_CASE_      		=     load_image(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
           '/img2img/sketch-mountains-input.jpg'		)
       SCREAMING_SNAKE_CASE_      		=     init_image.resize((768, 512)		)
       SCREAMING_SNAKE_CASE_      		=     load_numpy(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy'		)

       SCREAMING_SNAKE_CASE_      		=     'BAAI/AltDiffusion'
       SCREAMING_SNAKE_CASE_      		=     AltDiffusionImgaImgPipeline.from_pretrained(
           _lowerCAmelCase							,							safety_checker=_lowerCAmelCase							,							)
       pipe.to(_lowerCAmelCase		)
       pipe.set_progress_bar_config(disable=_lowerCAmelCase		)
       pipe.enable_attention_slicing()

       SCREAMING_SNAKE_CASE_      		=     'A fantasy landscape, trending on artstation'

       SCREAMING_SNAKE_CASE_      		=     torch.manual_seed(0		)
       SCREAMING_SNAKE_CASE_      		=     pipe(
           prompt=_lowerCAmelCase							,							image=_lowerCAmelCase							,							strength=0.75							,							guidance_scale=7.5							,							generator=_lowerCAmelCase							,							output_type='np'							,							)
       SCREAMING_SNAKE_CASE_      		=     output.images[0]

       assert image.shape == (512, 768, 3)
       # img2img is flaky across GPUs even in fp32, so using MAE here
       assert np.abs(expected_image - image		).max() < 1E-2