






'''simple docstring'''







import gc
import random
import tempfile
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion_safe import StableDiffusionPipelineSafe as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import require_torch_gpu




class   __SCREAMING_SNAKE_CASE       (unittest.TestCase       ):


  """simple docstring"""


  def    _a       (			self	):






         """simple docstring"""


         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()







  @property
  def    _a       (			self	):






         """simple docstring"""


         a_						     =			1
         a_						     =			3
         a_						     =			(32, 32)

         a_						     =			floats_tensor((batch_size, num_channels) + sizes   ,	rng=random.Random(0	)	).to(lowercase_	)
         return image







  @property
  def    _a       (			self	):






         """simple docstring"""


         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   ,	)
         return model







  @property
  def    _a       (			self	):






         """simple docstring"""


         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   ,	)
         return model







  @property
  def    _a       (			self	):






         """simple docstring"""


         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=1_000   ,	)
         return CLIPTextModel(lowercase_	)







  @property
  def    _a       (			self	):






         """simple docstring"""


         def extract(*UpperCamelCase__   ,	**UpperCamelCase__	):




                class   __SCREAMING_SNAKE_CASE       :


                  """simple docstring"""


                  def __init__(			self	):






                         """simple docstring"""


                         a_						     =			torch.ones([0]	)







                  def    _a       (			self   ,	UpperCamelCase__	):






                         """simple docstring"""


                         self.pixel_values.to(lowercase_	)
                         return self





                return Out()

         return extract







  def    _a       (			self	):






         """simple docstring"""


         a_						     =			'cpu'  # ensure determinism for the device-dependent torch.Generator
         a_						     =			self.dummy_cond_unet
         a_						     =			DDIMScheduler(
             beta_start=0.00_085   ,	beta_end=0.012   ,	beta_schedule='scaled_linear'   ,	clip_sample=lowercase_   ,	set_alpha_to_one=lowercase_   ,	)

         a_						     =			self.dummy_vae
         a_						     =			self.dummy_text_encoder
         a_						     =			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'	)

         # make sure here that pndm scheduler skips prk
         a_						     =			StableDiffusionPipeline(
             unet=lowercase_   ,	scheduler=lowercase_   ,	vae=lowercase_   ,	text_encoder=lowercase_   ,	tokenizer=lowercase_   ,	safety_checker=lowercase_   ,	feature_extractor=self.dummy_extractor   ,	)
         a_						     =			sd_pipe.to(lowercase_	)
         sd_pipe.set_progress_bar_config(disable=lowercase_	)

         a_						     =			'A painting of a squirrel eating a burger'

         a_						     =			torch.Generator(device=lowercase_	).manual_seed(0	)
         a_						     =			sd_pipe([prompt]   ,	generator=lowercase_   ,	guidance_scale=6.0   ,	num_inference_steps=2   ,	output_type='np'	)
         a_						     =			output.images

         a_						     =			torch.Generator(device=lowercase_	).manual_seed(0	)
         a_						     =			sd_pipe(
             [prompt]   ,	generator=lowercase_   ,	guidance_scale=6.0   ,	num_inference_steps=2   ,	output_type='np'   ,	return_dict=lowercase_   ,	)[0]

         a_						     =			image[0, -3:, -3:, -1]
         a_						     =			image_from_tuple[0, -3:, -3:, -1]

         assert image.shape == (1, 64, 64, 3)
         a_						     =			np.array([0.5_756, 0.6_118, 0.5_005, 0.5_041, 0.5_471, 0.4_726, 0.4_976, 0.4_865, 0.4_864]	)

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







  def    _a       (			self	):






         """simple docstring"""


         a_						     =			'cpu'  # ensure determinism for the device-dependent torch.Generator
         a_						     =			self.dummy_cond_unet
         a_						     =			PNDMScheduler(skip_prk_steps=lowercase_	)
         a_						     =			self.dummy_vae
         a_						     =			self.dummy_text_encoder
         a_						     =			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'	)

         # make sure here that pndm scheduler skips prk
         a_						     =			StableDiffusionPipeline(
             unet=lowercase_   ,	scheduler=lowercase_   ,	vae=lowercase_   ,	text_encoder=lowercase_   ,	tokenizer=lowercase_   ,	safety_checker=lowercase_   ,	feature_extractor=self.dummy_extractor   ,	)
         a_						     =			sd_pipe.to(lowercase_	)
         sd_pipe.set_progress_bar_config(disable=lowercase_	)

         a_						     =			'A painting of a squirrel eating a burger'
         a_						     =			torch.Generator(device=lowercase_	).manual_seed(0	)
         a_						     =			sd_pipe([prompt]   ,	generator=lowercase_   ,	guidance_scale=6.0   ,	num_inference_steps=2   ,	output_type='np'	)

         a_						     =			output.images

         a_						     =			torch.Generator(device=lowercase_	).manual_seed(0	)
         a_						     =			sd_pipe(
             [prompt]   ,	generator=lowercase_   ,	guidance_scale=6.0   ,	num_inference_steps=2   ,	output_type='np'   ,	return_dict=lowercase_   ,	)[0]

         a_						     =			image[0, -3:, -3:, -1]
         a_						     =			image_from_tuple[0, -3:, -3:, -1]

         assert image.shape == (1, 64, 64, 3)
         a_						     =			np.array([0.5_125, 0.5_716, 0.4_828, 0.5_060, 0.5_650, 0.4_768, 0.5_185, 0.4_895, 0.4_993]	)

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







  def    _a       (			self	):






         """simple docstring"""


         a_						     =			StableDiffusionPipeline.from_pretrained(
             'hf-internal-testing/tiny-stable-diffusion-lms-pipe'   ,	safety_checker=lowercase_	)
         assert isinstance(lowercase_   ,	lowercase_	)
         assert isinstance(pipe.scheduler   ,	lowercase_	)
         assert pipe.safety_checker is None

         a_						     =			pipe('example prompt'   ,	num_inference_steps=2	).images[0]
         assert image is not None

         # check that there's no error when saving a pipeline with one of the models being None
         with tempfile.TemporaryDirectory() as tmpdirname:
                pipe.save_pretrained(lowercase_	)
                a_						     =			StableDiffusionPipeline.from_pretrained(lowercase_	)

         # sanity check that the pipeline still works
         assert pipe.safety_checker is None
         a_						     =			pipe('example prompt'   ,	num_inference_steps=2	).images[0]
         assert image is not None







  @unittest.skipIf(torch_device != 'cuda'   ,	'This test requires a GPU'	)
  def    _a       (			self	):






         """simple docstring"""


         a_						     =			self.dummy_cond_unet
         a_						     =			PNDMScheduler(skip_prk_steps=lowercase_	)
         a_						     =			self.dummy_vae
         a_						     =			self.dummy_text_encoder
         a_						     =			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'	)

         # put models in fp16
         a_						     =			unet.half()
         a_						     =			vae.half()
         a_						     =			bert.half()

         # make sure here that pndm scheduler skips prk
         a_						     =			StableDiffusionPipeline(
             unet=lowercase_   ,	scheduler=lowercase_   ,	vae=lowercase_   ,	text_encoder=lowercase_   ,	tokenizer=lowercase_   ,	safety_checker=lowercase_   ,	feature_extractor=self.dummy_extractor   ,	)
         a_						     =			sd_pipe.to(lowercase_	)
         sd_pipe.set_progress_bar_config(disable=lowercase_	)

         a_						     =			'A painting of a squirrel eating a burger'
         a_						     =			sd_pipe([prompt]   ,	num_inference_steps=2   ,	output_type='np'	).images

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





@nightly
@require_torch_gpu
class   __SCREAMING_SNAKE_CASE       (unittest.TestCase       ):


  """simple docstring"""


  def    _a       (			self	):






         """simple docstring"""


         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()







  def    _a       (			self	):






         """simple docstring"""


         a_						     =			StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'   ,	safety_checker=lowercase_	)
         a_						     =			LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config	)
         a_						     =			sd_pipe.to(lowercase_	)
         sd_pipe.set_progress_bar_config(disable=lowercase_	)

         a_						     =			(
             'portrait of girl with smokey eyes makeup in abandoned hotel, grange clothes, redshift, wide high angle'
             ' coloured polaroid photograph with flash, kodak film, hyper real, stunning moody cinematography, with'
             ' anamorphic lenses, by maripol, fallen angels by wong kar - wai, style of suspiria and neon demon and'
             ' children from bahnhof zoo, detailed '
         )
         a_						     =			4_003_660_346
         a_						     =			7

         # without safety guidance (sld_guidance_scale = 0)
         a_						     =			torch.manual_seed(lowercase_	)
         a_						     =			sd_pipe(
             [prompt]   ,	generator=lowercase_   ,	guidance_scale=lowercase_   ,	num_inference_steps=50   ,	output_type='np'   ,	width=512   ,	height=512   ,	sld_guidance_scale=0   ,	)

         a_						     =			output.images
         a_						     =			image[0, -3:, -3:, -1]
         a_						     =			[0.2_278, 0.2_231, 0.2_249, 0.2_333, 0.2_303, 0.1_885, 0.2_273, 0.2_144, 0.2_176]

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

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

         # without safety guidance (strong configuration)
         a_						     =			torch.manual_seed(lowercase_	)
         a_						     =			sd_pipe(
             [prompt]   ,	generator=lowercase_   ,	guidance_scale=lowercase_   ,	num_inference_steps=50   ,	output_type='np'   ,	width=512   ,	height=512   ,	sld_guidance_scale=2_000   ,	sld_warmup_steps=7   ,	sld_threshold=0.025   ,	sld_momentum_scale=0.5   ,	sld_mom_beta=0.7   ,	)

         a_						     =			output.images
         a_						     =			image[0, -3:, -3:, -1]
         a_						     =			[0.2_383, 0.2_276, 0.236, 0.2_192, 0.2_186, 0.2_053, 0.1_971, 0.1_901, 0.1_719]

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

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







  def    _a       (			self	):






         """simple docstring"""


         a_						     =			StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'   ,	safety_checker=lowercase_	)
         a_						     =			LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config	)
         a_						     =			sd_pipe.to(lowercase_	)
         sd_pipe.set_progress_bar_config(disable=lowercase_	)

         a_						     =			'padme amidala taking a bath artwork, safe for work, no nudity'
         a_						     =			2_734_971_755
         a_						     =			7

         a_						     =			torch.manual_seed(lowercase_	)
         a_						     =			sd_pipe(
             [prompt]   ,	generator=lowercase_   ,	guidance_scale=lowercase_   ,	num_inference_steps=50   ,	output_type='np'   ,	width=512   ,	height=512   ,	sld_guidance_scale=0   ,	)

         a_						     =			output.images
         a_						     =			image[0, -3:, -3:, -1]
         a_						     =			[0.3_502, 0.3_622, 0.3_396, 0.3_642, 0.3_478, 0.3_318, 0.35, 0.3_348, 0.3_297]

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

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

         a_						     =			torch.manual_seed(lowercase_	)
         a_						     =			sd_pipe(
             [prompt]   ,	generator=lowercase_   ,	guidance_scale=lowercase_   ,	num_inference_steps=50   ,	output_type='np'   ,	width=512   ,	height=512   ,	sld_guidance_scale=2_000   ,	sld_warmup_steps=7   ,	sld_threshold=0.025   ,	sld_momentum_scale=0.5   ,	sld_mom_beta=0.7   ,	)

         a_						     =			output.images
         a_						     =			image[0, -3:, -3:, -1]
         a_						     =			[0.5_531, 0.5_206, 0.4_895, 0.5_156, 0.5_182, 0.4_751, 0.4_802, 0.4_803, 0.4_443]

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

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







  def    _a       (			self	):






         """simple docstring"""


         a_						     =			StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'	)
         a_						     =			sd_pipe.to(lowercase_	)
         sd_pipe.set_progress_bar_config(disable=lowercase_	)

         a_						     =			(
             'the four horsewomen of the apocalypse, painting by tom of finland, gaston bussiere, craig mullins, j. c.'
             ' leyendecker'
         )
         a_						     =			1_044_355_234
         a_						     =			12

         a_						     =			torch.manual_seed(lowercase_	)
         a_						     =			sd_pipe(
             [prompt]   ,	generator=lowercase_   ,	guidance_scale=lowercase_   ,	num_inference_steps=50   ,	output_type='np'   ,	width=512   ,	height=512   ,	sld_guidance_scale=0   ,	)

         a_						     =			output.images
         a_						     =			image[0, -3:, -3:, -1]
         a_						     =			np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]	)

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

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

         a_						     =			torch.manual_seed(lowercase_	)
         a_						     =			sd_pipe(
             [prompt]   ,	generator=lowercase_   ,	guidance_scale=lowercase_   ,	num_inference_steps=50   ,	output_type='np'   ,	width=512   ,	height=512   ,	sld_guidance_scale=2_000   ,	sld_warmup_steps=7   ,	sld_threshold=0.025   ,	sld_momentum_scale=0.5   ,	sld_mom_beta=0.7   ,	)

         a_						     =			output.images
         a_						     =			image[0, -3:, -3:, -1]
         a_						     =			np.array([0.5_818, 0.6_285, 0.6_835, 0.6_019, 0.625, 0.6_754, 0.6_096, 0.6_334, 0.6_561]	)
         assert image.shape == (1, 512, 512, 3)

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





import enum
import os
from hashlib import shaaaa
from typing import Optional

from .. import config
from .logging import get_logger


lowerCamelCase__  =					get_logger(__name__)




class 					_UpperCAmelCase			(     enum.Enum  ):



 '''simple docstring'''




 __A			       =					'''all_checks'''
 __A			       =					'''basic_checks'''
 __A			       =					'''no_checks'''




class 					_UpperCAmelCase			(     lowerCAmelCase  ):



 '''simple docstring'''








class 					_UpperCAmelCase			(     lowerCAmelCase  ):



 '''simple docstring'''








class 					_UpperCAmelCase			(     lowerCAmelCase  ):



 '''simple docstring'''








class 					_UpperCAmelCase			(     lowerCAmelCase  ):



 '''simple docstring'''






def 				lowerCAmelCase__       (			a__    ,     a__    ,     a__=None )      ->int:




  '''simple docstring'''

  if expected_checksums is None:
    logger.info("Unable to verify checksums." )
    return
  if len(set(a__ ) - set(a__ ) ) > 0:
    raise ExpectedMoreDownloadedFiles(str(set(a__ ) - set(a__ ) ) )
  if len(set(a__ ) - set(a__ ) ) > 0:
    raise UnexpectedDownloadedFile(str(set(a__ ) - set(a__ ) ) )
  _UpperCamelCase           =    [url for url in expected_checksums if expected_checksums[url] != recorded_checksums[url]]
  _UpperCamelCase           =    " for " + verification_name if verification_name is not None else ""
  if len(a__ ) > 0:
    raise NonMatchingChecksumError(
        f'Checksums didn\'t match{for_verification_name}:\n'
        f'{bad_urls}\n'
        "Set `verification_mode='no_checks'` to skip checksums verification and ignore this error" )
  logger.info("All the checksums matched successfully" + for_verification_name )




class 					_UpperCAmelCase			(     lowerCAmelCase  ):



 '''simple docstring'''








class 					_UpperCAmelCase			(     lowerCAmelCase  ):



 '''simple docstring'''








class 					_UpperCAmelCase			(     lowerCAmelCase  ):



 '''simple docstring'''








class 					_UpperCAmelCase			(     lowerCAmelCase  ):



 '''simple docstring'''






def 				lowerCAmelCase__       (			a__    ,     a__ )      ->Dict:




  '''simple docstring'''

  if expected_splits is None:
    logger.info("Unable to verify splits sizes." )
    return
  if len(set(a__ ) - set(a__ ) ) > 0:
    raise ExpectedMoreSplits(str(set(a__ ) - set(a__ ) ) )
  if len(set(a__ ) - set(a__ ) ) > 0:
    raise UnexpectedSplits(str(set(a__ ) - set(a__ ) ) )
  _UpperCamelCase           =    [
      {"expected": expected_splits[name], "recorded": recorded_splits[name]}
      for name in expected_splits
      if expected_splits[name].num_examples != recorded_splits[name].num_examples
  ]
  if len(a__ ) > 0:
    raise NonMatchingSplitsSizesError(str(a__ ) )
  logger.info("All the splits matched successfully." )


def 				lowerCAmelCase__       (			a__    ,     a__ = True )      ->dict:




  '''simple docstring'''

  if record_checksum:
    _UpperCamelCase           =    shaaaa()
    with open(a__    ,     "rb" ) as f:
      for chunk in iter(lambda: f.read(1 << 20 )    ,     B"" ):
        m.update(a__ )
      _UpperCamelCase           =    m.hexdigest()
  else:
    _UpperCamelCase           =    None
  return {"num_bytes": os.path.getsize(a__ ), "checksum": checksum}







def 				lowerCAmelCase__       (			a__ )      ->Any:




  '''simple docstring'''

  if dataset_size and config.IN_MEMORY_MAX_SIZE:
    return dataset_size < config.IN_MEMORY_MAX_SIZE
  else:
    return False




