





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        _lowerCAmelCase    (			UpperCAmelCase_							,    UpperCAmelCase_							,    UpperCAmelCase_							,    unittest.TestCase							):


  '''simple docstring'''
  a_							:       Union[str, Any]      					=StableUnCLIPPipeline
  a_							:       int      					=TEXT_TO_IMAGE_PARAMS
  a_							:       str      					=TEXT_TO_IMAGE_BATCH_PARAMS
  a_							:       Optional[Any]      					=TEXT_TO_IMAGE_IMAGE_PARAMS
  a_							:       Dict      					=TEXT_TO_IMAGE_IMAGE_PARAMS

  # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
  a_							:       Optional[Any]      					=False







  def UpperCamelCase_ (      self      :					List[Any]							):


    '''simple docstring'''





    _snake_case   :				int														=	32
    _snake_case   :				Tuple														=	embedder_hidden_size

    # prior components

    torch.manual_seed(0							)
    _snake_case   :				Optional[Any]														=	CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'							)

    torch.manual_seed(0							)
    _snake_case   :				Union[str, Any]														=	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=10_00      ,				)							)

    torch.manual_seed(0							)
    _snake_case   :				str														=	PriorTransformer(
        num_attention_heads=2      ,				attention_head_dim=12      ,				embedding_dim=UpperCamelCase      ,				num_layers=1      ,				)

    torch.manual_seed(0							)
    _snake_case   :				Tuple														=	DDPMScheduler(
        variance_type='fixed_small_log'      ,				prediction_type='sample'      ,				num_train_timesteps=10_00      ,				clip_sample=UpperCamelCase      ,				clip_sample_range=5.0      ,				beta_schedule='squaredcos_cap_v2'      ,				)

    # regular denoising components

    torch.manual_seed(0							)
    _snake_case   :				int														=	StableUnCLIPImageNormalizer(embedding_dim=UpperCamelCase							)
    _snake_case   :				List[str]														=	DDPMScheduler(beta_schedule='squaredcos_cap_v2'							)

    torch.manual_seed(0							)
    _snake_case   :				List[Any]														=	CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'							)

    torch.manual_seed(0							)
    _snake_case   :				List[str]														=	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=10_00      ,				)							)

    torch.manual_seed(0							)
    _snake_case   :				int														=	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							)
    _snake_case   :				List[str]														=	DDIMScheduler(
        beta_schedule='scaled_linear'      ,				beta_start=0.0_00_85      ,				beta_end=0.0_12      ,				prediction_type='v_prediction'      ,				set_alpha_to_one=UpperCamelCase      ,				steps_offset=1      ,				)

    torch.manual_seed(0							)
    _snake_case   :				Tuple														=	AutoencoderKL()

    _snake_case   :				Optional[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 UpperCamelCase_ (      self      :					Dict      ,				UpperCamelCase      :					Tuple      ,				UpperCamelCase      :					Tuple=0							):


    '''simple docstring'''





    if str(UpperCamelCase							).startswith('mps'							):
      _snake_case   :				Union[str, Any]														=	torch.manual_seed(UpperCamelCase							)
    else:
      _snake_case   :				Tuple														=	torch.Generator(device=UpperCamelCase							).manual_seed(UpperCamelCase							)
    _snake_case   :				Any														=	{
        '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 UpperCamelCase_ (      self      :					List[Any]							):


    '''simple docstring'''





    _snake_case   :				int														=	torch_device == 'cpu'

    self._test_attention_slicing_forward_pass(test_max_difference=UpperCamelCase							)







  def UpperCamelCase_ (      self      :					int							):


    '''simple docstring'''





    _snake_case   :				str														=	torch_device in ['cpu', 'mps']

    self._test_inference_batch_single_identical(test_max_difference=UpperCamelCase							)




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


  '''simple docstring'''







  def UpperCamelCase_ (      self      :					List[str]							):


    '''simple docstring'''





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







  def UpperCamelCase_ (      self      :					Tuple							):


    '''simple docstring'''





    _snake_case   :				Optional[int]														=	load_numpy(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy'							)

    _snake_case   :				int														=	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()

    _snake_case   :				Dict														=	torch.Generator(device='cpu'							).manual_seed(0							)
    _snake_case   :				Tuple														=	pipe('anime turle'      ,				generator=UpperCamelCase      ,				output_type='np'							)

    _snake_case   :				List[str]														=	output.images[0]

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

    assert_mean_pixel_difference(UpperCamelCase      ,				UpperCamelCase							)







  def UpperCamelCase_ (      self      :					str							):


    '''simple docstring'''





    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_peak_memory_stats()

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

    _snake_case   :				Union[str, Any]														=	pipe(
        'anime turtle'      ,				prior_num_inference_steps=2      ,				num_inference_steps=2      ,				output_type='np'      ,				)

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











from unittest import TestCase

from datasets import Dataset
from minhash_deduplication import deduplicate_dataset, make_duplicate_clusters


def     lowerCamelCase_	(      )->		Any:
  _snake_case   :				List[str]														=	{
      'repo_name': ['test_repo1', 'test_repo2', 'test_repo3'],
      'path': ['test_1.py', 'test_2.py', 'unit_test.py'],
      'content': ['a ' * 20, 'a ' * 30, 'b ' * 7],
  }
  _snake_case   :				Optional[Any]														=	Dataset.from_dict(lowerCAmelCase				)
  return dataset







class        _lowerCAmelCase    (			UpperCAmelCase_							):


  '''simple docstring'''







  def UpperCamelCase_ (      self      :					Dict							):


    '''simple docstring'''





    _snake_case   :				Union[str, Any]														=	get_dataset()
    _snake_case   :				Tuple														=	make_duplicate_clusters(UpperCamelCase      ,				0.85							)
    self.assertEqual(len(duplicate_clusters[0]							)      ,				2							)







  def UpperCamelCase_ (      self      :					Dict							):


    '''simple docstring'''





    _snake_case   :				List[str]														=	get_dataset()
    _snake_case							,					_snake_case   :				str														=	deduplicate_dataset(UpperCamelCase							)
    self.assertEqual(len(UpperCamelCase							)      ,				2							)
    print(UpperCamelCase							)
    self.assertEqual(duplicate_clusters[0][0]['copies']      ,				2							)
    self.assertEqual(duplicate_clusters[0][0]['is_extreme']      ,				UpperCamelCase							)




