





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 					__snake_case	(  __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     unittest.TestCase    ):


    '''simple docstring'''







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



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			32
        SCREAMING_SNAKE_CASE__											=			embedder_hidden_size

        # image encoding components

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

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

        # regular denoising components

        torch.manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			StableUnCLIPImageNormalizer(embedding_dim=A_     )
        SCREAMING_SNAKE_CASE__											=			DDPMScheduler(beta_schedule='''squaredcos_cap_v2'''     )

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

        torch.manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			CLIPTextModel(
            CLIPTextConfig(
                bos_token_id=0   ,				eos_token_id=2   ,				hidden_size=A_   ,				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=10_00   ,				)     )

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

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

        torch.manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			AutoencoderKL()

        SCREAMING_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  lowercase_					(       self   ,				A_   ,				A_=0   ,				A_=True     ):





        '''simple docstring'''

        if str(A_     ).startswith('''mps'''     ):
            SCREAMING_SNAKE_CASE__											=			torch.manual_seed(A_     )
        else:
            SCREAMING_SNAKE_CASE__											=			torch.Generator(device=A_     ).manual_seed(A_     )

        SCREAMING_SNAKE_CASE__											=			floats_tensor((1, 3, 32, 32)   ,				rng=random.Random(A_     )     ).to(A_     )

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

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



    @skip_mps
    def  lowercase_					(       self     ):





        '''simple docstring'''

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

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

        assert image.shape == (1, 32, 32, 3)
        SCREAMING_SNAKE_CASE__											=			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  lowercase_					(       self     ):





        '''simple docstring'''

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

        self._test_attention_slicing_forward_pass(test_max_difference=A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

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

        self._test_inference_batch_single_identical(test_max_difference=A_     )




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





        '''simple docstring'''

        self._test_xformers_attention_forwardGenerator_pass(test_max_difference=A_     )





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


    '''simple docstring'''







    def  lowercase_					(       self     ):





        '''simple docstring'''

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



    def  lowercase_					(       self     ):





        '''simple docstring'''

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

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

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

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

        SCREAMING_SNAKE_CASE__											=			output.images[0]

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

        assert_mean_pixel_difference(A_   ,				A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

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

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

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

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

        SCREAMING_SNAKE_CASE__											=			output.images[0]

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

        assert_mean_pixel_difference(A_   ,				A_     )




    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_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()

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

        SCREAMING_SNAKE_CASE__											=			pipe(
            A_   ,				'''anime turtle'''   ,				num_inference_steps=2   ,				output_type='''np'''   ,				)

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







'''simple docstring'''




import unittest

from transformers.testing_utils import require_bsa
from transformers.utils import is_bsa_available

from ...test_feature_extraction_common import FeatureExtractionSavingTestMixin


if is_bsa_available():
							from transformers import MarkupLMFeatureExtractor




class      UpperCamelCase__					(					unittest.TestCase			):




					"""simple docstring"""


					def __init__(						self				:			str      ,						lowerCamelCase_				:			str						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				List[str]		=			parent



					def lowerCamelCase_			(						self				:			Union[str, Any]						):

							'''simple docstring'''



							return {}







def 		__A		(  ):


		"""simple docstring"""



		SCREAMING_SNAKE_CASE :				Any		=			"""<HTML>

    <HEAD>
    <TITLE>sample document</TITLE>
    </HEAD>

    <BODY BGCOLOR=\"FFFFFF\">
    <HR>
    <a href=\"http://google.com\">Goog</a>
    <H1>This is one header</H1>
    <H2>This is a another Header</H2>
    <P>Travel from
        <P>
        <B>SFO to JFK</B>
        <BR>
        <B><I>on May 2, 2015 at 2:00 pm. For details go to confirm.com </I></B>
        <HR>
        <div style=\"color:#0000FF\">
            <h3>Traveler <b> name </b> is
            <p> John Doe </p>
        </div>"""

		SCREAMING_SNAKE_CASE :				int		=			"""
    <!DOCTYPE html>
    <html>
    <body>

    <h1>My First Heading</h1>
    <p>My first paragraph.</p>

    </body>
    </html>
    """

		return [html_string_a, html_string_a]




@require_bsa
class      UpperCamelCase__					(					lowercase_						,					unittest.TestCase			):




					"""simple docstring"""
					SCREAMING_SNAKE_CASE__				=	MarkupLMFeatureExtractor if is_bsa_available() else None


					def lowerCamelCase_			(						self				:			Optional[int]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				int		=			MarkupLMFeatureExtractionTester(self						)


					@property
					def lowerCamelCase_			(						self				:			int						):

							'''simple docstring'''



							return self.feature_extract_tester.prepare_feat_extract_dict()



					def lowerCamelCase_			(						self				:			int						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			self.feature_extraction_class()

							# Test not batched input
							SCREAMING_SNAKE_CASE :				Optional[int]		=			get_html_strings()[0]
							SCREAMING_SNAKE_CASE :				List[Any]		=			feature_extractor(lowerCamelCase_						)

							# fmt: off
							SCREAMING_SNAKE_CASE :				Optional[int]		=			[["""sample document""", """Goog""", """This is one header""", """This is a another Header""", """Travel from""", """SFO to JFK""", """on May 2, 2015 at 2:00 pm. For details go to confirm.com""", """Traveler""", """name""", """is""", """John Doe"""]]
							SCREAMING_SNAKE_CASE :				List[str]		=			[["""/html/head/title""", """/html/body/a""", """/html/body/h1""", """/html/body/h2""", """/html/body/p""", """/html/body/p/p/b[1]""", """/html/body/p/p/b[2]/i""", """/html/body/p/p/div/h3""", """/html/body/p/p/div/h3/b""", """/html/body/p/p/div/h3""", """/html/body/p/p/div/h3/p"""]]
							# fmt: on

							self.assertEqual(encoding.nodes      ,						lowerCamelCase_						)
							self.assertEqual(encoding.xpaths      ,						lowerCamelCase_						)

							# Test batched
							SCREAMING_SNAKE_CASE :				str		=			get_html_strings()
							SCREAMING_SNAKE_CASE :				str		=			feature_extractor(lowerCamelCase_						)

							# fmt: off
							SCREAMING_SNAKE_CASE :				int		=			expected_nodes + [["""My First Heading""", """My first paragraph."""]]
							SCREAMING_SNAKE_CASE :				List[Any]		=			expected_xpaths + [["""/html/body/h1""", """/html/body/p"""]]

							self.assertEqual(len(encoding.nodes						)      ,						2						)
							self.assertEqual(len(encoding.xpaths						)      ,						2						)

							self.assertEqual(encoding.nodes      ,						lowerCamelCase_						)
							self.assertEqual(encoding.xpaths      ,						lowerCamelCase_						)



