'''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.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   _lowercase(     unittest.TestCase				):



   """simple docstring"""





   def 			snake_case	(							self:    List[Any]	):
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()





   @property
   def 			snake_case	(							self:    Dict	):
    __UpperCAmelCase         =					1
    __UpperCAmelCase         =					3
    __UpperCAmelCase         =					(32, 32)

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





   @property
   def 			snake_case	(							self:    Any	):
    torch.manual_seed(0	)
    __UpperCAmelCase         =					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 			snake_case	(							self:    str	):
    torch.manual_seed(0	)
    __UpperCAmelCase         =					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 			snake_case	(							self:    Union[str, Any]	):
    torch.manual_seed(0	)
    __UpperCAmelCase         =					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 			snake_case	(							self:    str	):
    def extract(*a:    List[str]			,**a:    Any	):

     class   _lowercase:



        """simple docstring"""





        def __init__(							self:    List[Any]	):
         __UpperCAmelCase         =					torch.ones([0]	)





        def 			snake_case	(							self:    List[str]			,a:    List[str]	):
         self.pixel_values.to(a	)
         return self




     return Out()

    return extract





   def 			snake_case	(							self:    Optional[int]	):
    __UpperCAmelCase         =					'cpu'  # ensure determinism for the device-dependent torch.Generator
    __UpperCAmelCase         =					self.dummy_cond_unet
    __UpperCAmelCase         =					DDIMScheduler(
        beta_start=0.00085			,beta_end=0.012			,beta_schedule='scaled_linear'			,clip_sample=a			,set_alpha_to_one=a			,)

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

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

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

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

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

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

    assert image.shape == (1, 64, 64, 3)
    __UpperCAmelCase         =					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 			snake_case	(							self:    Optional[int]	):
    __UpperCAmelCase         =					'cpu'  # ensure determinism for the device-dependent torch.Generator
    __UpperCAmelCase         =					self.dummy_cond_unet
    __UpperCAmelCase         =					PNDMScheduler(skip_prk_steps=a	)
    __UpperCAmelCase         =					self.dummy_vae
    __UpperCAmelCase         =					self.dummy_text_encoder
    __UpperCAmelCase         =					CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'	)

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

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

    __UpperCAmelCase         =					output.images

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

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

    assert image.shape == (1, 64, 64, 3)
    __UpperCAmelCase         =					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 			snake_case	(							self:    List[str]	):
    __UpperCAmelCase         =					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

    __UpperCAmelCase         =					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	)
     __UpperCAmelCase         =					StableDiffusionPipeline.from_pretrained(a	)

    # sanity check that the pipeline still works
    assert pipe.safety_checker is None
    __UpperCAmelCase         =					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 			snake_case	(							self:    Dict	):
    __UpperCAmelCase         =					self.dummy_cond_unet
    __UpperCAmelCase         =					PNDMScheduler(skip_prk_steps=a	)
    __UpperCAmelCase         =					self.dummy_vae
    __UpperCAmelCase         =					self.dummy_text_encoder
    __UpperCAmelCase         =					CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'	)

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

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

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

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




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



   """simple docstring"""





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





   def 			snake_case	(							self:    Optional[Any]	):
    __UpperCAmelCase         =					'cuda'
    __UpperCAmelCase         =					StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'	)
    __UpperCAmelCase         =					pipe.to(a	)
    pipe.set_progress_bar_config(disable=a	)

    __UpperCAmelCase         =					'a photo of a cat'
    __UpperCAmelCase         =					{
        '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,
    }

    __UpperCAmelCase         =					3
    __UpperCAmelCase         =					7

    # no sega enabled
    __UpperCAmelCase         =					torch.Generator(a	)
    generator.manual_seed(a	)
    __UpperCAmelCase         =					pipe(
        [prompt]			,generator=a			,guidance_scale=a			,num_inference_steps=50			,output_type='np'			,width=512			,height=512			,)

    __UpperCAmelCase         =					output.images
    __UpperCAmelCase         =					image[0, -3:, -3:, -1]
    __UpperCAmelCase         =					[
        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	)
    __UpperCAmelCase         =					pipe(
        [prompt]			,generator=a			,guidance_scale=a			,num_inference_steps=50			,output_type='np'			,width=512			,height=512			,**a			,)

    __UpperCAmelCase         =					output.images
    __UpperCAmelCase         =					image[0, -3:, -3:, -1]
    __UpperCAmelCase         =					[
        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 			snake_case	(							self:    Union[str, Any]	):
    __UpperCAmelCase         =					'cuda'
    __UpperCAmelCase         =					StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'	)
    __UpperCAmelCase         =					pipe.to(a	)
    pipe.set_progress_bar_config(disable=a	)

    __UpperCAmelCase         =					'an image of a crowded boulevard, realistic, 4k'
    __UpperCAmelCase         =					{
        '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,
    }

    __UpperCAmelCase         =					9
    __UpperCAmelCase         =					7

    # no sega enabled
    __UpperCAmelCase         =					torch.Generator(a	)
    generator.manual_seed(a	)
    __UpperCAmelCase         =					pipe(
        [prompt]			,generator=a			,guidance_scale=a			,num_inference_steps=50			,output_type='np'			,width=512			,height=512			,)

    __UpperCAmelCase         =					output.images
    __UpperCAmelCase         =					image[0, -3:, -3:, -1]
    __UpperCAmelCase         =					[
        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	)
    __UpperCAmelCase         =					pipe(
        [prompt]			,generator=a			,guidance_scale=a			,num_inference_steps=50			,output_type='np'			,width=512			,height=512			,**a			,)

    __UpperCAmelCase         =					output.images
    __UpperCAmelCase         =					image[0, -3:, -3:, -1]
    __UpperCAmelCase         =					[
        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 			snake_case	(							self:    List[str]	):
    __UpperCAmelCase         =					'cuda'
    __UpperCAmelCase         =					StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'	)
    __UpperCAmelCase         =					pipe.to(a	)
    pipe.set_progress_bar_config(disable=a	)

    __UpperCAmelCase         =					'a castle next to a river'
    __UpperCAmelCase         =					{
        '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,
    }

    __UpperCAmelCase         =					48
    __UpperCAmelCase         =					7

    # no sega enabled
    __UpperCAmelCase         =					torch.Generator(a	)
    generator.manual_seed(a	)
    __UpperCAmelCase         =					pipe(
        [prompt]			,generator=a			,guidance_scale=a			,num_inference_steps=50			,output_type='np'			,width=512			,height=512			,)

    __UpperCAmelCase         =					output.images
    __UpperCAmelCase         =					image[0, -3:, -3:, -1]
    __UpperCAmelCase         =					[
        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	)
    __UpperCAmelCase         =					pipe(
        [prompt]			,generator=a			,guidance_scale=a			,num_inference_steps=50			,output_type='np'			,width=512			,height=512			,**a			,)

    __UpperCAmelCase         =					output.images
    __UpperCAmelCase         =					image[0, -3:, -3:, -1]
    __UpperCAmelCase         =					[
        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 			snake_case	(							self:    List[str]	):
    __UpperCAmelCase         =					'cuda'
    __UpperCAmelCase         =					StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'			,torch_dtype=torch.floataa	)
    __UpperCAmelCase         =					pipe.to(a	)
    pipe.set_progress_bar_config(disable=a	)

    __UpperCAmelCase         =					'a photo of a cat'
    __UpperCAmelCase         =					{
        '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,
    }

    __UpperCAmelCase         =					3
    __UpperCAmelCase         =					7

    # no sega enabled
    __UpperCAmelCase         =					torch.Generator(a	)
    generator.manual_seed(a	)
    __UpperCAmelCase         =					pipe(
        [prompt]			,generator=a			,guidance_scale=a			,num_inference_steps=50			,output_type='np'			,width=512			,height=512			,)

    __UpperCAmelCase         =					output.images
    __UpperCAmelCase         =					image[0, -3:, -3:, -1]
    __UpperCAmelCase         =					[
        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	)
    __UpperCAmelCase         =					pipe(
        [prompt]			,generator=a			,guidance_scale=a			,num_inference_steps=50			,output_type='np'			,width=512			,height=512			,**a			,)

    __UpperCAmelCase         =					output.images
    __UpperCAmelCase         =					image[0, -3:, -3:, -1]
    __UpperCAmelCase         =					[
        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'''




from typing import Union

import torch
from PIL import Image
from torchvision import transforms as tfms
from tqdm.auto import tqdm
from transformers import CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)

class   _lowercase(     _lowerCamelCase				):



   """simple docstring"""





   def __init__(							self:    List[str]			,a:    AutoencoderKL			,a:    CLIPTextModel			,a:    CLIPTokenizer			,a:    UNetaDConditionModel			,a:    Union[PNDMScheduler, LMSDiscreteScheduler, DDIMScheduler]			,):
    super().__init__()

    self.register_modules(vae=a			,text_encoder=a			,tokenizer=a			,unet=a			,scheduler=a	)





   def 			snake_case	(							self:    Union[str, Any]			,a:    Any	):
    with torch.no_grad():
     __UpperCAmelCase         =					self.vae.encode(tfms.ToTensor()(a	).unsqueeze(0	).to(self.device	) * 2 - 1	)
     __UpperCAmelCase         =					0.18215 * latent.latent_dist.sample()
    return latent





   def 			snake_case	(							self:    Any			,a:    List[str]	):
    __UpperCAmelCase         =					(1 / 0.18215) * latent
    with torch.no_grad():
     __UpperCAmelCase         =					self.vae.decode(a	).sample
    __UpperCAmelCase         =					(img / 2 + 0.5).clamp(0			,1	)
    __UpperCAmelCase         =					img.detach().cpu().permute(0			,2			,3			,1	).numpy()
    __UpperCAmelCase         =					(img * 255).round().astype('uint8'	)
    return Image.fromarray(img[0]	)





   def 			snake_case	(							self:    Any			,a:    str	):
    __UpperCAmelCase         =					self.tokenizer(
        a			,padding='max_length'			,max_length=self.tokenizer.model_max_length			,truncation=a			,return_tensors='pt'			,)

    __UpperCAmelCase         =					self.text_encoder(text_input.input_ids.to(self.device	)	)[0]

    __UpperCAmelCase         =					self.tokenizer(
        ''			,padding='max_length'			,max_length=self.tokenizer.model_max_length			,truncation=a			,return_tensors='pt'			,)

    __UpperCAmelCase         =					self.text_encoder(uncond_input.input_ids.to(self.device	)	)[0]

    return torch.cat([uncond_embedding, text_embedding]	)





   def __call__(							self:    Optional[Any]			,a:    Image.Image			,a:    str			,a:    float = 0.3			,a:    float = 0.6			,a:    float = 0.5			,a:    int = 42			,a:    int = 50			,a:    float = 7.5			,):
    __UpperCAmelCase         =					steps - int(kmin * steps	)
    __UpperCAmelCase         =					steps - int(kmax * steps	)

    __UpperCAmelCase         =					self.prep_text(a	)

    self.scheduler.set_timesteps(a	)

    __UpperCAmelCase  ,			__UpperCAmelCase         =					img.size
    __UpperCAmelCase         =					self.encode(a	)

    torch.manual_seed(a	)
    __UpperCAmelCase         =					torch.randn(
        (1, self.unet.config.in_channels, height // 8, width // 8)			,).to(self.device	)

    __UpperCAmelCase         =					self.scheduler.add_noise(
        a			,a			,timesteps=self.scheduler.timesteps[tmax]			,)

    __UpperCAmelCase         =					torch.cat([latents] * 2	)

    __UpperCAmelCase         =					self.scheduler.scale_model_input(a			,self.scheduler.timesteps[tmax]	)

    with torch.no_grad():
     __UpperCAmelCase         =					self.unet(
         a			,self.scheduler.timesteps[tmax]			,encoder_hidden_states=a			,).sample

    __UpperCAmelCase  ,			__UpperCAmelCase         =					pred.chunk(2	)
    __UpperCAmelCase         =					pred_uncond + guidance_scale * (pred_text - pred_uncond)

    __UpperCAmelCase         =					self.scheduler.step(a			,self.scheduler.timesteps[tmax]			,a	).prev_sample

    for i, t in enumerate(tqdm(self.scheduler.timesteps	)	):
     if i > tmax:
      if i < tmin:  # layout generation phase
       __UpperCAmelCase         =					self.scheduler.add_noise(
           a			,a			,timesteps=a			,)

       __UpperCAmelCase         =					(mix_factor * latents) + (
           1 - mix_factor
       ) * orig_latents  # interpolating between layout noise and conditionally generated noise to preserve layout sematics
       __UpperCAmelCase         =					torch.cat([input] * 2	)

      else:  # content generation phase
       __UpperCAmelCase         =					torch.cat([latents] * 2	)

      __UpperCAmelCase         =					self.scheduler.scale_model_input(a			,a	)

      with torch.no_grad():
       __UpperCAmelCase         =					self.unet(
           a			,a			,encoder_hidden_states=a			,).sample

      __UpperCAmelCase  ,			__UpperCAmelCase         =					pred.chunk(2	)
      __UpperCAmelCase         =					pred_uncond + guidance_scale * (pred_text - pred_uncond)

      __UpperCAmelCase         =					self.scheduler.step(a			,a			,a	).prev_sample

    return self.decode(a	)



