





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      A_	(a_					,							a_					,							unittest.TestCase      ):
	UpperCAmelCase__ =      StableDiffusionParadigmsPipeline
	UpperCAmelCase__ =      TEXT_TO_IMAGE_PARAMS
	UpperCAmelCase__ =      TEXT_TO_IMAGE_BATCH_PARAMS
	UpperCAmelCase__ =      TEXT_TO_IMAGE_IMAGE_PARAMS
	UpperCAmelCase__ =      TEXT_TO_IMAGE_IMAGE_PARAMS
	def 					_lowercase    (   self		):






				'''simple docstring'''







				torch.manual_seed(0		)
				UpperCAmelCase       		=    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=_A							, )
				UpperCAmelCase       		=    DDIMParallelScheduler(
				    beta_start=0.0_00_85							, beta_end=0.0_12							, beta_schedule='''scaled_linear'''							, clip_sample=_A							, set_alpha_to_one=_A							, )
				torch.manual_seed(0		)
				UpperCAmelCase       		=    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		)
				UpperCAmelCase       		=    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							, )
				UpperCAmelCase       		=    CLIPTextModel(_A		)
				UpperCAmelCase       		=    CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)

				UpperCAmelCase       		=    {
				    '''unet''': unet,
				    '''scheduler''': scheduler,
				    '''vae''': vae,
				    '''text_encoder''': text_encoder,
				    '''tokenizer''': tokenizer,
				    '''safety_checker''': None,
				    '''feature_extractor''': None,
				}
				return components
	def 					_lowercase    (   self							, _A							, _A=0		):






				'''simple docstring'''







				if str(_A		).startswith('''mps'''		):
							UpperCAmelCase       		=    torch.manual_seed(_A		)
				else:
							UpperCAmelCase       		=    torch.Generator(device=_A		).manual_seed(_A		)
				UpperCAmelCase       		=    {
				    '''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 					_lowercase    (   self		):






				'''simple docstring'''







				UpperCAmelCase       		=    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
				UpperCAmelCase       		=    self.get_dummy_components()
				UpperCAmelCase       		=    StableDiffusionParadigmsPipeline(**_A		)
				UpperCAmelCase       		=    sd_pipe.to(_A		)
				sd_pipe.set_progress_bar_config(disable=_A		)

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

				UpperCAmelCase       		=    np.array([0.47_73, 0.54_17, 0.47_23, 0.49_25, 0.56_31, 0.47_52, 0.52_40, 0.49_35, 0.50_23]		)

				assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2
	def 					_lowercase    (   self		):






				'''simple docstring'''







				UpperCAmelCase       		=    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
				UpperCAmelCase       		=    self.get_dummy_components()
				torch.manual_seed(0		)
				UpperCAmelCase       		=    DDPMParallelScheduler()
				torch.manual_seed(0		)
				UpperCAmelCase       		=    StableDiffusionParadigmsPipeline(**_A		)
				UpperCAmelCase       		=    sd_pipe.to(_A		)
				sd_pipe.set_progress_bar_config(disable=_A		)

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

				UpperCAmelCase       		=    np.array([0.35_73, 0.44_20, 0.49_60, 0.47_99, 0.37_96, 0.38_79, 0.48_19, 0.43_65, 0.44_68]		)

				assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2
	def 					_lowercase    (   self		):






				'''simple docstring'''







				super().test_inference_batch_consistent(batch_sizes=[1, 2]		)
	def 					_lowercase    (   self		):






				'''simple docstring'''







				super().test_inference_batch_single_identical(batch_size=2							, expected_max_diff=3E-3		)







	def 					_lowercase    (   self		):






				'''simple docstring'''







				UpperCAmelCase       		=    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
				UpperCAmelCase       		=    self.get_dummy_components()
				UpperCAmelCase       		=    StableDiffusionParadigmsPipeline(**_A		)
				UpperCAmelCase       		=    sd_pipe.to(_A		)
				sd_pipe.set_progress_bar_config(disable=_A		)

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

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

				UpperCAmelCase       		=    np.array([0.47_71, 0.54_20, 0.46_83, 0.49_18, 0.56_36, 0.47_25, 0.52_30, 0.49_23, 0.50_15]		)

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

@slow
@require_torch_gpu
class      A_	(unittest.TestCase      ):
	def 					_lowercase    (   self		):






				'''simple docstring'''







				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()
	def 					_lowercase    (   self							, _A=0		):






				'''simple docstring'''







				UpperCAmelCase       		=    torch.Generator(device=_A		).manual_seed(_A		)
				UpperCAmelCase       		=    {
				    '''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 					_lowercase    (   self		):






				'''simple docstring'''







				UpperCAmelCase       		=    '''stabilityai/stable-diffusion-2-base'''
				UpperCAmelCase       		=    DDIMParallelScheduler.from_pretrained(_A							, subfolder='''scheduler'''		)
				UpperCAmelCase       		=    StableDiffusionParadigmsPipeline.from_pretrained(_A							, scheduler=_A							, safety_checker=_A		)
				pipe.to(_A		)
				pipe.set_progress_bar_config(disable=_A		)
				pipe.enable_attention_slicing()

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

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

				UpperCAmelCase       		=    np.array([0.96_22, 0.96_02, 0.97_48, 0.95_91, 0.96_30, 0.96_91, 0.96_61, 0.96_31, 0.97_41]		)

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










from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available, is_tokenizers_available


__A    :   Any          =       {}

try:
				if not is_sentencepiece_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__A    :   Union[str, Any]          =       ["CpmTokenizer"]

try:
				if not is_tokenizers_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__A    :   List[Any]          =       ["CpmTokenizerFast"]


if TYPE_CHECKING:
				try:
								if not is_sentencepiece_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .tokenization_cpm import CpmTokenizer

				try:
								if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .tokenization_cpm_fast import CpmTokenizerFast

else:
				import sys

				__A    :   str          =       _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



