






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__      (	lowerCAmelCase__	,       lowerCAmelCase__	,       unittest.TestCase		):
		UpperCAmelCase_							:       Optional[int] 					= StableDiffusionParadigmsPipeline
		UpperCAmelCase_							:       Tuple 					= TEXT_TO_IMAGE_PARAMS
		UpperCAmelCase_							:       Dict 					= TEXT_TO_IMAGE_BATCH_PARAMS
		UpperCAmelCase_							:       List[Any] 					= TEXT_TO_IMAGE_IMAGE_PARAMS
		UpperCAmelCase_							:       Tuple 					= TEXT_TO_IMAGE_IMAGE_PARAMS






		def 					_lowerCamelCase					(      self						)				->				Optional[int]:
				torch.manual_seed(0						)
				__A    =							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							, attention_head_dim=(2, 4)							, use_linear_projection=lowercase__							, )
				__A    =							DDIMParallelScheduler(
				    beta_start=0.0_0085							, beta_end=0.012							, beta_schedule="scaled_linear"							, clip_sample=lowercase__							, set_alpha_to_one=lowercase__							, )
				torch.manual_seed(0						)
				__A    =							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						)
				__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="gelu"							, projection_dim=512							, )
				__A    =							CLIPTextModel(lowercase__						)
				__A    =							CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"						)

				__A    =							{
				    "unet": unet,
				    "scheduler": scheduler,
				    "vae": vae,
				    "text_encoder": text_encoder,
				    "tokenizer": tokenizer,
				    "safety_checker": None,
				    "feature_extractor": None,
				}
				return components






		def 					_lowerCamelCase					(      self							, lowercase__							, lowercase__=0						)				->				Union[str, Any]:
				if str(lowercase__						).startswith("mps"						):
						__A    =							torch.manual_seed(lowercase__						)
				else:
						__A    =							torch.Generator(device=lowercase__						).manual_seed(lowercase__						)
				__A    =							{
				    "prompt": "a photograph of an astronaut riding a horse",
				    "generator": generator,
				    "num_inference_steps": 10,
				    "guidance_scale": 6.0,
				    "output_type": "numpy",
				    "parallel": 3,
				    "debug": True,
				}
				return inputs






		def 					_lowerCamelCase					(      self						)				->				Tuple:
				__A    =							"cpu"  # ensure determinism for the device-dependent torch.Generator
				__A    =							self.get_dummy_components()
				__A    =							StableDiffusionParadigmsPipeline(**lowercase__						)
				__A    =							sd_pipe.to(lowercase__						)
				sd_pipe.set_progress_bar_config(disable=lowercase__						)

				__A    =							self.get_dummy_inputs(lowercase__						)
				__A    =							sd_pipe(**lowercase__						).images
				__A    =							image[0, -3:, -3:, -1]
				assert image.shape == (1, 64, 64, 3)

				__A    =							np.array([0.4773, 0.5417, 0.4723, 0.4925, 0.5631, 0.4752, 0.5240, 0.4935, 0.5023]						)

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






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

				__A    =							self.get_dummy_inputs(lowercase__						)
				__A    =							sd_pipe(**lowercase__						).images
				__A    =							image[0, -3:, -3:, -1]
				assert image.shape == (1, 64, 64, 3)

				__A    =							np.array([0.3573, 0.4420, 0.4960, 0.4799, 0.3796, 0.3879, 0.4819, 0.4365, 0.4468]						)

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






		def 					_lowerCamelCase					(      self						)				->				Union[str, Any]:
				super().test_inference_batch_consistent(batch_sizes=[1, 2]						)






		def 					_lowerCamelCase					(      self						)				->				List[str]:
				super().test_inference_batch_single_identical(batch_size=2							, expected_max_diff=3e-3						)






		def 					_lowerCamelCase					(      self						)				->				Optional[Any]:
				__A    =							"cpu"  # ensure determinism for the device-dependent torch.Generator
				__A    =							self.get_dummy_components()
				__A    =							StableDiffusionParadigmsPipeline(**lowercase__						)
				__A    =							sd_pipe.to(lowercase__						)
				sd_pipe.set_progress_bar_config(disable=lowercase__						)

				__A    =							self.get_dummy_inputs(lowercase__						)
				__A    =							"french fries"
				__A    =							sd_pipe(**lowercase__							, negative_prompt=lowercase__						)
				__A    =							output.images
				__A    =							image[0, -3:, -3:, -1]

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

				__A    =							np.array([0.4771, 0.5420, 0.4683, 0.4918, 0.5636, 0.4725, 0.5230, 0.4923, 0.5015]						)

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





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






		def 					_lowerCamelCase					(      self						)				->				Optional[int]:
				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()






		def 					_lowerCamelCase					(      self							, lowercase__=0						)				->				Optional[Any]:
				__A    =							torch.Generator(device=lowercase__						).manual_seed(lowercase__						)
				__A    =							{
				    "prompt": "a photograph of an astronaut riding a horse",
				    "generator": generator,
				    "num_inference_steps": 10,
				    "guidance_scale": 7.5,
				    "output_type": "numpy",
				    "parallel": 3,
				    "debug": True,
				}
				return inputs






		def 					_lowerCamelCase					(      self						)				->				Dict:
				__A    =							"stabilityai/stable-diffusion-2-base"
				__A    =							DDIMParallelScheduler.from_pretrained(lowercase__							, subfolder="scheduler"						)
				__A    =							StableDiffusionParadigmsPipeline.from_pretrained(lowercase__							, scheduler=lowercase__							, safety_checker=lowercase__						)
				pipe.to(lowercase__						)
				pipe.set_progress_bar_config(disable=lowercase__						)
				pipe.enable_attention_slicing()

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

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

				__A    =							np.array([0.9622, 0.9602, 0.9748, 0.9591, 0.9630, 0.9691, 0.9661, 0.9631, 0.9741]						)

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



"""simple docstring"""







from typing import List

import PIL
import torch
from PIL import Image

from ...configuration_utils import ConfigMixin
from ...models.modeling_utils import ModelMixin
from ...utils import PIL_INTERPOLATION






class      __a      (       lowerCAmelCase_					,						lowerCAmelCase_      ):







       def __init__(       self		):



          '''simple docstring'''

          super().__init__()

          self.register_buffer("watermark_image"				,torch.zeros((62, 62, 4)		)		)
          _UpperCAmelCase 						=       None







       def  _snake_case      (       self				,lowerCamelCase_				,lowerCamelCase_=None		):



          '''simple docstring'''

          _UpperCAmelCase 						=       images[0].height
          _UpperCAmelCase 						=       images[0].width

          _UpperCAmelCase 						=       sample_size or h

          _UpperCAmelCase 						=       min(h / sample_size				,w / sample_size		)
          _UpperCAmelCase						,				_UpperCAmelCase 						=       (int(h / coef		), int(w / coef		)) if coef < 1 else (h, w)

          _UpperCAmelCase						,				_UpperCAmelCase 						=       1024**2, img_w * img_h
          _UpperCAmelCase 						=       (Sa / Sa) ** 0.5
          _UpperCAmelCase						,				_UpperCAmelCase						,				_UpperCAmelCase 						=       int(K * 62		), img_w - int(14 * K		), img_h - int(14 * K		)

          if self.watermark_image_as_pil is None:
             _UpperCAmelCase 						=       self.watermark_image.to(torch.uinta		).cpu().numpy()
             _UpperCAmelCase 						=       Image.fromarray(lowerCamelCase_				,mode="RGBA"		)
             _UpperCAmelCase 						=       watermark_image

          _UpperCAmelCase 						=       self.watermark_image_as_pil.resize(
              (wm_size, wm_size)				,PIL_INTERPOLATION["bicubic"]				,reducing_gap=lowerCamelCase_		)

          for pil_img in images:
             pil_img.paste(lowerCamelCase_				,box=(wm_x - wm_size, wm_y - wm_size, wm_x, wm_y)				,mask=wm_img.split()[-1]		)

          return images

