






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




							def 						__snake_case					(   self     :   int	):
											lowerCAmelCase__      			=						32
											lowerCAmelCase__      			=						embedder_hidden_size

											# image encoding components

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

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

											# regular denoising components

											torch.manual_seed(0	)
											lowerCAmelCase__      			=						StableUnCLIPImageNormalizer(embedding_dim=lowerCAmelCase_	)
											lowerCAmelCase__      			=						DDPMScheduler(beta_schedule='''squaredcos_cap_v2'''	)

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

											torch.manual_seed(0	)
											lowerCAmelCase__      			=						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=1_000 ,							)	)

											torch.manual_seed(0	)
											lowerCAmelCase__      			=						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	)
											lowerCAmelCase__      			=						DDIMScheduler(
											    beta_schedule='''scaled_linear''' ,							beta_start=0.00_085 ,							beta_end=0.012 ,							prediction_type='''v_prediction''' ,							set_alpha_to_one=lowerCAmelCase_ ,							steps_offset=1 ,							)

											torch.manual_seed(0	)
											lowerCAmelCase__      			=						AutoencoderKL()

											lowerCAmelCase__      			=						{
											    # 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 						__snake_case					(   self     :   str ,							SCREAMING_SNAKE_CASE_     :   Optional[int] ,							SCREAMING_SNAKE_CASE_     :   List[str]=0 ,							SCREAMING_SNAKE_CASE_     :   str=True	):







											if str(lowerCAmelCase_	).startswith('''mps'''	):
															lowerCAmelCase__      			=						torch.manual_seed(lowerCAmelCase_	)
											else:
															lowerCAmelCase__      			=						torch.Generator(device=lowerCAmelCase_	).manual_seed(lowerCAmelCase_	)

											lowerCAmelCase__      			=						floats_tensor((1, 3, 32, 32) ,							rng=random.Random(lowerCAmelCase_	)	).to(lowerCAmelCase_	)

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

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




							@skip_mps
							def 						__snake_case					(   self     :   str	):
											lowerCAmelCase__      			=						'''cpu'''  # ensure determinism for the device-dependent torch.Generator
											lowerCAmelCase__      			=						self.get_dummy_components()
											lowerCAmelCase__      			=						StableUnCLIPImgaImgPipeline(**lowerCAmelCase_	)
											lowerCAmelCase__      			=						sd_pipe.to(lowerCAmelCase_	)
											sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_	)

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

											assert image.shape == (1, 32, 32, 3)
											lowerCAmelCase__      			=						np.array([0.3_872, 0.7_224, 0.5_601, 0.4_741, 0.6_872, 0.5_814, 0.4_636, 0.3_867, 0.5_078]	)

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




							def 						__snake_case					(   self     :   Union[str, Any]	):
											lowerCAmelCase__      			=						torch_device in ['''cpu''', '''mps''']

											self._test_attention_slicing_forward_pass(test_max_difference=lowerCAmelCase_	)




							def 						__snake_case					(   self     :   int	):
											lowerCAmelCase__      			=						torch_device in ['''cpu''', '''mps''']

											self._test_inference_batch_single_identical(test_max_difference=lowerCAmelCase_	)




							@unittest.skipIf(
							    torch_device != '''cuda''' or not is_xformers_available() ,							reason='''XFormers attention is only available with CUDA and `xformers` installed''' ,							)
							def 						__snake_case					(   self     :   int	):
											self._test_xformers_attention_forwardGenerator_pass(test_max_difference=lowerCAmelCase_	)



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




							def 						__snake_case					(   self     :   Tuple	):
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()




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

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

											lowerCAmelCase__      			=						StableUnCLIPImgaImgPipeline.from_pretrained(
											    '''fusing/stable-unclip-2-1-l-img2img''' ,							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()

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

											lowerCAmelCase__      			=						output.images[0]

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

											assert_mean_pixel_difference(lowerCAmelCase_ ,							lowerCAmelCase_	)




							def 						__snake_case					(   self     :   List[str]	):
											lowerCAmelCase__      			=						load_image(
											    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'''	)

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

											lowerCAmelCase__      			=						StableUnCLIPImgaImgPipeline.from_pretrained(
											    '''fusing/stable-unclip-2-1-h-img2img''' ,							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()

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

											lowerCAmelCase__      			=						output.images[0]

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

											assert_mean_pixel_difference(lowerCAmelCase_ ,							lowerCAmelCase_	)




							def 						__snake_case					(   self     :   Dict	):
											lowerCAmelCase__      			=						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()

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

											lowerCAmelCase__      			=						pipe(
											    lowerCAmelCase_ ,							'''anime turtle''' ,							num_inference_steps=2 ,							output_type='''np''' ,							)

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













from typing import TYPE_CHECKING

from ....utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_UpperCAmelCase			:		Optional[Any]   				=    {
    "configuration_mctct": ["MCTCT_PRETRAINED_CONFIG_ARCHIVE_MAP", "MCTCTConfig"],
    "feature_extraction_mctct": ["MCTCTFeatureExtractor"],
    "processing_mctct": ["MCTCTProcessor"],
}


try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_UpperCAmelCase			:		str   				=    [
			    "MCTCT_PRETRAINED_MODEL_ARCHIVE_LIST",
			    "MCTCTForCTC",
			    "MCTCTModel",
			    "MCTCTPreTrainedModel",
			]


if TYPE_CHECKING:
			from .configuration_mctct import MCTCT_PRETRAINED_CONFIG_ARCHIVE_MAP, MCTCTConfig
			from .feature_extraction_mctct import MCTCTFeatureExtractor
			from .processing_mctct import MCTCTProcessor

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_mctct import MCTCT_PRETRAINED_MODEL_ARCHIVE_LIST, MCTCTForCTC, MCTCTModel, MCTCTPreTrainedModel

else:
			import sys

			_UpperCAmelCase			:		List[str]   				=    _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





