

'''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			):

  '''simple docstring'''





  __lowerCAmelCase       :			Optional[int]     	= StableUnCLIPPipeline
  __lowerCAmelCase       :			Dict     	= TEXT_TO_IMAGE_PARAMS
  __lowerCAmelCase       :			Dict     	= TEXT_TO_IMAGE_BATCH_PARAMS
  __lowerCAmelCase       :			Optional[int]     	= TEXT_TO_IMAGE_IMAGE_PARAMS
  __lowerCAmelCase       :			Any     	= TEXT_TO_IMAGE_IMAGE_PARAMS

  # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
  __lowerCAmelCase       :			List[str]     	= False
  def     snake_case_      (  self  )				->					str:


    """simple docstring"""
    UpperCAmelCase    				=						3_2
    UpperCAmelCase    				=						embedder_hidden_size

    # prior components

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

    torch.manual_seed(0  )
    UpperCAmelCase    				=						CLIPTextModelWithProjection(
        CLIPTextConfig(
            bos_token_id=0 ,			eos_token_id=2 ,			hidden_size=a_ ,			projection_dim=a_ ,			intermediate_size=3_7 ,			layer_norm_eps=1E-05 ,			num_attention_heads=4 ,			num_hidden_layers=5 ,			pad_token_id=1 ,			vocab_size=1_0_0_0 ,			)  )

    torch.manual_seed(0  )
    UpperCAmelCase    				=						PriorTransformer(
        num_attention_heads=2 ,			attention_head_dim=1_2 ,			embedding_dim=a_ ,			num_layers=1 ,			)

    torch.manual_seed(0  )
    UpperCAmelCase    				=						DDPMScheduler(
        variance_type='fixed_small_log' ,			prediction_type='sample' ,			num_train_timesteps=1_0_0_0 ,			clip_sample=a_ ,			clip_sample_range=5.0 ,			beta_schedule='squaredcos_cap_v2' ,			)

    # regular denoising components

    torch.manual_seed(0  )
    UpperCAmelCase    				=						StableUnCLIPImageNormalizer(embedding_dim=a_  )
    UpperCAmelCase    				=						DDPMScheduler(beta_schedule='squaredcos_cap_v2'  )

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

    torch.manual_seed(0  )
    UpperCAmelCase    				=						CLIPTextModel(
        CLIPTextConfig(
            bos_token_id=0 ,			eos_token_id=2 ,			hidden_size=a_ ,			projection_dim=3_2 ,			intermediate_size=3_7 ,			layer_norm_eps=1E-05 ,			num_attention_heads=4 ,			num_hidden_layers=5 ,			pad_token_id=1 ,			vocab_size=1_0_0_0 ,			)  )

    torch.manual_seed(0  )
    UpperCAmelCase    				=						UNetaDConditionModel(
        sample_size=3_2 ,			in_channels=4 ,			out_channels=4 ,			down_block_types=('CrossAttnDownBlock2D', 'DownBlock2D') ,			up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D') ,			block_out_channels=(3_2, 6_4) ,			attention_head_dim=(2, 4) ,			class_embed_type='projection' ,			projection_class_embeddings_input_dim=embedder_projection_dim * 2 ,			cross_attention_dim=a_ ,			layers_per_block=1 ,			upcast_attention=a_ ,			use_linear_projection=a_ ,			)

    torch.manual_seed(0  )
    UpperCAmelCase    				=						DDIMScheduler(
        beta_schedule='scaled_linear' ,			beta_start=0.0_0085 ,			beta_end=0.012 ,			prediction_type='v_prediction' ,			set_alpha_to_one=a_ ,			steps_offset=1 ,			)

    torch.manual_seed(0  )
    UpperCAmelCase    				=						AutoencoderKL()

    UpperCAmelCase    				=						{
        # 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 ,			a_ ,			a_=0  )				->					List[Any]:


    """simple docstring"""
    if str(a_  ).startswith('mps'  ):
      UpperCAmelCase    				=						torch.manual_seed(a_  )
    else:
      UpperCAmelCase    				=						torch.Generator(device=a_  ).manual_seed(a_  )
    UpperCAmelCase    				=						{
        '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  )				->					Any:


    """simple docstring"""
    UpperCAmelCase    				=						torch_device == 'cpu'

    self._test_attention_slicing_forward_pass(test_max_difference=a_  )





  def     snake_case_      (  self  )				->					List[str]:


    """simple docstring"""
    UpperCAmelCase    				=						torch_device in ['cpu', 'mps']

    self._test_inference_batch_single_identical(test_max_difference=a_  )




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

  '''simple docstring'''





  def     snake_case_      (  self  )				->					int:


    """simple docstring"""
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()
  def     snake_case_      (  self  )				->					Tuple:


    """simple docstring"""
    UpperCAmelCase    				=						load_numpy(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy'  )

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

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

    UpperCAmelCase    				=						output.images[0]

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

    assert_mean_pixel_difference(a_ ,			a_  )





  def     snake_case_      (  self  )				->					List[Any]:


    """simple docstring"""
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_peak_memory_stats()

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

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

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





'''simple docstring'''




from ..utils import DummyObject, requires_backends



class 							lowercase_						(		metaclass=a			):

  '''simple docstring'''





  __lowerCAmelCase       :			str     	= ["keras_nlp"]
  def __init__(  self ,			*a_ ,			**a_  )				->					Union[str, Any]:


    """simple docstring"""
    requires_backends(self ,			['keras_nlp']  )


