





import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()





class 	lowerCAmelCase__(       _a     ,   _a     ,   unittest.TestCase			):






       '''simple docstring'''


       __snake_case       	=			StableDiffusionXLPipeline
       __snake_case       	=			TEXT_TO_IMAGE_PARAMS
       __snake_case       	=			TEXT_TO_IMAGE_BATCH_PARAMS
       __snake_case       	=			TEXT_TO_IMAGE_IMAGE_PARAMS
       __snake_case       	=			TEXT_TO_IMAGE_IMAGE_PARAMS




       def 			UpperCamelCase_	(					self							)	->   List[Any]:
            torch.manual_seed(0							)
            _SCREAMING_SNAKE_CASE				:		str  			=							UNetaDConditionModel(
                block_out_channels=(3_2, 6_4)			,  layers_per_block=2			,  sample_size=3_2			,  in_channels=4			,  out_channels=4			,  down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")			,  up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")			,  attention_head_dim=(2, 4)			,  use_linear_projection=lowerCAmelCase_			,  addition_embed_type="text_time"			,  addition_time_embed_dim=8			,  transformer_layers_per_block=(1, 2)			,  projection_class_embeddings_input_dim=8_0			,  cross_attention_dim=6_4			,  )
            _SCREAMING_SNAKE_CASE				:		Tuple  			=							EulerDiscreteScheduler(
                beta_start=0.0_0085			,  beta_end=0.012			,  steps_offset=1			,  beta_schedule="scaled_linear"			,  timestep_spacing="leading"			,  )
            torch.manual_seed(0							)
            _SCREAMING_SNAKE_CASE				:		Dict  			=							AutoencoderKL(
                block_out_channels=[3_2, 6_4]			,  in_channels=3			,  out_channels=3			,  down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]			,  up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]			,  latent_channels=4			,  sample_size=1_2_8			,  )
            torch.manual_seed(0							)
            _SCREAMING_SNAKE_CASE				:		Any  			=							CLIPTextConfig(
                bos_token_id=0			,  eos_token_id=2			,  hidden_size=3_2			,  intermediate_size=3_7			,  layer_norm_eps=1E-05			,  num_attention_heads=4			,  num_hidden_layers=5			,  pad_token_id=1			,  vocab_size=1_0_0_0			,  hidden_act="gelu"			,  projection_dim=3_2			,  )
            _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							CLIPTextModel(lowerCAmelCase_							)
            _SCREAMING_SNAKE_CASE				:		List[str]  			=							CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"			,  local_files_only=lowerCAmelCase_							)

            _SCREAMING_SNAKE_CASE				:		str  			=							CLIPTextModelWithProjection(lowerCAmelCase_							)
            _SCREAMING_SNAKE_CASE				:		Union[str, Any]  			=							CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"			,  local_files_only=lowerCAmelCase_							)

            _SCREAMING_SNAKE_CASE				:		int  			=							{
                "unet": unet,
                "scheduler": scheduler,
                "vae": vae,
                "text_encoder": text_encoder,
                "tokenizer": tokenizer,
                "text_encoder_2": text_encoder_a,
                "tokenizer_2": tokenizer_a,
                # "safety_checker": None,
                # "feature_extractor": None,
            }
            return components




       def 			UpperCamelCase_	(					self			,  __lowerCamelCase			,  __lowerCamelCase=0							)	->   List[str]:


            if str(lowerCAmelCase_							).startswith("mps"							):
                 _SCREAMING_SNAKE_CASE				:		int  			=							torch.manual_seed(lowerCAmelCase_							)
            else:
                 _SCREAMING_SNAKE_CASE				:		List[Any]  			=							torch.Generator(device=lowerCAmelCase_							).manual_seed(lowerCAmelCase_							)
            _SCREAMING_SNAKE_CASE				:		int  			=							{
                "prompt": "A painting of a squirrel eating a burger",
                "generator": generator,
                "num_inference_steps": 2,
                "guidance_scale": 5.0,
                "output_type": "numpy",
            }
            return inputs




       def 			UpperCamelCase_	(					self							)	->   Any:
            _SCREAMING_SNAKE_CASE				:		List[str]  			=							"cpu"  # ensure determinism for the device-dependent torch.Generator
            _SCREAMING_SNAKE_CASE				:		List[Any]  			=							self.get_dummy_components()
            _SCREAMING_SNAKE_CASE				:		List[str]  			=							StableDiffusionXLPipeline(**lowerCAmelCase_							)
            _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							sd_pipe.to(lowerCAmelCase_							)
            sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

            _SCREAMING_SNAKE_CASE				:		Any  			=							self.get_dummy_inputs(lowerCAmelCase_							)
            _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							sd_pipe(**lowerCAmelCase_							).images
            _SCREAMING_SNAKE_CASE				:		int  			=							image[0, -3:, -3:, -1]

            assert image.shape == (1, 6_4, 6_4, 3)
            _SCREAMING_SNAKE_CASE				:		List[str]  			=							np.array([0.5873, 0.6128, 0.4797, 0.5122, 0.5674, 0.4639, 0.5227, 0.5149, 0.4747]							)

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




       def 			UpperCamelCase_	(					self							)	->   Any:
            _SCREAMING_SNAKE_CASE				:		Union[str, Any]  			=							self.get_dummy_components()
            _SCREAMING_SNAKE_CASE				:		Dict  			=							StableDiffusionXLPipeline(**lowerCAmelCase_							)
            _SCREAMING_SNAKE_CASE				:		Tuple  			=							sd_pipe.to(lowerCAmelCase_							)
            _SCREAMING_SNAKE_CASE				:		int  			=							sd_pipe.to(lowerCAmelCase_							)
            sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

            # forward without prompt embeds
            _SCREAMING_SNAKE_CASE				:		int  			=							self.get_dummy_inputs(lowerCAmelCase_							)
            _SCREAMING_SNAKE_CASE				:		Dict  			=							3 * ["this is a negative prompt"]
            _SCREAMING_SNAKE_CASE				:		List[str]  			=							negative_prompt
            _SCREAMING_SNAKE_CASE				:		Tuple  			=							3 * [inputs["prompt"]]

            _SCREAMING_SNAKE_CASE				:		Optional[Any]  			=							sd_pipe(**lowerCAmelCase_							)
            _SCREAMING_SNAKE_CASE				:		int  			=							output.images[0, -3:, -3:, -1]

            # forward with prompt embeds
            _SCREAMING_SNAKE_CASE				:		str  			=							self.get_dummy_inputs(lowerCAmelCase_							)
            _SCREAMING_SNAKE_CASE				:		int  			=							3 * ["this is a negative prompt"]
            _SCREAMING_SNAKE_CASE				:		List[str]  			=							3 * [inputs.pop("prompt"							)]

            (
                (
                _SCREAMING_SNAKE_CASE							
            )			,	(
                _SCREAMING_SNAKE_CASE							
            )			,	(
                _SCREAMING_SNAKE_CASE							
            )			,	(
                _SCREAMING_SNAKE_CASE							
            )			,								
            )				:		str  			=							sd_pipe.encode_prompt(lowerCAmelCase_			,  negative_prompt=lowerCAmelCase_							)

            _SCREAMING_SNAKE_CASE				:		Optional[Any]  			=							sd_pipe(
                **lowerCAmelCase_			,  prompt_embeds=lowerCAmelCase_			,  negative_prompt_embeds=lowerCAmelCase_			,  pooled_prompt_embeds=lowerCAmelCase_			,  negative_pooled_prompt_embeds=lowerCAmelCase_			,  )
            _SCREAMING_SNAKE_CASE				:		Any  			=							output.images[0, -3:, -3:, -1]

            # make sure that it's equal
            assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()							).max() < 1E-4




       def 			UpperCamelCase_	(					self							)	->   Optional[Any]:
            super().test_attention_slicing_forward_pass(expected_max_diff=3E-3							)







       def 			UpperCamelCase_	(					self							)	->   Optional[int]:
            super().test_inference_batch_single_identical(expected_max_diff=3E-3							)





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






       '''simple docstring'''


       def 			UpperCamelCase_	(					self							)	->   str:
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()




       def 			UpperCamelCase_	(					self			,  __lowerCamelCase			,  __lowerCamelCase="cpu"			,  __lowerCamelCase=torch.floataa			,  __lowerCamelCase=0							)	->   List[Any]:
            _SCREAMING_SNAKE_CASE				:		Tuple  			=							torch.Generator(device=lowerCAmelCase_							).manual_seed(lowerCAmelCase_							)
            _SCREAMING_SNAKE_CASE				:		Union[str, Any]  			=							np.random.RandomState(lowerCAmelCase_							).standard_normal((1, 4, 6_4, 6_4)							)
            _SCREAMING_SNAKE_CASE				:		Union[str, Any]  			=							torch.from_numpy(lowerCAmelCase_							).to(device=lowerCAmelCase_			,  dtype=lowerCAmelCase_							)
            _SCREAMING_SNAKE_CASE				:		List[str]  			=							{
                "prompt": "a photograph of an astronaut riding a horse",
                "latents": latents,
                "generator": generator,
                "num_inference_steps": 3,
                "guidance_scale": 7.5,
                "output_type": "numpy",
            }
            return inputs







       def 			UpperCamelCase_	(					self							)	->   Tuple:
            _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-base"							)
            pipe.to(lowerCAmelCase_							)
            pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

            _SCREAMING_SNAKE_CASE				:		List[Any]  			=							self.get_inputs(lowerCAmelCase_							)
            _SCREAMING_SNAKE_CASE				:		Dict  			=							pipe(**lowerCAmelCase_							).images
            _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							image[0, -3:, -3:, -1].flatten()

            assert image.shape == (1, 5_1_2, 5_1_2, 3)
            _SCREAMING_SNAKE_CASE				:		Tuple  			=							np.array([0.4_9493, 0.4_7896, 0.4_0798, 0.5_4214, 0.5_3212, 0.4_8202, 0.4_7656, 0.4_6329, 0.4_8506]							)
            assert np.abs(image_slice - expected_slice							).max() < 7E-3





'''simple docstring'''


# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


_snake_case      :      Tuple   							=		{
    'configuration_convnextv2': [
        'CONVNEXTV2_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'ConvNextV2Config',
    ]
}

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				_snake_case      :      Tuple   							=		[
				    'CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST',
				    'ConvNextV2ForImageClassification',
				    'ConvNextV2Model',
				    'ConvNextV2PreTrainedModel',
				    'ConvNextV2Backbone',
				]


if TYPE_CHECKING:
				from .configuration_convnextva import (
				    CONVNEXTV2_PRETRAINED_CONFIG_ARCHIVE_MAP,
				    ConvNextVaConfig,
				)

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_convnextva import (
								    CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST,
								    ConvNextVaBackbone,
								    ConvNextVaForImageClassification,
								    ConvNextVaModel,
								    ConvNextVaPreTrainedModel,
								)

else:
				import sys

				_snake_case      :      Any   							=		_LazyModule(__name__, globals()['__file__'], _import_structure)

