

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 			lowercase    (	lowercase_     ,  lowercase_     ,  lowercase_     ,  unittest.TestCase      ):
				__SCREAMING_SNAKE_CASE :			List[Any]	      =      StableUnCLIPImgaImgPipeline
				__SCREAMING_SNAKE_CASE :			Any	      =      TEXT_GUIDED_IMAGE_VARIATION_PARAMS
				__SCREAMING_SNAKE_CASE :			Union[str, Any]	      =      TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
				__SCREAMING_SNAKE_CASE :			Tuple	      =      frozenset(
				    []      )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
				__SCREAMING_SNAKE_CASE :			List[str]	      =      frozenset([]      )




				def    a							( self	):
											snake_case_          =   32
											snake_case_          =   embedder_hidden_size

											# image encoding components

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

											torch.manual_seed(0	)
											snake_case_          =   CLIPVisionModelWithProjection(
											    CLIPVisionConfig(
											        hidden_size=snake_case				,       projection_dim=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	)
											snake_case_          =   StableUnCLIPImageNormalizer(embedding_dim=snake_case	)
											snake_case_          =   DDPMScheduler(beta_schedule='squaredcos_cap_v2'	)

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

											torch.manual_seed(0	)
											snake_case_          =   CLIPTextModel(
											    CLIPTextConfig(
											        bos_token_id=0				,       eos_token_id=2				,       hidden_size=snake_case				,       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=1000				,       )	)

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

											torch.manual_seed(0	)
											snake_case_          =   DDIMScheduler(
											    beta_schedule='scaled_linear'				,       beta_start=0.0_00_85				,       beta_end=0.0_12				,       prediction_type='v_prediction'				,       set_alpha_to_one=snake_case				,       steps_offset=1				,       )

											torch.manual_seed(0	)
											snake_case_          =   AutoencoderKL()

											snake_case_          =   {
											    # 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				,       snake_case				,       snake_case=0				,       snake_case=True	):
											if str(snake_case	).startswith('mps'	):
																		snake_case_          =   torch.manual_seed(snake_case	)
											else:
																		snake_case_          =   torch.Generator(device=snake_case	).manual_seed(snake_case	)

											snake_case_          =   floats_tensor((1, 3, 32, 32)				,       rng=random.Random(snake_case	)	).to(snake_case	)

											if pil_image:
																		snake_case_          =   input_image * 0.5 + 0.5
																		snake_case_          =   input_image.clamp(0				,       1	)
																		snake_case_          =   input_image.cpu().permute(0				,       2				,       3				,       1	).float().numpy()
																		snake_case_          =   DiffusionPipeline.numpy_to_pil(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	):
											snake_case_          =   'cpu'  # ensure determinism for the device-dependent torch.Generator
											snake_case_          =   self.get_dummy_components()
											snake_case_          =   StableUnCLIPImgaImgPipeline(**snake_case	)
											snake_case_          =   sd_pipe.to(snake_case	)
											sd_pipe.set_progress_bar_config(disable=snake_case	)

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

											assert image.shape == (1, 32, 32, 3)
											snake_case_          =   np.array([0.38_72, 0.72_24, 0.56_01, 0.47_41, 0.68_72, 0.58_14, 0.46_36, 0.38_67, 0.50_78]	)

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




				def    a							( self	):
											snake_case_          =   torch_device in ['cpu', 'mps']

											self._test_attention_slicing_forward_pass(test_max_difference=snake_case	)




				def    a							( self	):
											snake_case_          =   torch_device in ['cpu', 'mps']

											self._test_inference_batch_single_identical(test_max_difference=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	):
											self._test_xformers_attention_forwardGenerator_pass(test_max_difference=snake_case	)







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




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




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

											snake_case_          =   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'	)

											snake_case_          =   StableUnCLIPImgaImgPipeline.from_pretrained(
											    'fusing/stable-unclip-2-1-l-img2img'				,       torch_dtype=torch.floataa	)
											pipe.to(snake_case	)
											pipe.set_progress_bar_config(disable=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()

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

											snake_case_          =   output.images[0]

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

											assert_mean_pixel_difference(snake_case				,       snake_case	)




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

											snake_case_          =   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'	)

											snake_case_          =   StableUnCLIPImgaImgPipeline.from_pretrained(
											    'fusing/stable-unclip-2-1-h-img2img'				,       torch_dtype=torch.floataa	)
											pipe.to(snake_case	)
											pipe.set_progress_bar_config(disable=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()

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

											snake_case_          =   output.images[0]

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

											assert_mean_pixel_difference(snake_case				,       snake_case	)




				def    a							( self	):
											snake_case_          =   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()

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

											snake_case_          =   pipe(
											    snake_case				,       'anime turtle'				,       num_inference_steps=2				,       output_type='np'				,       )

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







import unittest

from transformers import is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_vision, slow, torch_device


if is_torch_available():
					import torch

					from transformers import AutoModelForImageClassification

if is_vision_available():
					from transformers import AutoImageProcessor







@require_torch
@require_vision
class 			lowercase    (	unittest.TestCase      ):




				@slow
				def    a							( self	):
											snake_case_          =   AutoImageProcessor.from_pretrained('microsoft/dit-base-finetuned-rvlcdip'	)
											snake_case_          =   AutoModelForImageClassification.from_pretrained('microsoft/dit-base-finetuned-rvlcdip'	)
											model.to(snake_case	)

											from datasets import load_dataset

											snake_case_          =   load_dataset('nielsr/rvlcdip-demo'	)

											snake_case_          =   dataset['train'][0]['image'].convert('RGB'	)

											snake_case_          =   image_processor(snake_case				,       return_tensors='pt'	).to(snake_case	)

											# forward pass
											with torch.no_grad():
																		snake_case_          =   model(**snake_case	)
																		snake_case_          =   outputs.logits

											snake_case_          =   torch.Size((1, 16)	)
											self.assertEqual(logits.shape				,       snake_case	)

											snake_case_          =   torch.tensor(
											    [-0.41_58, -0.40_92, -0.43_47]				,       device=snake_case				,       dtype=torch.float				,       )
											self.assertTrue(torch.allclose(logits[0, :3]				,       snake_case				,       atol=1e-4	)	)




