





'''simple docstring'''






import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()


class   _lowercase						(							a		,					a		,					a		,					unittest.TestCase				):
   _UpperCamelCase       		=   StableUnCLIPPipeline
   _UpperCamelCase       		=   TEXT_TO_IMAGE_PARAMS
   _UpperCamelCase       		=   TEXT_TO_IMAGE_BATCH_PARAMS
   _UpperCamelCase       		=   TEXT_TO_IMAGE_IMAGE_PARAMS
   _UpperCamelCase       		=   TEXT_TO_IMAGE_IMAGE_PARAMS

   # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
   _UpperCamelCase       		=   False


   def 					snake_case				( self					):
       A       :       List[Any]							=						32
       A       :       Optional[Any]							=						embedder_hidden_size

       # prior components

       torch.manual_seed(0					)
       A       :       str							=						CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)

       torch.manual_seed(0					)
       A       :       List[str]							=						CLIPTextModelWithProjection(
           CLIPTextConfig(
               bos_token_id=0    ,   eos_token_id=2    ,   hidden_size=_UpperCAmelCase    ,   projection_dim=_UpperCAmelCase    ,   intermediate_size=37    ,   layer_norm_eps=1E-0_5    ,   num_attention_heads=4    ,   num_hidden_layers=5    ,   pad_token_id=1    ,   vocab_size=1_000    ,   )					)

       torch.manual_seed(0					)
       A       :       Optional[int]							=						PriorTransformer(
           num_attention_heads=2    ,   attention_head_dim=12    ,   embedding_dim=_UpperCAmelCase    ,   num_layers=1    ,   )

       torch.manual_seed(0					)
       A       :       Tuple							=						DDPMScheduler(
           variance_type='''fixed_small_log'''    ,   prediction_type='''sample'''    ,   num_train_timesteps=1_000    ,   clip_sample=_UpperCAmelCase    ,   clip_sample_range=5.0    ,   beta_schedule='''squaredcos_cap_v2'''    ,   )

       # regular denoising components

       torch.manual_seed(0					)
       A       :       Any							=						StableUnCLIPImageNormalizer(embedding_dim=_UpperCAmelCase					)
       A       :       List[str]							=						DDPMScheduler(beta_schedule='''squaredcos_cap_v2'''					)

       torch.manual_seed(0					)
       A       :       Tuple							=						CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)

       torch.manual_seed(0					)
       A       :       Dict							=						CLIPTextModel(
           CLIPTextConfig(
               bos_token_id=0    ,   eos_token_id=2    ,   hidden_size=_UpperCAmelCase    ,   projection_dim=32    ,   intermediate_size=37    ,   layer_norm_eps=1E-0_5    ,   num_attention_heads=4    ,   num_hidden_layers=5    ,   pad_token_id=1    ,   vocab_size=1_000    ,   )					)

       torch.manual_seed(0					)
       A       :       List[str]							=						UNetaDConditionModel(
           sample_size=32    ,   in_channels=4    ,   out_channels=4    ,   down_block_types=('''CrossAttnDownBlock2D''', '''DownBlock2D''')    ,   up_block_types=('''UpBlock2D''', '''CrossAttnUpBlock2D''')    ,   block_out_channels=(32, 64)    ,   attention_head_dim=(2, 4)    ,   class_embed_type='''projection'''    ,   projection_class_embeddings_input_dim=embedder_projection_dim * 2    ,   cross_attention_dim=_UpperCAmelCase    ,   layers_per_block=1    ,   upcast_attention=_UpperCAmelCase    ,   use_linear_projection=_UpperCAmelCase    ,   )

       torch.manual_seed(0					)
       A       :       Any							=						DDIMScheduler(
           beta_schedule='''scaled_linear'''    ,   beta_start=0.00085    ,   beta_end=0.012    ,   prediction_type='''v_prediction'''    ,   set_alpha_to_one=_UpperCAmelCase    ,   steps_offset=1    ,   )

       torch.manual_seed(0					)
       A       :       List[Any]							=						AutoencoderKL()

       A       :       List[Any]							=						{
           # prior components
           '''prior_tokenizer''': prior_tokenizer,
           '''prior_text_encoder''': prior_text_encoder,
           '''prior''': prior,
           '''prior_scheduler''': prior_scheduler,
           # image noising components
           '''image_normalizer''': image_normalizer,
           '''image_noising_scheduler''': image_noising_scheduler,
           # regular denoising components
           '''tokenizer''': tokenizer,
           '''text_encoder''': text_encoder,
           '''unet''': unet,
           '''scheduler''': scheduler,
           '''vae''': vae,
       }

       return components


   def 					snake_case				( self    ,   _UpperCAmelCase    ,   _UpperCAmelCase=0					):
       if str(_UpperCAmelCase					).startswith('''mps'''					):
           A       :       Tuple							=						torch.manual_seed(_UpperCAmelCase					)
       else:
           A       :       int							=						torch.Generator(device=_UpperCAmelCase					).manual_seed(_UpperCAmelCase					)
       A       :       List[str]							=						{
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''prior_num_inference_steps''': 2,
           '''output_type''': '''numpy''',
       }
       return inputs


   def 					snake_case				( self					):
       A       :       Any							=						torch_device == '''cpu'''

       self._test_attention_slicing_forward_pass(test_max_difference=_UpperCAmelCase					)


   def 					snake_case				( self					):
       A       :       Dict							=						torch_device in ['''cpu''', '''mps''']

       self._test_inference_batch_single_identical(test_max_difference=_UpperCAmelCase					)







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


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


   def 					snake_case				( self					):
       A       :       Optional[Any]							=						load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy'''					)

       A       :       Dict							=						StableUnCLIPPipeline.from_pretrained('''fusing/stable-unclip-2-1-l'''    ,   torch_dtype=torch.floataa					)
       pipe.to(_UpperCAmelCase					)
       pipe.set_progress_bar_config(disable=_UpperCAmelCase					)
       # stable unclip will oom when integration tests are run on a V100,
       # so turn on memory savings
       pipe.enable_attention_slicing()
       pipe.enable_sequential_cpu_offload()

       A       :       int							=						torch.Generator(device='''cpu'''					).manual_seed(0					)
       A       :       Union[str, Any]							=						pipe('''anime turle'''    ,   generator=_UpperCAmelCase    ,   output_type='''np'''					)

       A       :       str							=						output.images[0]

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

       assert_mean_pixel_difference(_UpperCAmelCase    ,   _UpperCAmelCase					)


   def 					snake_case				( self					):
       torch.cuda.empty_cache()
       torch.cuda.reset_max_memory_allocated()
       torch.cuda.reset_peak_memory_stats()

       A       :       List[str]							=						StableUnCLIPPipeline.from_pretrained('''fusing/stable-unclip-2-1-l'''    ,   torch_dtype=torch.floataa					)
       A       :       Any							=						pipe.to(_UpperCAmelCase					)
       pipe.set_progress_bar_config(disable=_UpperCAmelCase					)
       pipe.enable_attention_slicing()
       pipe.enable_sequential_cpu_offload()

       A       :       List[Any]							=						pipe(
           '''anime turtle'''    ,   prior_num_inference_steps=2    ,   num_inference_steps=2    ,   output_type='''np'''    ,   )

       A       :       List[str]							=						torch.cuda.max_memory_allocated()
       # make sure that less than 7 GB is allocated
       assert mem_bytes < 7 * 10**9







'''simple docstring'''






import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
     import os

     import jax.numpy as jnp
     from jax import jit

     from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
     from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

     snake_case_     =	"""0.12"""  # assumed parallelism: 8


if is_torch_available():
     import torch
def 							_lowerCamelCase(  UpperCamelCase__    : int    ,       UpperCamelCase__    : Optional[int]    ,       UpperCamelCase__    : Any=None		)		->  List[Any]:
    if rng is None:
        A       :       int							=						random.Random()

    A       :       int							=						1
    for dim in shape:
        total_dims *= dim

    A       :       Union[str, Any]							=						[]
    for _ in range(UpperCamelCase__		):
        values.append(rng.randint(0    ,       vocab_size - 1		)		)

    A       :       Union[str, Any]							=						np.array(UpperCamelCase__    ,       dtype=jnp.intaa		).reshape(UpperCamelCase__		)

    return output
def 							_lowerCamelCase(  UpperCamelCase__    : Dict    ,       UpperCamelCase__    : Optional[Any]=None		)		->  Dict:
    A       :       Union[str, Any]							=						ids_tensor(UpperCamelCase__    ,       vocab_size=2    ,       rng=UpperCamelCase__		)
    # make sure that at least one token is attended to for each batch
    A       :       Any							=						1
    return attn_mask


@require_flax
class   _lowercase						:
   _UpperCamelCase       		=   None
   _UpperCamelCase       		=   ()


   def 					snake_case				( self					):
       A,	A       :       Optional[int]							=						self.model_tester.prepare_config_and_inputs_for_common()

       # cut to half length & take max batch_size 3
       A       :       Dict							=						2
       A       :       List[str]							=						inputs['''input_ids'''].shape[-1] // 2
       A       :       int							=						inputs['''input_ids'''][:max_batch_size, :sequence_length]

       A       :       Optional[Any]							=						jnp.ones_like(_UpperCAmelCase					)
       A       :       int							=						attention_mask[:max_batch_size, :sequence_length]

       # generate max 5 tokens
       A       :       Optional[int]							=						input_ids.shape[-1] + 5
       if config.eos_token_id is not None and config.pad_token_id is None:
           # hack to allow generate for models such as GPT2 as is done in `generate()`
           A       :       str							=						config.eos_token_id
       return config, input_ids, attention_mask, max_length


   @is_pt_flax_cross_test
   def 					snake_case				( self					):
       A,	A,	A,	A       :       Optional[Any]							=						self._get_input_ids_and_config()
       A       :       Optional[int]							=						False
       A       :       Union[str, Any]							=						max_length
       A       :       str							=						0

       for model_class in self.all_generative_model_classes:
           A       :       int							=						model_class(_UpperCAmelCase					)

           A       :       Dict							=						model_class.__name__[4:]  # Skip the "Flax" at the beginning
           A       :       int							=						getattr(_UpperCAmelCase    ,   _UpperCAmelCase					)
           A       :       Union[str, Any]							=						pt_model_class(_UpperCAmelCase					).eval()
           A       :       Tuple							=						load_flax_weights_in_pytorch_model(_UpperCAmelCase    ,   flax_model.params					)

           A       :       Union[str, Any]							=						flax_model.generate(_UpperCAmelCase					).sequences
           A       :       Optional[int]							=						pt_model.generate(torch.tensor(_UpperCAmelCase    ,   dtype=torch.long					)					)

           if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
               A       :       List[Any]							=						flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

           self.assertListEqual(pt_generation_outputs.numpy().tolist()    ,   flax_generation_outputs.tolist()					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       str							=						self._get_input_ids_and_config()
       A       :       Optional[int]							=						False
       A       :       Dict							=						max_length

       for model_class in self.all_generative_model_classes:
           A       :       Any							=						model_class(_UpperCAmelCase					)

           A       :       Dict							=						model.generate(_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[-1]    ,   _UpperCAmelCase					)

           A       :       str							=						jit(model.generate					)
           A       :       List[Any]							=						jit_generate(_UpperCAmelCase					).sequences

           self.assertListEqual(generation_outputs.tolist()    ,   jit_generation_outputs.tolist()					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       List[Any]							=						self._get_input_ids_and_config()
       A       :       List[str]							=						True
       A       :       List[Any]							=						max_length

       for model_class in self.all_generative_model_classes:
           A       :       Any							=						model_class(_UpperCAmelCase					)

           A       :       Tuple							=						model.generate(_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[-1]    ,   _UpperCAmelCase					)

           A       :       Tuple							=						jit(model.generate					)
           A       :       int							=						jit_generate(_UpperCAmelCase					).sequences

           self.assertListEqual(generation_outputs.tolist()    ,   jit_generation_outputs.tolist()					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       List[str]							=						self._get_input_ids_and_config()
       A       :       Any							=						False
       A       :       str							=						max_length
       A       :       Optional[int]							=						2

       for model_class in self.all_generative_model_classes:
           A       :       Dict							=						model_class(_UpperCAmelCase					)

           A       :       Any							=						model.generate(_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[-1]    ,   _UpperCAmelCase					)

           A       :       List[Any]							=						jit(model.generate					)
           A       :       Optional[Any]							=						jit_generate(_UpperCAmelCase					).sequences

           self.assertListEqual(generation_outputs.tolist()    ,   jit_generation_outputs.tolist()					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       Optional[int]							=						self._get_input_ids_and_config()
       A       :       Dict							=						False
       A       :       List[Any]							=						max_length
       A       :       Optional[int]							=						2
       A       :       int							=						2

       for model_class in self.all_generative_model_classes:
           A       :       List[str]							=						model_class(_UpperCAmelCase					)

           A       :       List[str]							=						model.generate(_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[0]    ,   input_ids.shape[0] * config.num_return_sequences					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       List[Any]							=						self._get_input_ids_and_config()
       A       :       Any							=						True
       A       :       Any							=						max_length
       A       :       Any							=						0.8
       A       :       List[Any]							=						10
       A       :       List[Any]							=						0.3
       A       :       Union[str, Any]							=						1
       A       :       Any							=						8
       A       :       List[Any]							=						9

       for model_class in self.all_generative_model_classes:
           A       :       List[Any]							=						model_class(_UpperCAmelCase					)

           A       :       Any							=						model.generate(_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[-1]    ,   _UpperCAmelCase					)

           A       :       int							=						jit(model.generate					)
           A       :       List[Any]							=						jit_generate(_UpperCAmelCase					).sequences

           self.assertListEqual(generation_outputs.tolist()    ,   jit_generation_outputs.tolist()					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       Optional[Any]							=						self._get_input_ids_and_config()
       A       :       int							=						max_length
       A       :       Dict							=						1
       A       :       Tuple							=						8
       A       :       Optional[Any]							=						9

       for model_class in self.all_generative_model_classes:
           A       :       Dict							=						model_class(_UpperCAmelCase					)

           A       :       str							=						model.generate(_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[-1]    ,   _UpperCAmelCase					)

           A       :       int							=						jit(model.generate					)
           A       :       Optional[int]							=						jit_generate(_UpperCAmelCase					).sequences

           self.assertListEqual(generation_outputs.tolist()    ,   jit_generation_outputs.tolist()					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       int							=						self._get_input_ids_and_config()
       A       :       List[Any]							=						max_length
       A       :       Optional[Any]							=						2
       A       :       Tuple							=						1
       A       :       List[Any]							=						8
       A       :       Dict							=						9

       for model_class in self.all_generative_model_classes:
           A       :       Union[str, Any]							=						model_class(_UpperCAmelCase					)

           A       :       Any							=						model.generate(_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[-1]    ,   _UpperCAmelCase					)

           A       :       Union[str, Any]							=						jit(model.generate					)
           A       :       int							=						jit_generate(_UpperCAmelCase					).sequences

           self.assertListEqual(generation_outputs.tolist()    ,   jit_generation_outputs.tolist()					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       Dict							=						self._get_input_ids_and_config()

       # pad attention mask on the left
       A       :       Union[str, Any]							=						attention_mask.at[(0, 0)].set(0					)

       A       :       int							=						False
       A       :       List[Any]							=						max_length

       for model_class in self.all_generative_model_classes:
           A       :       List[str]							=						model_class(_UpperCAmelCase					)

           A       :       Dict							=						model.generate(_UpperCAmelCase    ,   attention_mask=_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[-1]    ,   _UpperCAmelCase					)

           A       :       Dict							=						jit(model.generate					)
           A       :       Optional[Any]							=						jit_generate(_UpperCAmelCase    ,   attention_mask=_UpperCAmelCase					).sequences

           self.assertListEqual(generation_outputs.tolist()    ,   jit_generation_outputs.tolist()					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       Optional[Any]							=						self._get_input_ids_and_config()

       # pad attention mask on the left
       A       :       Optional[Any]							=						attention_mask.at[(0, 0)].set(0					)

       A       :       Optional[int]							=						True
       A       :       int							=						max_length

       for model_class in self.all_generative_model_classes:
           A       :       int							=						model_class(_UpperCAmelCase					)

           A       :       Union[str, Any]							=						model.generate(_UpperCAmelCase    ,   attention_mask=_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[-1]    ,   _UpperCAmelCase					)

           A       :       Optional[Any]							=						jit(model.generate					)
           A       :       Dict							=						jit_generate(_UpperCAmelCase    ,   attention_mask=_UpperCAmelCase					).sequences

           self.assertListEqual(generation_outputs.tolist()    ,   jit_generation_outputs.tolist()					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       Optional[Any]							=						self._get_input_ids_and_config()

       # pad attention mask on the left
       A       :       Optional[int]							=						attention_mask.at[(0, 0)].set(0					)

       A       :       Any							=						2
       A       :       Any							=						max_length

       for model_class in self.all_generative_model_classes:
           A       :       Dict							=						model_class(_UpperCAmelCase					)

           A       :       Optional[Any]							=						model.generate(_UpperCAmelCase    ,   attention_mask=_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[-1]    ,   _UpperCAmelCase					)

           A       :       str							=						jit(model.generate					)
           A       :       List[str]							=						jit_generate(_UpperCAmelCase    ,   attention_mask=_UpperCAmelCase					).sequences

           self.assertListEqual(generation_outputs.tolist()    ,   jit_generation_outputs.tolist()					)







@require_flax
class   _lowercase						(							unittest.TestCase				):


   def 					snake_case				( self					):
       A       :       List[Any]							=						AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-bert'''					)
       A       :       Dict							=						FlaxAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''					)

       A       :       Union[str, Any]							=						'''Hello world'''
       A       :       Union[str, Any]							=						tokenizer(_UpperCAmelCase    ,   return_tensors='''np'''					).input_ids

       # typos are quickly detected (the correct argument is `do_sample`)
       with self.assertRaisesRegex(_UpperCAmelCase    ,   '''do_samples'''					):
           model.generate(_UpperCAmelCase    ,   do_samples=_UpperCAmelCase					)

       # arbitrary arguments that will not be used anywhere are also not accepted
       with self.assertRaisesRegex(_UpperCAmelCase    ,   '''foo'''					):
           A       :       List[Any]							=						{'''foo''': '''bar'''}
           model.generate(_UpperCAmelCase    ,   **_UpperCAmelCase					)
