



a_			=					[
    'VerificationMode',
    'Version',
    'disable_progress_bar',
    'enable_progress_bar',
    'is_progress_bar_enabled',
    'experimental',
]

from .info_utils import VerificationMode
from .logging import disable_progress_bar, enable_progress_bar, is_progress_bar_enabled
from .version import Version
from .experimental import experimental











import gc
import random
import unittest

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

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

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class  __lowerCAmelCase     (	lowerCAmelCase__  ,	lowerCAmelCase__  ,	unittest.TestCase   ):
 lowerCAmelCase__       					=  StableDiffusionXLImgaImgPipeline
 lowerCAmelCase__       					=  TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
 lowerCAmelCase__       					=  PipelineTesterMixin.required_optional_params - {"""latents"""}
 lowerCAmelCase__       					=  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 lowerCAmelCase__       					=  IMAGE_TO_IMAGE_IMAGE_PARAMS
 lowerCAmelCase__       					=  IMAGE_TO_IMAGE_IMAGE_PARAMS




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    torch.manual_seed(0				)
    __lowerCamelCase													=    UNetaDConditionModel(
        block_out_channels=(32, 64)	,				layers_per_block=2	,				sample_size=32	,				in_channels=4	,				out_channels=4	,				down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')	,				up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')	,				attention_head_dim=(2, 4)	,				use_linear_projection=__UpperCAmelCase	,				addition_embed_type='''text_time'''	,				addition_time_embed_dim=8	,				transformer_layers_per_block=(1, 2)	,				projection_class_embeddings_input_dim=80	,				cross_attention_dim=64	,				)
    __lowerCamelCase													=    EulerDiscreteScheduler(
        beta_start=0.00_085	,				beta_end=0.012	,				steps_offset=1	,				beta_schedule='''scaled_linear'''	,				timestep_spacing='''leading'''	,				)
    torch.manual_seed(0				)
    __lowerCamelCase													=    AutoencoderKL(
        block_out_channels=[32, 64]	,				in_channels=3	,				out_channels=3	,				down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']	,				up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']	,				latent_channels=4	,				sample_size=128	,				)
    torch.manual_seed(0				)
    __lowerCamelCase													=    CLIPTextConfig(
        bos_token_id=0	,				eos_token_id=2	,				hidden_size=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	,				hidden_act='''gelu'''	,				projection_dim=32	,				)
    __lowerCamelCase													=    CLIPTextModel(__UpperCAmelCase				)
    __lowerCamelCase													=    CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''	,				local_files_only=__UpperCAmelCase				)

    __lowerCamelCase													=    CLIPTextModelWithProjection(__UpperCAmelCase				)
    __lowerCamelCase													=    CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''	,				local_files_only=__UpperCAmelCase				)

    __lowerCamelCase													=    {
        '''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        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase=0				):



    '''simple docstring'''
    __lowerCamelCase													=    floats_tensor((1, 3, 32, 32)	,				rng=random.Random(__UpperCAmelCase				)				).to(__UpperCAmelCase				)
    __lowerCamelCase													=    image / 2 + 0.5
    if str(__UpperCAmelCase				).startswith('''mps'''				):
       __lowerCamelCase													=    torch.manual_seed(__UpperCAmelCase				)
    else:
       __lowerCamelCase													=    torch.Generator(device=__UpperCAmelCase				).manual_seed(__UpperCAmelCase				)
    __lowerCamelCase													=    {
        '''prompt''': '''A painting of a squirrel eating a burger''',
        '''image''': image,
        '''generator''': generator,
        '''num_inference_steps''': 2,
        '''guidance_scale''': 5.0,
        '''output_type''': '''numpy''',
        '''strength''': 0.75,
    }
    return inputs




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
    __lowerCamelCase													=    self.get_dummy_components()
    __lowerCamelCase													=    StableDiffusionXLImgaImgPipeline(**__UpperCAmelCase				)
    __lowerCamelCase													=    sd_pipe.to(__UpperCAmelCase				)
    sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase				)

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

    assert image.shape == (1, 32, 32, 3)

    __lowerCamelCase													=    np.array([0.4_656, 0.4_840, 0.4_439, 0.6_698, 0.5_574, 0.4_524, 0.5_799, 0.5_943, 0.5_165]				)

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




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    super().test_attention_slicing_forward_pass(expected_max_diff=3E-3				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    super().test_inference_batch_single_identical(expected_max_diff=3E-3				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    pass




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.get_dummy_components()
    __lowerCamelCase													=    StableDiffusionXLImgaImgPipeline(**__UpperCAmelCase				)
    __lowerCamelCase													=    sd_pipe.to(__UpperCAmelCase				)
    __lowerCamelCase													=    sd_pipe.to(__UpperCAmelCase				)
    sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase				)

    # forward without prompt embeds
    __lowerCamelCase													=    self.get_dummy_inputs(__UpperCAmelCase				)
    __lowerCamelCase													=    3 * ['''this is a negative prompt''']
    __lowerCamelCase													=    negative_prompt
    __lowerCamelCase													=    3 * [inputs['''prompt''']]

    __lowerCamelCase													=    sd_pipe(**__UpperCAmelCase				)
    __lowerCamelCase													=    output.images[0, -3:, -3:, -1]

    # forward with prompt embeds
    __lowerCamelCase													=    self.get_dummy_inputs(__UpperCAmelCase				)
    __lowerCamelCase													=    3 * ['''this is a negative prompt''']
    __lowerCamelCase													=    3 * [inputs.pop('''prompt'''				)]

    (
        (
        __lowerCamelCase
    )							,(
        __lowerCamelCase
    )							,(
        __lowerCamelCase
    )							,(
        __lowerCamelCase
    )							,
    )													=    sd_pipe.encode_prompt(__UpperCAmelCase	,				negative_prompt=__UpperCAmelCase				)

    __lowerCamelCase													=    sd_pipe(
        **__UpperCAmelCase	,				prompt_embeds=__UpperCAmelCase	,				negative_prompt_embeds=__UpperCAmelCase	,				pooled_prompt_embeds=__UpperCAmelCase	,				negative_pooled_prompt_embeds=__UpperCAmelCase	,				)
    __lowerCamelCase													=    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



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




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()




 def        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase="cpu"	,				__UpperCAmelCase=torch.floataa	,				__UpperCAmelCase=0				):



    '''simple docstring'''
    __lowerCamelCase													=    torch.Generator(device=__UpperCAmelCase				).manual_seed(__UpperCAmelCase				)
    __lowerCamelCase													=    np.random.RandomState(__UpperCAmelCase				).standard_normal((1, 4, 64, 64)				)
    __lowerCamelCase													=    torch.from_numpy(__UpperCAmelCase				).to(device=__UpperCAmelCase	,				dtype=__UpperCAmelCase				)
    __lowerCamelCase													=    {
        '''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        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''				)
    pipe.to(__UpperCAmelCase				)
    pipe.set_progress_bar_config(disable=__UpperCAmelCase				)

    __lowerCamelCase													=    self.get_inputs(__UpperCAmelCase				)
    __lowerCamelCase													=    pipe(**__UpperCAmelCase				).images
    __lowerCamelCase													=    image[0, -3:, -3:, -1].flatten()

    assert image.shape == (1, 512, 512, 3)
    __lowerCamelCase													=    np.array([0.49_493, 0.47_896, 0.40_798, 0.54_214, 0.53_212, 0.48_202, 0.47_656, 0.46_329, 0.48_506]				)
    assert np.abs(image_slice - expected_slice				).max() < 7E-3






