





import gc
import random
import unittest

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

import diffusers
from diffusers import (
    AutoencoderKL,
    EulerDiscreteScheduler,
    StableDiffusionLatentUpscalePipeline,
    StableDiffusionPipeline,
    UNetaDConditionModel,
)
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


def       UpperCamelCase__(			UpperCamelCase__     :			List[str]    )->Any:
       A__		   =		[tensor.shape for tensor in tensor_list]
       return all(shape == shapes[0] for shape in shapes[1:]    )




class        SCREAMING_SNAKE_CASE__					(		UpperCamelCase__							,   UpperCamelCase__							,   UpperCamelCase__							,   unittest.TestCase ):
 __SCREAMING_SNAKE_CASE							=      StableDiffusionLatentUpscalePipeline
 __SCREAMING_SNAKE_CASE							=      TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
     '''height''',
     '''width''',
     '''cross_attention_kwargs''',
     '''negative_prompt_embeds''',
     '''prompt_embeds''',
 }
 __SCREAMING_SNAKE_CASE							=      PipelineTesterMixin.required_optional_params - {'''num_images_per_prompt'''}
 __SCREAMING_SNAKE_CASE							=      TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 __SCREAMING_SNAKE_CASE							=      frozenset(
     [] )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
 __SCREAMING_SNAKE_CASE							=      frozenset([] )

 __SCREAMING_SNAKE_CASE							=      True



 @property
 def 							UpperCamelCase     (				self					):
        A__		   =		1
        A__		   =		4
        A__		   =		(16, 16)

        A__		   =		floats_tensor((batch_size, num_channels) + sizes,rng=random.Random(0					)					).to(__lowerCamelCase					)
        return image



 def 							UpperCamelCase     (				self					):
        torch.manual_seed(0					)
        A__		   =		UNetaDConditionModel(
            act_fn='''gelu''',attention_head_dim=8,norm_num_groups=__lowerCamelCase,block_out_channels=[32, 32, 64, 64],time_cond_proj_dim=160,conv_in_kernel=1,conv_out_kernel=1,cross_attention_dim=32,down_block_types=(
                '''KDownBlock2D''',
                '''KCrossAttnDownBlock2D''',
                '''KCrossAttnDownBlock2D''',
                '''KCrossAttnDownBlock2D''',
            ),in_channels=8,mid_block_type=__lowerCamelCase,only_cross_attention=__lowerCamelCase,out_channels=5,resnet_time_scale_shift='''scale_shift''',time_embedding_type='''fourier''',timestep_post_act='''gelu''',up_block_types=('''KCrossAttnUpBlock2D''', '''KCrossAttnUpBlock2D''', '''KCrossAttnUpBlock2D''', '''KUpBlock2D'''),)
        A__		   =		AutoencoderKL(
            block_out_channels=[32, 32, 64, 64],in_channels=3,out_channels=3,down_block_types=[
                '''DownEncoderBlock2D''',
                '''DownEncoderBlock2D''',
                '''DownEncoderBlock2D''',
                '''DownEncoderBlock2D''',
            ],up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''', '''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],latent_channels=4,)
        A__		   =		EulerDiscreteScheduler(prediction_type='''sample'''					)
        A__		   =		CLIPTextConfig(
            bos_token_id=0,eos_token_id=2,hidden_size=32,intermediate_size=37,layer_norm_eps=1E-05,num_attention_heads=4,num_hidden_layers=5,pad_token_id=1,vocab_size=1000,hidden_act='''quick_gelu''',projection_dim=512,)
        A__		   =		CLIPTextModel(__lowerCamelCase					)
        A__		   =		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)

        A__		   =		{
            '''unet''': model.eval(),
            '''vae''': vae.eval(),
            '''scheduler''': scheduler,
            '''text_encoder''': text_encoder,
            '''tokenizer''': tokenizer,
        }

        return components



 def 							UpperCamelCase     (				self,__lowerCamelCase,__lowerCamelCase=0					):
        if str(__lowerCamelCase					).startswith('''mps'''					):
               A__		   =		torch.manual_seed(__lowerCamelCase					)
        else:
               A__		   =		torch.Generator(device=__lowerCamelCase					).manual_seed(__lowerCamelCase					)
        A__		   =		{
            '''prompt''': '''A painting of a squirrel eating a burger''',
            '''image''': self.dummy_image.cpu(),
            '''generator''': generator,
            '''num_inference_steps''': 2,
            '''output_type''': '''numpy''',
        }
        return inputs



 def 							UpperCamelCase     (				self					):
        A__		   =		'''cpu'''

        A__		   =		self.get_dummy_components()
        A__		   =		self.pipeline_class(**__lowerCamelCase					)
        pipe.to(__lowerCamelCase					)
        pipe.set_progress_bar_config(disable=__lowerCamelCase					)

        A__		   =		self.get_dummy_inputs(__lowerCamelCase					)
        A__		   =		pipe(**__lowerCamelCase					).images
        A__		   =		image[0, -3:, -3:, -1]

        self.assertEqual(image.shape,(1, 256, 256, 3)					)
        A__		   =		np.array(
            [0.47222412, 0.41921633, 0.44717434, 0.46874192, 0.42588258, 0.46150726, 0.4677534, 0.45583832, 0.48579055]					)
        A__		   =		np.abs(image_slice.flatten() - expected_slice					).max()
        self.assertLessEqual(__lowerCamelCase,1E-3					)



 def 							UpperCamelCase     (				self					):
        super().test_attention_slicing_forward_pass(expected_max_diff=7E-3					)



 def 							UpperCamelCase     (				self					):
        super().test_cpu_offload_forward_pass(expected_max_diff=3E-3					)



 def 							UpperCamelCase     (				self					):
        super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3					)



 def 							UpperCamelCase     (				self					):
        super().test_inference_batch_single_identical(expected_max_diff=7E-3					)



 def 							UpperCamelCase     (				self					):
        super().test_pt_np_pil_outputs_equivalent(expected_max_diff=3E-3					)



 def 							UpperCamelCase     (				self					):
        super().test_save_load_local(expected_max_difference=3E-3					)



 def 							UpperCamelCase     (				self					):
        super().test_save_load_optional_components(expected_max_difference=3E-3					)




 def 							UpperCamelCase     (				self					):
        A__		   =		[
            '''DDIMScheduler''',
            '''DDPMScheduler''',
            '''PNDMScheduler''',
            '''HeunDiscreteScheduler''',
            '''EulerAncestralDiscreteScheduler''',
            '''KDPM2DiscreteScheduler''',
            '''KDPM2AncestralDiscreteScheduler''',
            '''DPMSolverSDEScheduler''',
        ]
        A__		   =		self.get_dummy_components()
        A__		   =		self.pipeline_class(**__lowerCamelCase					)

        # make sure that PNDM does not need warm-up
        pipe.scheduler.register_to_config(skip_prk_steps=__lowerCamelCase					)

        pipe.to(__lowerCamelCase					)
        pipe.set_progress_bar_config(disable=__lowerCamelCase					)
        A__		   =		self.get_dummy_inputs(__lowerCamelCase					)
        A__		   =		2

        A__		   =		[]
        for scheduler_enum in KarrasDiffusionSchedulers:
               if scheduler_enum.name in skip_schedulers:
                      # no sigma schedulers are not supported
                      # no schedulers
                      continue

               A__		   =		getattr(__lowerCamelCase,scheduler_enum.name					)
               A__		   =		scheduler_cls.from_config(pipe.scheduler.config					)
               A__		   =		pipe(**__lowerCamelCase					)[0]
               outputs.append(__lowerCamelCase					)

        assert check_same_shape(__lowerCamelCase					)




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



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



 def 							UpperCamelCase     (				self					):
        A__		   =		torch.manual_seed(33					)

        A__		   =		StableDiffusionPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4''',torch_dtype=torch.floataa					)
        pipe.to('''cuda'''					)

        A__		   =		StableDiffusionLatentUpscalePipeline.from_pretrained(
            '''stabilityai/sd-x2-latent-upscaler''',torch_dtype=torch.floataa					)
        upscaler.to('''cuda'''					)

        A__		   =		'''a photo of an astronaut high resolution, unreal engine, ultra realistic'''

        A__		   =		pipe(__lowerCamelCase,generator=__lowerCamelCase,output_type='''latent'''					).images

        A__		   =		upscaler(
            prompt=__lowerCamelCase,image=__lowerCamelCase,num_inference_steps=20,guidance_scale=0,generator=__lowerCamelCase,output_type='''np''',).images[0]

        A__		   =		load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/astronaut_1024.npy'''					)
        assert np.abs((expected_image - image).mean()					) < 5E-2




 def 							UpperCamelCase     (				self					):
        A__		   =		torch.manual_seed(33					)

        A__		   =		StableDiffusionLatentUpscalePipeline.from_pretrained(
            '''stabilityai/sd-x2-latent-upscaler''',torch_dtype=torch.floataa					)
        upscaler.to('''cuda'''					)

        A__		   =		'''the temple of fire by Ross Tran and Gerardo Dottori, oil on canvas'''

        A__		   =		load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_512.png'''					)

        A__		   =		upscaler(
            prompt=__lowerCamelCase,image=__lowerCamelCase,num_inference_steps=20,guidance_scale=0,generator=__lowerCamelCase,output_type='''np''',).images[0]

        A__		   =		load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_1024.npy'''					)
        assert np.abs((expected_image - image).max()					) < 5E-2












a__: str							   =						'0.21.0'

from .accelerator import Accelerator
from .big_modeling import (
    cpu_offload,
    cpu_offload_with_hook,
    disk_offload,
    dispatch_model,
    init_empty_weights,
    init_on_device,
    load_checkpoint_and_dispatch,
)
from .data_loader import skip_first_batches
from .launchers import debug_launcher, notebook_launcher
from .state import PartialState
from .utils import (
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    FullyShardedDataParallelPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    find_executable_batch_size,
    infer_auto_device_map,
    is_rich_available,
    load_checkpoint_in_model,
    synchronize_rng_states,
)


if is_rich_available():
       from .utils import rich





