import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModel,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableUnCLIPImgaImgPipeline, UNetaDConditionModel
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    floats_tensor,
    load_image,
    load_numpy,
    require_torch_gpu,
    skip_mps,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()



class     _A   (    UpperCamelCase_	, UpperCamelCase_	, UpperCamelCase_	, unittest.TestCase							):

  '''simple docstring'''





  SCREAMING_SNAKE_CASE_	:	str						=StableUnCLIPImgaImgPipeline
  SCREAMING_SNAKE_CASE_	:	Optional[int]						=TEXT_GUIDED_IMAGE_VARIATION_PARAMS
  SCREAMING_SNAKE_CASE_	:	int						=TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
  SCREAMING_SNAKE_CASE_	:	int						=frozenset(
      []							)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
  SCREAMING_SNAKE_CASE_	:	List[Any]						=frozenset([]							)






  def 		_a					(self	)		->      Dict:





       '''simple docstring'''





       UpperCamelCase__										=      32
       UpperCamelCase__										=      embedder_hidden_size

       # image encoding components

       UpperCamelCase__										=      CLIPImageProcessor(crop_size=32   ,				size=32	)

       torch.manual_seed(0	)
       UpperCamelCase__										=      CLIPVisionModelWithProjection(
           CLIPVisionConfig(
               hidden_size=SCREAMING_SNAKE_CASE_   ,				projection_dim=SCREAMING_SNAKE_CASE_   ,				num_hidden_layers=5   ,				num_attention_heads=4   ,				image_size=32   ,				intermediate_size=37   ,				patch_size=1   ,				)	)

       # regular denoising components

       torch.manual_seed(0	)
       UpperCamelCase__										=      StableUnCLIPImageNormalizer(embedding_dim=SCREAMING_SNAKE_CASE_	)
       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=SCREAMING_SNAKE_CASE_   ,				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	)
       UpperCamelCase__										=      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=SCREAMING_SNAKE_CASE_   ,				layers_per_block=1   ,				upcast_attention=SCREAMING_SNAKE_CASE_   ,				use_linear_projection=SCREAMING_SNAKE_CASE_   ,				)

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

       torch.manual_seed(0	)
       UpperCamelCase__										=      AutoencoderKL()

       UpperCamelCase__										=      {
           # image encoding components
           '''feature_extractor''': feature_extractor,
           '''image_encoder''': image_encoder.eval(),
           # image noising components
           '''image_normalizer''': image_normalizer.eval(),
           '''image_noising_scheduler''': image_noising_scheduler,
           # regular denoising components
           '''tokenizer''': tokenizer,
           '''text_encoder''': text_encoder.eval(),
           '''unet''': unet.eval(),
           '''scheduler''': scheduler,
           '''vae''': vae.eval(),
       }

       return components






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_=0   ,				SCREAMING_SNAKE_CASE_=True	)		->      Tuple:





       '''simple docstring'''





       if str(SCREAMING_SNAKE_CASE_	).startswith('''mps'''	):
            UpperCamelCase__										=      torch.manual_seed(SCREAMING_SNAKE_CASE_	)
       else:
            UpperCamelCase__										=      torch.Generator(device=SCREAMING_SNAKE_CASE_	).manual_seed(SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      floats_tensor((1, 3, 32, 32)   ,				rng=random.Random(SCREAMING_SNAKE_CASE_	)	).to(SCREAMING_SNAKE_CASE_	)

       if pil_image:
            UpperCamelCase__										=      input_image * 0.5 + 0.5
            UpperCamelCase__										=      input_image.clamp(0   ,				1	)
            UpperCamelCase__										=      input_image.cpu().permute(0   ,				2   ,				3   ,				1	).float().numpy()
            UpperCamelCase__										=      DiffusionPipeline.numpy_to_pil(SCREAMING_SNAKE_CASE_	)[0]

       return {
           "prompt": "An anime racoon running a marathon",
           "image": input_image,
           "generator": generator,
           "num_inference_steps": 2,
           "output_type": "np",
       }






  @skip_mps
  def 		_a					(self	)		->      Optional[Any]:





       '''simple docstring'''





       UpperCamelCase__										=      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
       UpperCamelCase__										=      self.get_dummy_components()
       UpperCamelCase__										=      StableUnCLIPImgaImgPipeline(**SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      sd_pipe.to(SCREAMING_SNAKE_CASE_	)
       sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      self.get_dummy_inputs(SCREAMING_SNAKE_CASE_	)
       inputs.update({'''image_embeds''': None}	)
       UpperCamelCase__										=      sd_pipe(**SCREAMING_SNAKE_CASE_	).images
       UpperCamelCase__										=      image[0, -3:, -3:, -1]

       assert image.shape == (1, 32, 32, 3)
       UpperCamelCase__										=      np.array([0.3872, 0.7224, 0.5601, 0.4741, 0.6872, 0.5814, 0.4636, 0.3867, 0.5078]	)

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






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





       '''simple docstring'''





       UpperCamelCase__										=      torch_device in ['''cpu''', '''mps''']

       self._test_attention_slicing_forward_pass(test_max_difference=SCREAMING_SNAKE_CASE_	)






  def 		_a					(self	)		->      Tuple:





       '''simple docstring'''





       UpperCamelCase__										=      torch_device in ['''cpu''', '''mps''']

       self._test_inference_batch_single_identical(test_max_difference=SCREAMING_SNAKE_CASE_	)






  @unittest.skipIf(
      torch_device != '''cuda''' or not is_xformers_available()   ,				reason='''XFormers attention is only available with CUDA and `xformers` installed'''   ,				)
  def 		_a					(self	)		->      Any:





       '''simple docstring'''





       self._test_xformers_attention_forwardGenerator_pass(test_max_difference=SCREAMING_SNAKE_CASE_	)



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

  '''simple docstring'''





  def 		_a					(self	)		->      Optional[Any]:





       '''simple docstring'''





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






  def 		_a					(self	)		->      str:





       '''simple docstring'''





       UpperCamelCase__										=      load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'''	)

       UpperCamelCase__										=      load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_img2img_anime_turtle_fp16.npy'''	)

       UpperCamelCase__										=      StableUnCLIPImgaImgPipeline.from_pretrained(
           '''fusing/stable-unclip-2-1-l-img2img'''   ,				torch_dtype=torch.floataa	)
       pipe.to(SCREAMING_SNAKE_CASE_	)
       pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)
       # 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(SCREAMING_SNAKE_CASE_   ,				'''anime turle'''   ,				generator=SCREAMING_SNAKE_CASE_   ,				output_type='''np'''	)

       UpperCamelCase__										=      output.images[0]

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

       assert_mean_pixel_difference(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)






  def 		_a					(self	)		->      Optional[int]:





       '''simple docstring'''





       UpperCamelCase__										=      load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'''	)

       UpperCamelCase__										=      load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_h_img2img_anime_turtle_fp16.npy'''	)

       UpperCamelCase__										=      StableUnCLIPImgaImgPipeline.from_pretrained(
           '''fusing/stable-unclip-2-1-h-img2img'''   ,				torch_dtype=torch.floataa	)
       pipe.to(SCREAMING_SNAKE_CASE_	)
       pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)
       # 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(SCREAMING_SNAKE_CASE_   ,				'''anime turle'''   ,				generator=SCREAMING_SNAKE_CASE_   ,				output_type='''np'''	)

       UpperCamelCase__										=      output.images[0]

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

       assert_mean_pixel_difference(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)






  def 		_a					(self	)		->      Any:





       '''simple docstring'''





       UpperCamelCase__										=      load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'''	)

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

       UpperCamelCase__										=      StableUnCLIPImgaImgPipeline.from_pretrained(
           '''fusing/stable-unclip-2-1-h-img2img'''   ,				torch_dtype=torch.floataa	)
       UpperCamelCase__										=      pipe.to(SCREAMING_SNAKE_CASE_	)
       pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)
       pipe.enable_attention_slicing()
       pipe.enable_sequential_cpu_offload()

       UpperCamelCase__										=      pipe(
           SCREAMING_SNAKE_CASE_   ,				'''anime turtle'''   ,				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 * 10**9




import torch
from torch import nn
from transformers import CLIPPreTrainedModel, CLIPVisionModel

from ...models.attention import BasicTransformerBlock
from ...utils import logging


__magic_name__								=logging.get_logger(__name__)  # pylint: disable=invalid-name



class     _A   (    __UpperCamelCase							):






  def __init__(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_=768	)		->      Dict:





       '''simple docstring'''





       super().__init__(SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      proj_size

       UpperCamelCase__										=      CLIPVisionModel(SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      PaintByExampleMapper(SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      nn.LayerNorm(config.hidden_size	)
       UpperCamelCase__										=      nn.Linear(config.hidden_size   ,				self.proj_size	)

       # uncondition for scaling
       UpperCamelCase__										=      nn.Parameter(torch.randn((1, 1, self.proj_size)	)	)






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_=False	)		->      List[str]:





       '''simple docstring'''





       UpperCamelCase__										=      self.model(pixel_values=SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      clip_output.pooler_output
       UpperCamelCase__										=      self.mapper(latent_states[:, None]	)
       UpperCamelCase__										=      self.final_layer_norm(SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      self.proj_out(SCREAMING_SNAKE_CASE_	)
       if return_uncond_vector:
            return latent_states, self.uncond_vector

       return latent_states



class     _A   (    nn.Module							):






  def __init__(self   ,				SCREAMING_SNAKE_CASE_	)		->      int:





       '''simple docstring'''





       super().__init__()
       UpperCamelCase__										=      (config.num_hidden_layers + 1) // 5
       UpperCamelCase__										=      config.hidden_size
       UpperCamelCase__										=      1
       UpperCamelCase__										=      nn.ModuleList(
           [
               BasicTransformerBlock(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				activation_fn='''gelu'''   ,				attention_bias=SCREAMING_SNAKE_CASE_	)
               for _ in range(SCREAMING_SNAKE_CASE_	)
           ]	)






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_	)		->      str:





       '''simple docstring'''





       for block in self.blocks:
            UpperCamelCase__										=      block(SCREAMING_SNAKE_CASE_	)

       return hidden_states



