

"""simple docstring"""

import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    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       _snake_case  (    lowerCAmelCase__    ,      lowerCAmelCase__    ,      unittest.TestCase							):
						__snake_case			      =  StableDiffusionParadigmsPipeline
						__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				)      ->    int:
											torch.manual_seed(0				)
											__lowercase        =      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''')  ,				cross_attention_dim=3_2  ,				attention_head_dim=(2, 4)  ,				use_linear_projection=_lowercase  ,				)
											__lowercase        =      DDIMParallelScheduler(
											    beta_start=0.0_0_0_8_5  ,				beta_end=0.0_1_2  ,				beta_schedule='''scaled_linear'''  ,				clip_sample=_lowercase  ,				set_alpha_to_one=_lowercase  ,				)
											torch.manual_seed(0				)
											__lowercase        =      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				)
											__lowercase        =      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=5_1_2  ,				)
											__lowercase        =      CLIPTextModel(_lowercase				)
											__lowercase        =      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				)

											__lowercase        =      {
											    '''unet''': unet,
											    '''scheduler''': scheduler,
											    '''vae''': vae,
											    '''text_encoder''': text_encoder,
											    '''tokenizer''': tokenizer,
											    '''safety_checker''': None,
											    '''feature_extractor''': None,
											}
											return components







						def 							_UpperCamelCase     (self  ,				UpperCAmelCase_  ,				UpperCAmelCase_=0				)      ->    List[str]:


											if str(_lowercase				).startswith('''mps'''				):
																__lowercase        =      torch.manual_seed(_lowercase				)
											else:
																__lowercase        =      torch.Generator(device=_lowercase				).manual_seed(_lowercase				)
											__lowercase        =      {
											    '''prompt''': '''a photograph of an astronaut riding a horse''',
											    '''generator''': generator,
											    '''num_inference_steps''': 1_0,
											    '''guidance_scale''': 6.0,
											    '''output_type''': '''numpy''',
											    '''parallel''': 3,
											    '''debug''': True,
											}
											return inputs







						def 							_UpperCamelCase     (self				)      ->    Union[str, Any]:
											__lowercase        =      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
											__lowercase        =      self.get_dummy_components()
											__lowercase        =      StableDiffusionParadigmsPipeline(**_lowercase				)
											__lowercase        =      sd_pipe.to(_lowercase				)
											sd_pipe.set_progress_bar_config(disable=_lowercase				)

											__lowercase        =      self.get_dummy_inputs(_lowercase				)
											__lowercase        =      sd_pipe(**_lowercase				).images
											__lowercase        =      image[0, -3:, -3:, -1]
											assert image.shape == (1, 6_4, 6_4, 3)

											__lowercase        =      np.array([0.4_7_7_3, 0.5_4_1_7, 0.4_7_2_3, 0.4_9_2_5, 0.5_6_3_1, 0.4_7_5_2, 0.5_2_4_0, 0.4_9_3_5, 0.5_0_2_3]				)

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







						def 							_UpperCamelCase     (self				)      ->    Optional[Any]:
											__lowercase        =      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
											__lowercase        =      self.get_dummy_components()
											torch.manual_seed(0				)
											__lowercase        =      DDPMParallelScheduler()
											torch.manual_seed(0				)
											__lowercase        =      StableDiffusionParadigmsPipeline(**_lowercase				)
											__lowercase        =      sd_pipe.to(_lowercase				)
											sd_pipe.set_progress_bar_config(disable=_lowercase				)

											__lowercase        =      self.get_dummy_inputs(_lowercase				)
											__lowercase        =      sd_pipe(**_lowercase				).images
											__lowercase        =      image[0, -3:, -3:, -1]
											assert image.shape == (1, 6_4, 6_4, 3)

											__lowercase        =      np.array([0.3_5_7_3, 0.4_4_2_0, 0.4_9_6_0, 0.4_7_9_9, 0.3_7_9_6, 0.3_8_7_9, 0.4_8_1_9, 0.4_3_6_5, 0.4_4_6_8]				)

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







						def 							_UpperCamelCase     (self				)      ->    int:
											super().test_inference_batch_consistent(batch_sizes=[1, 2]				)







						def 							_UpperCamelCase     (self				)      ->    Tuple:
											super().test_inference_batch_single_identical(batch_size=2  ,				expected_max_diff=3e-3				)







						def 							_UpperCamelCase     (self				)      ->    Union[str, Any]:
											__lowercase        =      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
											__lowercase        =      self.get_dummy_components()
											__lowercase        =      StableDiffusionParadigmsPipeline(**_lowercase				)
											__lowercase        =      sd_pipe.to(_lowercase				)
											sd_pipe.set_progress_bar_config(disable=_lowercase				)

											__lowercase        =      self.get_dummy_inputs(_lowercase				)
											__lowercase        =      '''french fries'''
											__lowercase        =      sd_pipe(**_lowercase  ,				negative_prompt=_lowercase				)
											__lowercase        =      output.images
											__lowercase        =      image[0, -3:, -3:, -1]

											assert image.shape == (1, 6_4, 6_4, 3)

											__lowercase        =      np.array([0.4_7_7_1, 0.5_4_2_0, 0.4_6_8_3, 0.4_9_1_8, 0.5_6_3_6, 0.4_7_2_5, 0.5_2_3_0, 0.4_9_2_3, 0.5_0_1_5]				)

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



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







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







						def 							_UpperCamelCase     (self  ,				UpperCAmelCase_=0				)      ->    List[str]:
											__lowercase        =      torch.Generator(device=_lowercase				).manual_seed(_lowercase				)
											__lowercase        =      {
											    '''prompt''': '''a photograph of an astronaut riding a horse''',
											    '''generator''': generator,
											    '''num_inference_steps''': 1_0,
											    '''guidance_scale''': 7.5,
											    '''output_type''': '''numpy''',
											    '''parallel''': 3,
											    '''debug''': True,
											}
											return inputs







						def 							_UpperCamelCase     (self				)      ->    Optional[Any]:
											__lowercase        =      '''stabilityai/stable-diffusion-2-base'''
											__lowercase        =      DDIMParallelScheduler.from_pretrained(_lowercase  ,				subfolder='''scheduler'''				)
											__lowercase        =      StableDiffusionParadigmsPipeline.from_pretrained(_lowercase  ,				scheduler=_lowercase  ,				safety_checker=_lowercase				)
											pipe.to(_lowercase				)
											pipe.set_progress_bar_config(disable=_lowercase				)
											pipe.enable_attention_slicing()

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

											assert image.shape == (1, 5_1_2, 5_1_2, 3)

											__lowercase        =      np.array([0.9_6_2_2, 0.9_6_0_2, 0.9_7_4_8, 0.9_5_9_1, 0.9_6_3_0, 0.9_6_9_1, 0.9_6_6_1, 0.9_6_3_1, 0.9_7_4_1]				)

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








from ..utils import DummyObject, requires_backends

class 		lowercase_       (	metaclass=lowerCAmelCase__				):
					__UpperCamelCase     	= ["torch", "transformers", "k_diffusion"]




					def __init__(   self:		Tuple,     *_lowercase:		Dict,     **_lowercase:		Optional[int]):


									'''simple docstring'''


									requires_backends(self,     ["""torch""", """transformers""", """k_diffusion"""])




					@classmethod
					def     _lowercase						(   cls:		Union[str, Any],     *_lowercase:		Union[str, Any],     **_lowercase:		Optional[int]):


									'''simple docstring'''


									requires_backends(cls,     ["""torch""", """transformers""", """k_diffusion"""])







					@classmethod
					def     _lowercase						(   cls:		List[str],     *_lowercase:		int,     **_lowercase:		List[Any]):


									'''simple docstring'''


									requires_backends(cls,     ["""torch""", """transformers""", """k_diffusion"""])





