



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.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()






class       lowerCamelCase_					(       unittest.TestCase  ):
  """simple docstring"""




  def 		_lowercase  (				self   :			Tuple     )      ->      Union[str, Any]:
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()




  @property
  def 		_lowercase  (				self   :			Optional[Any]     )      ->      List[str]:
    __lowerCamelCase : List[str]					 =				1
    __lowerCamelCase : str					 =				3
    __lowerCamelCase : Tuple					 =				(32, 32)

    __lowerCamelCase : Dict					 =				floats_tensor((batch_size, num_channels) + sizes    ,			rng=random.Random(0     )     ).to(_a     )
    return image




  @property
  def 		_lowercase  (				self   :			Optional[int]     )      ->      Union[str, Any]:
    torch.manual_seed(0     )
    __lowerCamelCase : 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    ,			)
    return model




  @property
  def 		_lowercase  (				self   :			Optional[int]     )      ->      str:
    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    ,			)
    return model




  @property
  def 		_lowercase  (				self   :			Dict     )      ->      Any:
    torch.manual_seed(0     )
    __lowerCamelCase : 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 		_lowercase  (				self   :			Dict     )      ->      Union[str, Any]:



    def extract(*_a   :			Dict    ,			**_a   :			Union[str, Any]     ):






      class       lowerCamelCase_					:
        """simple docstring"""




        def __init__(				self   :			str     )      ->      Optional[int]:
          __lowerCamelCase : Tuple					 =				torch.ones([0]     )




        def 		_lowercase  (				self   :			Any    ,			_a   :			Any     )      ->      Union[str, Any]:
          self.pixel_values.to(_a     )
          return self




      return Out()

    return extract




  def 		_lowercase  (				self   :			int     )      ->      List[Any]:
    __lowerCamelCase : List[Any]					 =				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
    __lowerCamelCase : Optional[int]					 =				self.dummy_cond_unet
    __lowerCamelCase : Tuple					 =				DDIMScheduler(
        beta_start=0.00085    ,			beta_end=0.012    ,			beta_schedule='scaled_linear'    ,			clip_sample=_a    ,			set_alpha_to_one=_a    ,			)

    __lowerCamelCase : List[Any]					 =				self.dummy_vae
    __lowerCamelCase : str					 =				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 : Union[str, Any]					 =				sd_pipe.to(_a     )
    sd_pipe.set_progress_bar_config(disable=_a     )

    __lowerCamelCase : Dict					 =				'''A painting of a squirrel eating a burger'''

    __lowerCamelCase : List[str]					 =				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'     )
    __lowerCamelCase : Tuple					 =				output.images

    __lowerCamelCase : Optional[int]					 =				torch.Generator(device=_a     ).manual_seed(0     )
    __lowerCamelCase : Optional[Any]					 =				sd_pipe(
        [prompt]    ,			generator=_a    ,			guidance_scale=6.0    ,			num_inference_steps=2    ,			output_type='np'    ,			return_dict=_a    ,			)[0]

    __lowerCamelCase : int					 =				image[0, -3:, -3:, -1]
    __lowerCamelCase : Optional[Any]					 =				image_from_tuple[0, -3:, -3:, -1]

    assert image.shape == (1, 64, 64, 3)
    __lowerCamelCase : List[Any]					 =				np.array([0.5753, 0.6114, 0.5001, 0.5034, 0.5470, 0.4729, 0.4971, 0.4867, 0.4867]     )

    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 		_lowercase  (				self   :			str     )      ->      Union[str, Any]:
    __lowerCamelCase : str					 =				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
    __lowerCamelCase : str					 =				self.dummy_cond_unet
    __lowerCamelCase : Optional[int]					 =				PNDMScheduler(skip_prk_steps=_a     )
    __lowerCamelCase : Union[str, Any]					 =				self.dummy_vae
    __lowerCamelCase : List[str]					 =				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 : Optional[Any]					 =				StableDiffusionPipeline(
        unet=_a    ,			scheduler=_a    ,			vae=_a    ,			text_encoder=_a    ,			tokenizer=_a    ,			safety_checker=_a    ,			feature_extractor=self.dummy_extractor    ,			)
    __lowerCamelCase : int					 =				sd_pipe.to(_a     )
    sd_pipe.set_progress_bar_config(disable=_a     )

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

    __lowerCamelCase : str					 =				output.images

    __lowerCamelCase : Any					 =				torch.Generator(device=_a     ).manual_seed(0     )
    __lowerCamelCase : int					 =				sd_pipe(
        [prompt]    ,			generator=_a    ,			guidance_scale=6.0    ,			num_inference_steps=2    ,			output_type='np'    ,			return_dict=_a    ,			)[0]

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

    assert image.shape == (1, 64, 64, 3)
    __lowerCamelCase : str					 =				np.array([0.5122, 0.5712, 0.4825, 0.5053, 0.5646, 0.4769, 0.5179, 0.4894, 0.4994]     )

    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 		_lowercase  (				self   :			str     )      ->      Tuple:
    __lowerCamelCase : str					 =				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 : List[Any]					 =				StableDiffusionPipeline.from_pretrained(_a     )

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

    # put models in fp16
    __lowerCamelCase : Optional[Any]					 =				unet.half()
    __lowerCamelCase : Optional[int]					 =				vae.half()
    __lowerCamelCase : Optional[int]					 =				bert.half()

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

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

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






@nightly
@require_torch_gpu
class       lowerCamelCase_					(       unittest.TestCase  ):
  """simple docstring"""




  def 		_lowercase  (				self   :			Dict     )      ->      Union[str, Any]:
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()




  def 		_lowercase  (				self   :			Dict     )      ->      Union[str, Any]:
    __lowerCamelCase : Optional[int]					 =				'''cuda'''
    __lowerCamelCase : Optional[int]					 =				StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'     )
    __lowerCamelCase : Tuple					 =				pipe.to(_a     )
    pipe.set_progress_bar_config(disable=_a     )

    __lowerCamelCase : List[Any]					 =				'''a photo of a cat'''
    __lowerCamelCase : Optional[int]					 =				{
        '''editing_prompt''': ['''sunglasses'''],
        '''reverse_editing_direction''': [False],
        '''edit_warmup_steps''': 10,
        '''edit_guidance_scale''': 6,
        '''edit_threshold''': 0.95,
        '''edit_momentum_scale''': 0.5,
        '''edit_mom_beta''': 0.6,
    }

    __lowerCamelCase : int					 =				3
    __lowerCamelCase : Union[str, Any]					 =				7

    # no sega enabled
    __lowerCamelCase : str					 =				torch.Generator(_a     )
    generator.manual_seed(_a     )
    __lowerCamelCase : int					 =				pipe(
        [prompt]    ,			generator=_a    ,			guidance_scale=_a    ,			num_inference_steps=50    ,			output_type='np'    ,			width=512    ,			height=512    ,			)

    __lowerCamelCase : Optional[int]					 =				output.images
    __lowerCamelCase : List[str]					 =				image[0, -3:, -3:, -1]
    __lowerCamelCase : Tuple					 =				[
        0.34673113,
        0.38492733,
        0.37597352,
        0.34086335,
        0.35650748,
        0.35579205,
        0.3384763,
        0.34340236,
        0.3573271,
    ]

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

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

    # with sega enabled
    # generator = torch.manual_seed(seed)
    generator.manual_seed(_a     )
    __lowerCamelCase : str					 =				pipe(
        [prompt]    ,			generator=_a    ,			guidance_scale=_a    ,			num_inference_steps=50    ,			output_type='np'    ,			width=512    ,			height=512    ,			**_a    ,			)

    __lowerCamelCase : List[Any]					 =				output.images
    __lowerCamelCase : List[Any]					 =				image[0, -3:, -3:, -1]
    __lowerCamelCase : Any					 =				[
        0.41887826,
        0.37728766,
        0.30138272,
        0.41416335,
        0.41664985,
        0.36283392,
        0.36191246,
        0.43364465,
        0.43001732,
    ]

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

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




  def 		_lowercase  (				self   :			Dict     )      ->      Tuple:
    __lowerCamelCase : str					 =				'''cuda'''
    __lowerCamelCase : Any					 =				StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'     )
    __lowerCamelCase : Optional[int]					 =				pipe.to(_a     )
    pipe.set_progress_bar_config(disable=_a     )

    __lowerCamelCase : int					 =				'''an image of a crowded boulevard, realistic, 4k'''
    __lowerCamelCase : int					 =				{
        '''editing_prompt''': '''crowd, crowded, people''',
        '''reverse_editing_direction''': True,
        '''edit_warmup_steps''': 10,
        '''edit_guidance_scale''': 8.3,
        '''edit_threshold''': 0.9,
        '''edit_momentum_scale''': 0.5,
        '''edit_mom_beta''': 0.6,
    }

    __lowerCamelCase : List[Any]					 =				9
    __lowerCamelCase : Any					 =				7

    # no sega enabled
    __lowerCamelCase : List[Any]					 =				torch.Generator(_a     )
    generator.manual_seed(_a     )
    __lowerCamelCase : Dict					 =				pipe(
        [prompt]    ,			generator=_a    ,			guidance_scale=_a    ,			num_inference_steps=50    ,			output_type='np'    ,			width=512    ,			height=512    ,			)

    __lowerCamelCase : Optional[Any]					 =				output.images
    __lowerCamelCase : Any					 =				image[0, -3:, -3:, -1]
    __lowerCamelCase : Union[str, Any]					 =				[
        0.43497998,
        0.91814065,
        0.7540739,
        0.55580205,
        0.8467265,
        0.5389691,
        0.62574506,
        0.58897763,
        0.50926757,
    ]

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

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

    # with sega enabled
    # generator = torch.manual_seed(seed)
    generator.manual_seed(_a     )
    __lowerCamelCase : Any					 =				pipe(
        [prompt]    ,			generator=_a    ,			guidance_scale=_a    ,			num_inference_steps=50    ,			output_type='np'    ,			width=512    ,			height=512    ,			**_a    ,			)

    __lowerCamelCase : Optional[int]					 =				output.images
    __lowerCamelCase : Dict					 =				image[0, -3:, -3:, -1]
    __lowerCamelCase : Optional[int]					 =				[
        0.3089719,
        0.30500144,
        0.29016042,
        0.30630964,
        0.325687,
        0.29419225,
        0.2908091,
        0.28723598,
        0.27696294,
    ]

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

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




  def 		_lowercase  (				self   :			Tuple     )      ->      Dict:
    __lowerCamelCase : List[str]					 =				'''cuda'''
    __lowerCamelCase : List[Any]					 =				StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'     )
    __lowerCamelCase : Tuple					 =				pipe.to(_a     )
    pipe.set_progress_bar_config(disable=_a     )

    __lowerCamelCase : Dict					 =				'''a castle next to a river'''
    __lowerCamelCase : Optional[Any]					 =				{
        '''editing_prompt''': ['''boat on a river, boat''', '''monet, impression, sunrise'''],
        '''reverse_editing_direction''': False,
        '''edit_warmup_steps''': [15, 18],
        '''edit_guidance_scale''': 6,
        '''edit_threshold''': [0.9, 0.8],
        '''edit_momentum_scale''': 0.5,
        '''edit_mom_beta''': 0.6,
    }

    __lowerCamelCase : Union[str, Any]					 =				48
    __lowerCamelCase : List[str]					 =				7

    # no sega enabled
    __lowerCamelCase : List[str]					 =				torch.Generator(_a     )
    generator.manual_seed(_a     )
    __lowerCamelCase : Union[str, Any]					 =				pipe(
        [prompt]    ,			generator=_a    ,			guidance_scale=_a    ,			num_inference_steps=50    ,			output_type='np'    ,			width=512    ,			height=512    ,			)

    __lowerCamelCase : List[str]					 =				output.images
    __lowerCamelCase : Optional[int]					 =				image[0, -3:, -3:, -1]
    __lowerCamelCase : Tuple					 =				[
        0.75163555,
        0.76037145,
        0.61785,
        0.9189673,
        0.8627701,
        0.85189694,
        0.8512813,
        0.87012076,
        0.8312857,
    ]

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

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

    # with sega enabled
    # generator = torch.manual_seed(seed)
    generator.manual_seed(_a     )
    __lowerCamelCase : int					 =				pipe(
        [prompt]    ,			generator=_a    ,			guidance_scale=_a    ,			num_inference_steps=50    ,			output_type='np'    ,			width=512    ,			height=512    ,			**_a    ,			)

    __lowerCamelCase : List[str]					 =				output.images
    __lowerCamelCase : str					 =				image[0, -3:, -3:, -1]
    __lowerCamelCase : List[Any]					 =				[
        0.73553365,
        0.7537271,
        0.74341905,
        0.66480356,
        0.6472925,
        0.63039416,
        0.64812905,
        0.6749717,
        0.6517102,
    ]

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

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




  def 		_lowercase  (				self   :			str     )      ->      int:
    __lowerCamelCase : Dict					 =				'''cuda'''
    __lowerCamelCase : Any					 =				StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'    ,			torch_dtype=torch.floataa     )
    __lowerCamelCase : Tuple					 =				pipe.to(_a     )
    pipe.set_progress_bar_config(disable=_a     )

    __lowerCamelCase : Optional[Any]					 =				'''a photo of a cat'''
    __lowerCamelCase : Union[str, Any]					 =				{
        '''editing_prompt''': ['''sunglasses'''],
        '''reverse_editing_direction''': [False],
        '''edit_warmup_steps''': 10,
        '''edit_guidance_scale''': 6,
        '''edit_threshold''': 0.95,
        '''edit_momentum_scale''': 0.5,
        '''edit_mom_beta''': 0.6,
    }

    __lowerCamelCase : Optional[int]					 =				3
    __lowerCamelCase : Optional[Any]					 =				7

    # no sega enabled
    __lowerCamelCase : Union[str, Any]					 =				torch.Generator(_a     )
    generator.manual_seed(_a     )
    __lowerCamelCase : Optional[int]					 =				pipe(
        [prompt]    ,			generator=_a    ,			guidance_scale=_a    ,			num_inference_steps=50    ,			output_type='np'    ,			width=512    ,			height=512    ,			)

    __lowerCamelCase : List[Any]					 =				output.images
    __lowerCamelCase : Dict					 =				image[0, -3:, -3:, -1]
    __lowerCamelCase : Dict					 =				[
        0.34887695,
        0.3876953,
        0.375,
        0.34423828,
        0.3581543,
        0.35717773,
        0.3383789,
        0.34570312,
        0.359375,
    ]

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

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

    # with sega enabled
    # generator = torch.manual_seed(seed)
    generator.manual_seed(_a     )
    __lowerCamelCase : List[Any]					 =				pipe(
        [prompt]    ,			generator=_a    ,			guidance_scale=_a    ,			num_inference_steps=50    ,			output_type='np'    ,			width=512    ,			height=512    ,			**_a    ,			)

    __lowerCamelCase : Union[str, Any]					 =				output.images
    __lowerCamelCase : Dict					 =				image[0, -3:, -3:, -1]
    __lowerCamelCase : Union[str, Any]					 =				[
        0.42285156,
        0.36914062,
        0.29077148,
        0.42041016,
        0.41918945,
        0.35498047,
        0.3618164,
        0.4423828,
        0.43115234,
    ]

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

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











'''simple docstring'''





import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
   import torch

   from transformers import CamembertModel






@require_torch
@require_sentencepiece
@require_tokenizers
class       lowerCamelCase_					(       unittest.TestCase  ):
  """simple docstring"""




  @slow
  def 		_lowercase  (				self   :			str     )      ->      List[Any]:
    __lowerCamelCase : str					 =				CamembertModel.from_pretrained('camembert-base'     )
    model.to(_a     )

    __lowerCamelCase : List[Any]					 =				torch.tensor(
        [[5, 121, 11, 660, 16, 730, 2_5543, 110, 83, 6]]    ,			device=_a    ,			dtype=torch.long    ,			)  # J'aime le camembert !
    with torch.no_grad():
      __lowerCamelCase : int					 =				model(_a     )['last_hidden_state']
    __lowerCamelCase : Optional[int]					 =				torch.Size((1, 10, 768)     )
    self.assertEqual(output.shape    ,			_a     )
    # compare the actual values for a slice.
    __lowerCamelCase : Tuple					 =				torch.tensor(
        [[[-0.0254, 0.0235, 0.1027], [0.0606, -0.1811, -0.0418], [-0.1561, -0.1127, 0.2687]]]    ,			device=_a    ,			dtype=torch.float    ,			)
    # camembert = torch.hub.load('pytorch/fairseq', 'camembert.v0')
    # camembert.eval()
    # expected_slice = roberta.model.forward(input_ids)[0][:, :3, :3].detach()

    self.assertTrue(torch.allclose(output[:, :3, :3]    ,			_a    ,			atol=1e-4     )     )






