

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 	UpperCAmelCase__		(	A__						,					A__						,					A__						,					unittest.TestCase	):






  """simple docstring"""


  a									= StableUnCLIPPipeline
  a									= TEXT_TO_IMAGE_PARAMS
  a									= TEXT_TO_IMAGE_BATCH_PARAMS
  a									= TEXT_TO_IMAGE_IMAGE_PARAMS
  a									= TEXT_TO_IMAGE_IMAGE_PARAMS

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






  def 				lowercase_		(							self						:		Dict )		->							List[str]:
       SCREAMING_SNAKE_CASE__   =   32
       SCREAMING_SNAKE_CASE__   =   embedder_hidden_size

       # prior components

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

       torch.manual_seed(0 )
       SCREAMING_SNAKE_CASE__   =   CLIPTextModelWithProjection(
           CLIPTextConfig(
               bos_token_id=0     ,       eos_token_id=2     ,       hidden_size=__lowerCamelCase     ,       projection_dim=__lowerCamelCase     ,       intermediate_size=37     ,       layer_norm_eps=1e-05     ,       num_attention_heads=4     ,       num_hidden_layers=5     ,       pad_token_id=1     ,       vocab_size=1000     ,       ) )

       torch.manual_seed(0 )
       SCREAMING_SNAKE_CASE__   =   PriorTransformer(
           num_attention_heads=2     ,       attention_head_dim=12     ,       embedding_dim=__lowerCamelCase     ,       num_layers=1     ,       )

       torch.manual_seed(0 )
       SCREAMING_SNAKE_CASE__   =   DDPMScheduler(
           variance_type='''fixed_small_log'''     ,       prediction_type='''sample'''     ,       num_train_timesteps=1000     ,       clip_sample=__lowerCamelCase     ,       clip_sample_range=5.0     ,       beta_schedule='''squaredcos_cap_v2'''     ,       )

       # regular denoising components

       torch.manual_seed(0 )
       SCREAMING_SNAKE_CASE__   =   StableUnCLIPImageNormalizer(embedding_dim=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   DDPMScheduler(beta_schedule='''squaredcos_cap_v2''' )

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

       torch.manual_seed(0 )
       SCREAMING_SNAKE_CASE__   =   CLIPTextModel(
           CLIPTextConfig(
               bos_token_id=0     ,       eos_token_id=2     ,       hidden_size=__lowerCamelCase     ,       projection_dim=32     ,       intermediate_size=37     ,       layer_norm_eps=1e-05     ,       num_attention_heads=4     ,       num_hidden_layers=5     ,       pad_token_id=1     ,       vocab_size=1000     ,       ) )

       torch.manual_seed(0 )
       SCREAMING_SNAKE_CASE__   =   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=__lowerCamelCase     ,       layers_per_block=1     ,       upcast_attention=__lowerCamelCase     ,       use_linear_projection=__lowerCamelCase     ,       )

       torch.manual_seed(0 )
       SCREAMING_SNAKE_CASE__   =   DDIMScheduler(
           beta_schedule='''scaled_linear'''     ,       beta_start=0.00085     ,       beta_end=0.012     ,       prediction_type='''v_prediction'''     ,       set_alpha_to_one=__lowerCamelCase     ,       steps_offset=1     ,       )

       torch.manual_seed(0 )
       SCREAMING_SNAKE_CASE__   =   AutoencoderKL()

       SCREAMING_SNAKE_CASE__   =   {
           # 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 				lowercase_		(							self						:		List[str]     ,       __lowerCamelCase						:		List[str]     ,       __lowerCamelCase						:		Any=0 )		->							Any:
       if str(__lowerCamelCase ).startswith('''mps''' ):
            SCREAMING_SNAKE_CASE__   =   torch.manual_seed(__lowerCamelCase )
       else:
            SCREAMING_SNAKE_CASE__   =   torch.Generator(device=__lowerCamelCase ).manual_seed(__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   {
           '''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 				lowercase_		(							self						:		List[Any] )		->							List[str]:
       SCREAMING_SNAKE_CASE__   =   torch_device == '''cpu'''

       self._test_attention_slicing_forward_pass(test_max_difference=__lowerCamelCase )







  def 				lowercase_		(							self						:		Optional[Any] )		->							Optional[Any]:
       SCREAMING_SNAKE_CASE__   =   torch_device in ['''cpu''', '''mps''']

       self._test_inference_batch_single_identical(test_max_difference=__lowerCamelCase )






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






  """simple docstring"""


  def 				lowercase_		(							self						:		str )		->							List[str]:
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()






  def 				lowercase_		(							self						:		Optional[int] )		->							Union[str, Any]:
       SCREAMING_SNAKE_CASE__   =   load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy''' )

       SCREAMING_SNAKE_CASE__   =   StableUnCLIPPipeline.from_pretrained('''fusing/stable-unclip-2-1-l'''     ,       torch_dtype=torch.floataa )
       pipe.to(__lowerCamelCase )
       pipe.set_progress_bar_config(disable=__lowerCamelCase )
       # 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()

       SCREAMING_SNAKE_CASE__   =   torch.Generator(device='''cpu''' ).manual_seed(0 )
       SCREAMING_SNAKE_CASE__   =   pipe('''anime turle'''     ,       generator=__lowerCamelCase     ,       output_type='''np''' )

       SCREAMING_SNAKE_CASE__   =   output.images[0]

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

       assert_mean_pixel_difference(__lowerCamelCase     ,       __lowerCamelCase )







  def 				lowercase_		(							self						:		Dict )		->							List[str]:
       torch.cuda.empty_cache()
       torch.cuda.reset_max_memory_allocated()
       torch.cuda.reset_peak_memory_stats()

       SCREAMING_SNAKE_CASE__   =   StableUnCLIPPipeline.from_pretrained('''fusing/stable-unclip-2-1-l'''     ,       torch_dtype=torch.floataa )
       SCREAMING_SNAKE_CASE__   =   pipe.to(__lowerCamelCase )
       pipe.set_progress_bar_config(disable=__lowerCamelCase )
       pipe.enable_attention_slicing()
       pipe.enable_sequential_cpu_offload()

       SCREAMING_SNAKE_CASE__   =   pipe(
           '''anime turtle'''     ,       prior_num_inference_steps=2     ,       num_inference_steps=2     ,       output_type='''np'''     ,       )

       SCREAMING_SNAKE_CASE__   =   torch.cuda.max_memory_allocated()
       # make sure that less than 7 GB is allocated
       assert mem_bytes < 7 * 10**9





import warnings

from .generation import TFGenerationMixin






class 	UpperCAmelCase__		(	A__	):






  """simple docstring"""


  warnings.warn(
      "Importing `TFGenerationMixin` from `src/transformers/generation_tf_utils.py` is deprecated and will "
      "be removed in Transformers v5. Import as `from transformers import TFGenerationMixin` instead."						,					A__						,					)


