"""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 		__A							(				unittest.TestCase	):






     def 					__A						(	self						):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()






     @property
     def 					__A						(	self						):
       _lowerCAmelCase		: Optional[int]						    =  1
       _lowerCAmelCase		: Dict						    =  3
       _lowerCAmelCase		: str						    =  (32, 32)

       _lowerCAmelCase		: Tuple						    =  floats_tensor((batch_size, num_channels) + sizes       ,  rng=random.Random(0						)						).to(a__						)
       return image






     @property
     def 					__A						(	self						):
       torch.manual_seed(0						)
       _lowerCAmelCase		: Optional[int]						    =  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						):
       torch.manual_seed(0						)
       _lowerCAmelCase		: str						    =  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						):
       torch.manual_seed(0						)
       _lowerCAmelCase		: Union[str, 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       ,  )
       return CLIPTextModel(a__						)






     @property
     def 					__A						(	self						):
       def extract(*a__       ,  **a__						):
         class 		__A							:






              def __init__(	self						):
                _lowerCAmelCase		: Dict						    =  torch.ones([0]						)






              def 					__A						(	self       ,  a__						):
                self.pixel_values.to(a__						)
                return self







         return Out()

       return extract






     def 					__A						(	self						):
       _lowerCAmelCase		: Dict						    =  """cpu"""  # ensure determinism for the device-dependent torch.Generator
       _lowerCAmelCase		: Any						    =  self.dummy_cond_unet
       _lowerCAmelCase		: List[Any]						    =  DDIMScheduler(
           beta_start=0.0_0_0_8_5       ,  beta_end=0.0_1_2       ,  beta_schedule="""scaled_linear"""       ,  clip_sample=a__       ,  set_alpha_to_one=a__       ,  )

       _lowerCAmelCase		: Optional[int]						    =  self.dummy_vae
       _lowerCAmelCase		: Optional[Any]						    =  self.dummy_text_encoder
       _lowerCAmelCase		: List[Any]						    =  CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

       # make sure here that pndm scheduler skips prk
       _lowerCAmelCase		: List[Any]						    =  StableDiffusionPipeline(
           unet=a__       ,  scheduler=a__       ,  vae=a__       ,  text_encoder=a__       ,  tokenizer=a__       ,  safety_checker=a__       ,  feature_extractor=self.dummy_extractor       ,  )
       _lowerCAmelCase		: str						    =  sd_pipe.to(a__						)
       sd_pipe.set_progress_bar_config(disable=a__						)

       _lowerCAmelCase		: Tuple						    =  """A painting of a squirrel eating a burger"""

       _lowerCAmelCase		: List[str]						    =  torch.Generator(device=a__						).manual_seed(0						)
       _lowerCAmelCase		: Dict						    =  sd_pipe([prompt]       ,  generator=a__       ,  guidance_scale=6.0       ,  num_inference_steps=2       ,  output_type="""np"""						)
       _lowerCAmelCase		: Union[str, Any]						    =  output.images

       _lowerCAmelCase		: Tuple						    =  torch.Generator(device=a__						).manual_seed(0						)
       _lowerCAmelCase		: Dict						    =  sd_pipe(
           [prompt]       ,  generator=a__       ,  guidance_scale=6.0       ,  num_inference_steps=2       ,  output_type="""np"""       ,  return_dict=a__       ,  )[0]

       _lowerCAmelCase		: List[str]						    =  image[0, -3:, -3:, -1]
       _lowerCAmelCase		: int						    =  image_from_tuple[0, -3:, -3:, -1]

       assert image.shape == (1, 64, 64, 3)
       _lowerCAmelCase		: Any						    =  np.array([0.5_7_5_6, 0.6_1_1_8, 0.5_0_0_5, 0.5_0_4_1, 0.5_4_7_1, 0.4_7_2_6, 0.4_9_7_6, 0.4_8_6_5, 0.4_8_6_4]						)

       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						):
       _lowerCAmelCase		: Any						    =  """cpu"""  # ensure determinism for the device-dependent torch.Generator
       _lowerCAmelCase		: int						    =  self.dummy_cond_unet
       _lowerCAmelCase		: List[str]						    =  PNDMScheduler(skip_prk_steps=a__						)
       _lowerCAmelCase		: str						    =  self.dummy_vae
       _lowerCAmelCase		: Union[str, Any]						    =  self.dummy_text_encoder
       _lowerCAmelCase		: List[str]						    =  CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

       # make sure here that pndm scheduler skips prk
       _lowerCAmelCase		: str						    =  StableDiffusionPipeline(
           unet=a__       ,  scheduler=a__       ,  vae=a__       ,  text_encoder=a__       ,  tokenizer=a__       ,  safety_checker=a__       ,  feature_extractor=self.dummy_extractor       ,  )
       _lowerCAmelCase		: List[str]						    =  sd_pipe.to(a__						)
       sd_pipe.set_progress_bar_config(disable=a__						)

       _lowerCAmelCase		: Union[str, Any]						    =  """A painting of a squirrel eating a burger"""
       _lowerCAmelCase		: Optional[Any]						    =  torch.Generator(device=a__						).manual_seed(0						)
       _lowerCAmelCase		: Optional[int]						    =  sd_pipe([prompt]       ,  generator=a__       ,  guidance_scale=6.0       ,  num_inference_steps=2       ,  output_type="""np"""						)

       _lowerCAmelCase		: Dict						    =  output.images

       _lowerCAmelCase		: Optional[int]						    =  torch.Generator(device=a__						).manual_seed(0						)
       _lowerCAmelCase		: List[Any]						    =  sd_pipe(
           [prompt]       ,  generator=a__       ,  guidance_scale=6.0       ,  num_inference_steps=2       ,  output_type="""np"""       ,  return_dict=a__       ,  )[0]

       _lowerCAmelCase		: str						    =  image[0, -3:, -3:, -1]
       _lowerCAmelCase		: Dict						    =  image_from_tuple[0, -3:, -3:, -1]

       assert image.shape == (1, 64, 64, 3)
       _lowerCAmelCase		: Any						    =  np.array([0.5_1_2_5, 0.5_7_1_6, 0.4_8_2_8, 0.5_0_6_0, 0.5_6_5_0, 0.4_7_6_8, 0.5_1_8_5, 0.4_8_9_5, 0.4_9_9_3]						)

       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						):
       _lowerCAmelCase		: Union[str, Any]						    =  StableDiffusionPipeline.from_pretrained(
           """hf-internal-testing/tiny-stable-diffusion-lms-pipe"""       ,  safety_checker=a__						)
       assert isinstance(a__       ,  a__						)
       assert isinstance(pipe.scheduler       ,  a__						)
       assert pipe.safety_checker is None

       _lowerCAmelCase		: List[str]						    =  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(a__						)
         _lowerCAmelCase		: Tuple						    =  StableDiffusionPipeline.from_pretrained(a__						)

       # sanity check that the pipeline still works
       assert pipe.safety_checker is None
       _lowerCAmelCase		: List[Any]						    =  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						):
       _lowerCAmelCase		: Union[str, Any]						    =  self.dummy_cond_unet
       _lowerCAmelCase		: str						    =  PNDMScheduler(skip_prk_steps=a__						)
       _lowerCAmelCase		: List[Any]						    =  self.dummy_vae
       _lowerCAmelCase		: Optional[Any]						    =  self.dummy_text_encoder
       _lowerCAmelCase		: Union[str, Any]						    =  CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

       # put models in fp16
       _lowerCAmelCase		: Dict						    =  unet.half()
       _lowerCAmelCase		: List[str]						    =  vae.half()
       _lowerCAmelCase		: Optional[Any]						    =  bert.half()

       # make sure here that pndm scheduler skips prk
       _lowerCAmelCase		: List[str]						    =  StableDiffusionPipeline(
           unet=a__       ,  scheduler=a__       ,  vae=a__       ,  text_encoder=a__       ,  tokenizer=a__       ,  safety_checker=a__       ,  feature_extractor=self.dummy_extractor       ,  )
       _lowerCAmelCase		: Optional[Any]						    =  sd_pipe.to(a__						)
       sd_pipe.set_progress_bar_config(disable=a__						)

       _lowerCAmelCase		: Tuple						    =  """A painting of a squirrel eating a burger"""
       _lowerCAmelCase		: Optional[Any]						    =  sd_pipe([prompt]       ,  num_inference_steps=2       ,  output_type="""np"""						).images

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







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






     def 					__A						(	self						):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()






     def 					__A						(	self						):
       _lowerCAmelCase		: Any						    =  StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""       ,  safety_checker=a__						)
       _lowerCAmelCase		: List[str]						    =  LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config						)
       _lowerCAmelCase		: str						    =  sd_pipe.to(a__						)
       sd_pipe.set_progress_bar_config(disable=a__						)

       _lowerCAmelCase		: str						    =  (
           """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 """
       )
       _lowerCAmelCase		: Optional[Any]						    =  4003660346
       _lowerCAmelCase		: Union[str, Any]						    =  7

       # without safety guidance (sld_guidance_scale = 0)
       _lowerCAmelCase		: Union[str, Any]						    =  torch.manual_seed(a__						)
       _lowerCAmelCase		: Tuple						    =  sd_pipe(
           [prompt]       ,  generator=a__       ,  guidance_scale=a__       ,  num_inference_steps=50       ,  output_type="""np"""       ,  width=512       ,  height=512       ,  sld_guidance_scale=0       ,  )

       _lowerCAmelCase		: int						    =  output.images
       _lowerCAmelCase		: Tuple						    =  image[0, -3:, -3:, -1]
       _lowerCAmelCase		: List[Any]						    =  [0.2_2_7_8, 0.2_2_3_1, 0.2_2_4_9, 0.2_3_3_3, 0.2_3_0_3, 0.1_8_8_5, 0.2_2_7_3, 0.2_1_4_4, 0.2_1_7_6]

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

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

       # without safety guidance (strong configuration)
       _lowerCAmelCase		: Union[str, Any]						    =  torch.manual_seed(a__						)
       _lowerCAmelCase		: Optional[int]						    =  sd_pipe(
           [prompt]       ,  generator=a__       ,  guidance_scale=a__       ,  num_inference_steps=50       ,  output_type="""np"""       ,  width=512       ,  height=512       ,  sld_guidance_scale=2000       ,  sld_warmup_steps=7       ,  sld_threshold=0.0_2_5       ,  sld_momentum_scale=0.5       ,  sld_mom_beta=0.7       ,  )

       _lowerCAmelCase		: Optional[Any]						    =  output.images
       _lowerCAmelCase		: Dict						    =  image[0, -3:, -3:, -1]
       _lowerCAmelCase		: Any						    =  [0.2_3_8_3, 0.2_2_7_6, 0.2_3_6, 0.2_1_9_2, 0.2_1_8_6, 0.2_0_5_3, 0.1_9_7_1, 0.1_9_0_1, 0.1_7_1_9]

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

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






     def 					__A						(	self						):
       _lowerCAmelCase		: Optional[Any]						    =  StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""       ,  safety_checker=a__						)
       _lowerCAmelCase		: List[str]						    =  LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config						)
       _lowerCAmelCase		: Optional[Any]						    =  sd_pipe.to(a__						)
       sd_pipe.set_progress_bar_config(disable=a__						)

       _lowerCAmelCase		: Optional[Any]						    =  """padme amidala taking a bath artwork, safe for work, no nudity"""
       _lowerCAmelCase		: int						    =  2734971755
       _lowerCAmelCase		: Optional[Any]						    =  7

       _lowerCAmelCase		: Tuple						    =  torch.manual_seed(a__						)
       _lowerCAmelCase		: Dict						    =  sd_pipe(
           [prompt]       ,  generator=a__       ,  guidance_scale=a__       ,  num_inference_steps=50       ,  output_type="""np"""       ,  width=512       ,  height=512       ,  sld_guidance_scale=0       ,  )

       _lowerCAmelCase		: Dict						    =  output.images
       _lowerCAmelCase		: Dict						    =  image[0, -3:, -3:, -1]
       _lowerCAmelCase		: Optional[Any]						    =  [0.3_5_0_2, 0.3_6_2_2, 0.3_3_9_6, 0.3_6_4_2, 0.3_4_7_8, 0.3_3_1_8, 0.3_5, 0.3_3_4_8, 0.3_2_9_7]

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

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

       _lowerCAmelCase		: Optional[Any]						    =  torch.manual_seed(a__						)
       _lowerCAmelCase		: str						    =  sd_pipe(
           [prompt]       ,  generator=a__       ,  guidance_scale=a__       ,  num_inference_steps=50       ,  output_type="""np"""       ,  width=512       ,  height=512       ,  sld_guidance_scale=2000       ,  sld_warmup_steps=7       ,  sld_threshold=0.0_2_5       ,  sld_momentum_scale=0.5       ,  sld_mom_beta=0.7       ,  )

       _lowerCAmelCase		: Optional[int]						    =  output.images
       _lowerCAmelCase		: List[str]						    =  image[0, -3:, -3:, -1]
       _lowerCAmelCase		: Dict						    =  [0.5_5_3_1, 0.5_2_0_6, 0.4_8_9_5, 0.5_1_5_6, 0.5_1_8_2, 0.4_7_5_1, 0.4_8_0_2, 0.4_8_0_3, 0.4_4_4_3]

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

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






     def 					__A						(	self						):
       _lowerCAmelCase		: Optional[Any]						    =  StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""						)
       _lowerCAmelCase		: Optional[int]						    =  sd_pipe.to(a__						)
       sd_pipe.set_progress_bar_config(disable=a__						)

       _lowerCAmelCase		: Optional[Any]						    =  (
           """the four horsewomen of the apocalypse, painting by tom of finland, gaston bussiere, craig mullins, j. c."""
           """ leyendecker"""
       )
       _lowerCAmelCase		: str						    =  1044355234
       _lowerCAmelCase		: int						    =  12

       _lowerCAmelCase		: Optional[int]						    =  torch.manual_seed(a__						)
       _lowerCAmelCase		: Tuple						    =  sd_pipe(
           [prompt]       ,  generator=a__       ,  guidance_scale=a__       ,  num_inference_steps=50       ,  output_type="""np"""       ,  width=512       ,  height=512       ,  sld_guidance_scale=0       ,  )

       _lowerCAmelCase		: Tuple						    =  output.images
       _lowerCAmelCase		: List[str]						    =  image[0, -3:, -3:, -1]
       _lowerCAmelCase		: Tuple						    =  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

       _lowerCAmelCase		: Optional[int]						    =  torch.manual_seed(a__						)
       _lowerCAmelCase		: List[str]						    =  sd_pipe(
           [prompt]       ,  generator=a__       ,  guidance_scale=a__       ,  num_inference_steps=50       ,  output_type="""np"""       ,  width=512       ,  height=512       ,  sld_guidance_scale=2000       ,  sld_warmup_steps=7       ,  sld_threshold=0.0_2_5       ,  sld_momentum_scale=0.5       ,  sld_mom_beta=0.7       ,  )

       _lowerCAmelCase		: Optional[int]						    =  output.images
       _lowerCAmelCase		: List[Any]						    =  image[0, -3:, -3:, -1]
       _lowerCAmelCase		: List[Any]						    =  np.array([0.5_8_1_8, 0.6_2_8_5, 0.6_8_3_5, 0.6_0_1_9, 0.6_2_5, 0.6_7_5_4, 0.6_0_9_6, 0.6_3_3_4, 0.6_5_6_1]						)
       assert image.shape == (1, 512, 512, 3)

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





"""simple docstring"""




import pickle
import unittest

import torch

from accelerate import Accelerator
from accelerate.state import AcceleratorState
from accelerate.test_utils import require_cpu
@require_cpu
class 		__A							(				unittest.TestCase	):






     def 					__A						(	self						):
       _lowerCAmelCase		: str						    =  torch.nn.Linear(10       ,  10						)
       _lowerCAmelCase		: Optional[Any]						    =  torch.optim.SGD(model.parameters()       ,  0.1						)
       _lowerCAmelCase		: Optional[Any]						    =  Accelerator()
       _lowerCAmelCase		: Tuple						    =  accelerator.prepare(a__						)
       try:
         pickle.loads(pickle.dumps(a__						)						)
       except Exception as e:
         self.fail(F"Accelerated optimizer pickling failed with {e}"						)
       AcceleratorState._reset_state()




