"""simple docstring"""





import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    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     (					_UpperCamelCase				, _UpperCamelCase				, unittest.TestCase		):
						snake_case__			  =			StableDiffusionSAGPipeline
						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
						snake_case__			  =			False






						def       lowerCamelCase__     (       self      :			str    ):
													torch.manual_seed(0    )
													__lowerCamelCase					:		List[Any]		   =     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")							,					cross_attention_dim=32							,					)
													__lowerCamelCase					:		Any		   =     DDIMScheduler(
													    beta_start=0.0_0_0_8_5							,					beta_end=0.0_1_2							,					beta_schedule="scaled_linear"							,					clip_sample=UpperCAmelCase							,					set_alpha_to_one=UpperCAmelCase							,					)
													torch.manual_seed(0    )
													__lowerCamelCase					:		Any		   =     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							,					)
													torch.manual_seed(0    )
													__lowerCamelCase					:		List[Any]		   =     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							,					)
													__lowerCamelCase					:		int		   =     CLIPTextModel(UpperCAmelCase    )
													__lowerCamelCase					:		Tuple		   =     CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"    )

													__lowerCamelCase					:		Optional[int]		   =     {
													    "unet": unet,
													    "scheduler": scheduler,
													    "vae": vae,
													    "text_encoder": text_encoder,
													    "tokenizer": tokenizer,
													    "safety_checker": None,
													    "feature_extractor": None,
													}
													return components






						def       lowerCamelCase__     (       self      :			Optional[Any]							,					UpperCAmelCase      :			Tuple							,					UpperCAmelCase      :			List[Any]=0    ):
													if str(UpperCAmelCase    ).startswith("mps"    ):
																				__lowerCamelCase					:		int		   =     torch.manual_seed(UpperCAmelCase    )
													else:
																				__lowerCamelCase					:		List[str]		   =     torch.Generator(device=UpperCAmelCase    ).manual_seed(UpperCAmelCase    )
													__lowerCamelCase					:		Union[str, Any]		   =     {
													    "prompt": ".",
													    "generator": generator,
													    "num_inference_steps": 2,
													    "guidance_scale": 1.0,
													    "sag_scale": 1.0,
													    "output_type": "numpy",
													}
													return inputs







						def       lowerCamelCase__     (       self      :			Optional[Any]    ):
													super().test_inference_batch_single_identical(expected_max_diff=3E-3    )

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






						def       lowerCamelCase__     (       self      :			Optional[int]    ):
													# clean up the VRAM after each test
													super().tearDown()
													gc.collect()
													torch.cuda.empty_cache()






						def       lowerCamelCase__     (       self      :			str    ):
													__lowerCamelCase					:		str		   =     StableDiffusionSAGPipeline.from_pretrained("CompVis/stable-diffusion-v1-4"    )
													__lowerCamelCase					:		int		   =     sag_pipe.to(UpperCAmelCase    )
													sag_pipe.set_progress_bar_config(disable=UpperCAmelCase    )

													__lowerCamelCase					:		Tuple		   =     "."
													__lowerCamelCase					:		Tuple		   =     torch.manual_seed(0    )
													__lowerCamelCase					:		Any		   =     sag_pipe(
													    [prompt]							,					generator=UpperCAmelCase							,					guidance_scale=7.5							,					sag_scale=1.0							,					num_inference_steps=20							,					output_type="np"    )

													__lowerCamelCase					:		str		   =     output.images

													__lowerCamelCase					:		Dict		   =     image[0, -3:, -3:, -1]

													assert image.shape == (1, 512, 512, 3)
													__lowerCamelCase					:		List[Any]		   =     np.array([0.1_5_6_8, 0.1_7_3_8, 0.1_6_9_5, 0.1_6_9_3, 0.1_5_0_7, 0.1_7_0_5, 0.1_5_4_7, 0.1_7_5_1, 0.1_9_4_9]    )

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






						def       lowerCamelCase__     (       self      :			str    ):
													__lowerCamelCase					:		Dict		   =     StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"    )
													__lowerCamelCase					:		Tuple		   =     sag_pipe.to(UpperCAmelCase    )
													sag_pipe.set_progress_bar_config(disable=UpperCAmelCase    )

													__lowerCamelCase					:		Optional[int]		   =     "."
													__lowerCamelCase					:		int		   =     torch.manual_seed(0    )
													__lowerCamelCase					:		Optional[Any]		   =     sag_pipe(
													    [prompt]							,					generator=UpperCAmelCase							,					guidance_scale=7.5							,					sag_scale=1.0							,					num_inference_steps=20							,					output_type="np"    )

													__lowerCamelCase					:		List[str]		   =     output.images

													__lowerCamelCase					:		Union[str, Any]		   =     image[0, -3:, -3:, -1]

													assert image.shape == (1, 512, 512, 3)
													__lowerCamelCase					:		List[str]		   =     np.array([0.3_4_5_9, 0.2_8_7_6, 0.2_5_3_7, 0.3_0_0_2, 0.2_6_7_1, 0.2_1_6_0, 0.3_0_2_6, 0.2_2_6_2, 0.2_3_7_1]    )

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







						def       lowerCamelCase__     (       self      :			Dict    ):
													__lowerCamelCase					:		Optional[int]		   =     StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"    )
													__lowerCamelCase					:		List[Any]		   =     sag_pipe.to(UpperCAmelCase    )
													sag_pipe.set_progress_bar_config(disable=UpperCAmelCase    )

													__lowerCamelCase					:		Tuple		   =     "."
													__lowerCamelCase					:		List[str]		   =     torch.manual_seed(0    )
													__lowerCamelCase					:		Tuple		   =     sag_pipe(
													    [prompt]							,					width=768							,					height=512							,					generator=UpperCAmelCase							,					guidance_scale=7.5							,					sag_scale=1.0							,					num_inference_steps=20							,					output_type="np"							,					)

													__lowerCamelCase					:		int		   =     output.images

													assert image.shape == (1, 512, 768, 3)
"""simple docstring"""





from typing import Any, Callable, Dict, List, Optional, Union

import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker


__A  =				'''CompVis/stable-diffusion-v1-1'''
__A  =				'''CompVis/stable-diffusion-v1-2'''
__A  =				'''CompVis/stable-diffusion-v1-3'''
__A  =				'''CompVis/stable-diffusion-v1-4'''
class  _snake_case     (					a__		):






						def __init__(       self      :			int							,					UpperCAmelCase      :			AutoencoderKL							,					UpperCAmelCase      :			CLIPTextModel							,					UpperCAmelCase      :			CLIPTokenizer							,					UpperCAmelCase      :			UNetaDConditionModel							,					UpperCAmelCase      :			Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]							,					UpperCAmelCase      :			StableDiffusionSafetyChecker							,					UpperCAmelCase      :			CLIPImageProcessor							,					UpperCAmelCase      :			bool = True							,					):
													super()._init_()

													__lowerCamelCase					:		Optional[int]		   =     StableDiffusionPipeline.from_pretrained(UpperCAmelCase    )
													__lowerCamelCase					:		Optional[int]		   =     StableDiffusionPipeline.from_pretrained(UpperCAmelCase    )
													__lowerCamelCase					:		Tuple		   =     StableDiffusionPipeline.from_pretrained(UpperCAmelCase    )
													__lowerCamelCase					:		List[str]		   =     StableDiffusionPipeline(
													    vae=UpperCAmelCase							,					text_encoder=UpperCAmelCase							,					tokenizer=UpperCAmelCase							,					unet=UpperCAmelCase							,					scheduler=UpperCAmelCase							,					safety_checker=UpperCAmelCase							,					feature_extractor=UpperCAmelCase							,					requires_safety_checker=UpperCAmelCase							,					)

													self.register_modules(pipelinea=self.pipea							,					pipelinea=self.pipea							,					pipelinea=self.pipea							,					pipelinea=self.pipea    )






						@property
						def       lowerCamelCase__     (       self      :			Dict    ):
													return {k: getattr(self							,					UpperCAmelCase    ) for k in self.config.keys() if not k.startswith("_"    )}






						def       lowerCamelCase__     (       self      :			str							,					UpperCAmelCase      :			Optional[Union[str, int]] = "auto"    ):
													if slice_size == "auto":
																				# half the attention head size is usually a good trade-off between
																				# speed and memory
																				__lowerCamelCase					:		Optional[int]		   =     self.unet.config.attention_head_dim // 2
													self.unet.set_attention_slice(UpperCAmelCase    )






						def       lowerCamelCase__     (       self      :			Tuple    ):
													self.enable_attention_slicing(UpperCAmelCase    )






						@torch.no_grad()
						def       lowerCamelCase__     (       self      :			Dict							,					UpperCAmelCase      :			Union[str, List[str]]							,					UpperCAmelCase      :			int = 512							,					UpperCAmelCase      :			int = 512							,					UpperCAmelCase      :			int = 50							,					UpperCAmelCase      :			float = 7.5							,					UpperCAmelCase      :			Optional[Union[str, List[str]]] = None							,					UpperCAmelCase      :			Optional[int] = 1							,					UpperCAmelCase      :			float = 0.0							,					UpperCAmelCase      :			Optional[torch.Generator] = None							,					UpperCAmelCase      :			Optional[torch.FloatTensor] = None							,					UpperCAmelCase      :			Optional[str] = "pil"							,					UpperCAmelCase      :			bool = True							,					UpperCAmelCase      :			Optional[Callable[[int, int, torch.FloatTensor], None]] = None							,					UpperCAmelCase      :			int = 1							,					**UpperCAmelCase      :			Optional[Any]							,					):
													return self.pipea(
													    prompt=UpperCAmelCase							,					height=UpperCAmelCase							,					width=UpperCAmelCase							,					num_inference_steps=UpperCAmelCase							,					guidance_scale=UpperCAmelCase							,					negative_prompt=UpperCAmelCase							,					num_images_per_prompt=UpperCAmelCase							,					eta=UpperCAmelCase							,					generator=UpperCAmelCase							,					latents=UpperCAmelCase							,					output_type=UpperCAmelCase							,					return_dict=UpperCAmelCase							,					callback=UpperCAmelCase							,					callback_steps=UpperCAmelCase							,					**UpperCAmelCase							,					)






						@torch.no_grad()
						def       lowerCamelCase__     (       self      :			Dict							,					UpperCAmelCase      :			Union[str, List[str]]							,					UpperCAmelCase      :			int = 512							,					UpperCAmelCase      :			int = 512							,					UpperCAmelCase      :			int = 50							,					UpperCAmelCase      :			float = 7.5							,					UpperCAmelCase      :			Optional[Union[str, List[str]]] = None							,					UpperCAmelCase      :			Optional[int] = 1							,					UpperCAmelCase      :			float = 0.0							,					UpperCAmelCase      :			Optional[torch.Generator] = None							,					UpperCAmelCase      :			Optional[torch.FloatTensor] = None							,					UpperCAmelCase      :			Optional[str] = "pil"							,					UpperCAmelCase      :			bool = True							,					UpperCAmelCase      :			Optional[Callable[[int, int, torch.FloatTensor], None]] = None							,					UpperCAmelCase      :			int = 1							,					**UpperCAmelCase      :			Union[str, Any]							,					):
													return self.pipea(
													    prompt=UpperCAmelCase							,					height=UpperCAmelCase							,					width=UpperCAmelCase							,					num_inference_steps=UpperCAmelCase							,					guidance_scale=UpperCAmelCase							,					negative_prompt=UpperCAmelCase							,					num_images_per_prompt=UpperCAmelCase							,					eta=UpperCAmelCase							,					generator=UpperCAmelCase							,					latents=UpperCAmelCase							,					output_type=UpperCAmelCase							,					return_dict=UpperCAmelCase							,					callback=UpperCAmelCase							,					callback_steps=UpperCAmelCase							,					**UpperCAmelCase							,					)






						@torch.no_grad()
						def       lowerCamelCase__     (       self      :			Tuple							,					UpperCAmelCase      :			Union[str, List[str]]							,					UpperCAmelCase      :			int = 512							,					UpperCAmelCase      :			int = 512							,					UpperCAmelCase      :			int = 50							,					UpperCAmelCase      :			float = 7.5							,					UpperCAmelCase      :			Optional[Union[str, List[str]]] = None							,					UpperCAmelCase      :			Optional[int] = 1							,					UpperCAmelCase      :			float = 0.0							,					UpperCAmelCase      :			Optional[torch.Generator] = None							,					UpperCAmelCase      :			Optional[torch.FloatTensor] = None							,					UpperCAmelCase      :			Optional[str] = "pil"							,					UpperCAmelCase      :			bool = True							,					UpperCAmelCase      :			Optional[Callable[[int, int, torch.FloatTensor], None]] = None							,					UpperCAmelCase      :			int = 1							,					**UpperCAmelCase      :			Dict							,					):
													return self.pipea(
													    prompt=UpperCAmelCase							,					height=UpperCAmelCase							,					width=UpperCAmelCase							,					num_inference_steps=UpperCAmelCase							,					guidance_scale=UpperCAmelCase							,					negative_prompt=UpperCAmelCase							,					num_images_per_prompt=UpperCAmelCase							,					eta=UpperCAmelCase							,					generator=UpperCAmelCase							,					latents=UpperCAmelCase							,					output_type=UpperCAmelCase							,					return_dict=UpperCAmelCase							,					callback=UpperCAmelCase							,					callback_steps=UpperCAmelCase							,					**UpperCAmelCase							,					)






						@torch.no_grad()
						def       lowerCamelCase__     (       self      :			Optional[Any]							,					UpperCAmelCase      :			Union[str, List[str]]							,					UpperCAmelCase      :			int = 512							,					UpperCAmelCase      :			int = 512							,					UpperCAmelCase      :			int = 50							,					UpperCAmelCase      :			float = 7.5							,					UpperCAmelCase      :			Optional[Union[str, List[str]]] = None							,					UpperCAmelCase      :			Optional[int] = 1							,					UpperCAmelCase      :			float = 0.0							,					UpperCAmelCase      :			Optional[torch.Generator] = None							,					UpperCAmelCase      :			Optional[torch.FloatTensor] = None							,					UpperCAmelCase      :			Optional[str] = "pil"							,					UpperCAmelCase      :			bool = True							,					UpperCAmelCase      :			Optional[Callable[[int, int, torch.FloatTensor], None]] = None							,					UpperCAmelCase      :			int = 1							,					**UpperCAmelCase      :			List[str]							,					):
													return self.pipea(
													    prompt=UpperCAmelCase							,					height=UpperCAmelCase							,					width=UpperCAmelCase							,					num_inference_steps=UpperCAmelCase							,					guidance_scale=UpperCAmelCase							,					negative_prompt=UpperCAmelCase							,					num_images_per_prompt=UpperCAmelCase							,					eta=UpperCAmelCase							,					generator=UpperCAmelCase							,					latents=UpperCAmelCase							,					output_type=UpperCAmelCase							,					return_dict=UpperCAmelCase							,					callback=UpperCAmelCase							,					callback_steps=UpperCAmelCase							,					**UpperCAmelCase							,					)







						@torch.no_grad()
						def       lowerCamelCase__     (       self      :			str							,					UpperCAmelCase      :			Union[str, List[str]]							,					UpperCAmelCase      :			int = 512							,					UpperCAmelCase      :			int = 512							,					UpperCAmelCase      :			int = 50							,					UpperCAmelCase      :			float = 7.5							,					UpperCAmelCase      :			Optional[Union[str, List[str]]] = None							,					UpperCAmelCase      :			Optional[int] = 1							,					UpperCAmelCase      :			float = 0.0							,					UpperCAmelCase      :			Optional[torch.Generator] = None							,					UpperCAmelCase      :			Optional[torch.FloatTensor] = None							,					UpperCAmelCase      :			Optional[str] = "pil"							,					UpperCAmelCase      :			bool = True							,					UpperCAmelCase      :			Optional[Callable[[int, int, torch.FloatTensor], None]] = None							,					UpperCAmelCase      :			int = 1							,					**UpperCAmelCase      :			Tuple							,					):
													__lowerCamelCase					:		Optional[int]		   =     "cuda" if torch.cuda.is_available() else "cpu"
													self.to(UpperCAmelCase    )

													# Checks if the height and width are divisible by 8 or not
													if height % 8 != 0 or width % 8 != 0:
																				raise ValueError(F"""`height` and `width` must be divisible by 8 but are {height} and {width}."""    )

													# Get first result from Stable Diffusion Checkpoint v1.1
													__lowerCamelCase					:		str		   =     self.textaimg_sda_a(
													    prompt=UpperCAmelCase							,					height=UpperCAmelCase							,					width=UpperCAmelCase							,					num_inference_steps=UpperCAmelCase							,					guidance_scale=UpperCAmelCase							,					negative_prompt=UpperCAmelCase							,					num_images_per_prompt=UpperCAmelCase							,					eta=UpperCAmelCase							,					generator=UpperCAmelCase							,					latents=UpperCAmelCase							,					output_type=UpperCAmelCase							,					return_dict=UpperCAmelCase							,					callback=UpperCAmelCase							,					callback_steps=UpperCAmelCase							,					**UpperCAmelCase							,					)

													# Get first result from Stable Diffusion Checkpoint v1.2
													__lowerCamelCase					:		Optional[int]		   =     self.textaimg_sda_a(
													    prompt=UpperCAmelCase							,					height=UpperCAmelCase							,					width=UpperCAmelCase							,					num_inference_steps=UpperCAmelCase							,					guidance_scale=UpperCAmelCase							,					negative_prompt=UpperCAmelCase							,					num_images_per_prompt=UpperCAmelCase							,					eta=UpperCAmelCase							,					generator=UpperCAmelCase							,					latents=UpperCAmelCase							,					output_type=UpperCAmelCase							,					return_dict=UpperCAmelCase							,					callback=UpperCAmelCase							,					callback_steps=UpperCAmelCase							,					**UpperCAmelCase							,					)

													# Get first result from Stable Diffusion Checkpoint v1.3
													__lowerCamelCase					:		Any		   =     self.textaimg_sda_a(
													    prompt=UpperCAmelCase							,					height=UpperCAmelCase							,					width=UpperCAmelCase							,					num_inference_steps=UpperCAmelCase							,					guidance_scale=UpperCAmelCase							,					negative_prompt=UpperCAmelCase							,					num_images_per_prompt=UpperCAmelCase							,					eta=UpperCAmelCase							,					generator=UpperCAmelCase							,					latents=UpperCAmelCase							,					output_type=UpperCAmelCase							,					return_dict=UpperCAmelCase							,					callback=UpperCAmelCase							,					callback_steps=UpperCAmelCase							,					**UpperCAmelCase							,					)

													# Get first result from Stable Diffusion Checkpoint v1.4
													__lowerCamelCase					:		Optional[Any]		   =     self.textaimg_sda_a(
													    prompt=UpperCAmelCase							,					height=UpperCAmelCase							,					width=UpperCAmelCase							,					num_inference_steps=UpperCAmelCase							,					guidance_scale=UpperCAmelCase							,					negative_prompt=UpperCAmelCase							,					num_images_per_prompt=UpperCAmelCase							,					eta=UpperCAmelCase							,					generator=UpperCAmelCase							,					latents=UpperCAmelCase							,					output_type=UpperCAmelCase							,					return_dict=UpperCAmelCase							,					callback=UpperCAmelCase							,					callback_steps=UpperCAmelCase							,					**UpperCAmelCase							,					)

													# Get all result images into a single list and pass it via StableDiffusionPipelineOutput for final result
													return StableDiffusionPipelineOutput([resa[0], resa[0], resa[0], resa[0]]    )