





"""simple docstring"""


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        lowercase_			(       __lowerCAmelCase			,__lowerCAmelCase			,unittest.TestCase							):
       snake_case__			:							int										=       StableDiffusionParadigmsPipeline
       snake_case__			:							Dict										=       TEXT_TO_IMAGE_PARAMS
       snake_case__			:							List[str]										=       TEXT_TO_IMAGE_BATCH_PARAMS
       snake_case__			:							Dict										=       TEXT_TO_IMAGE_IMAGE_PARAMS
       snake_case__			:							Tuple										=       TEXT_TO_IMAGE_IMAGE_PARAMS

       def 						_SCREAMING_SNAKE_CASE(     self       ):






        '''simple docstring'''

        torch.manual_seed(0       )
        snake_case__      							=			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=__SCREAMING_SNAKE_CASE				,   )
        snake_case__      							=			DDIMParallelScheduler(
            beta_start=0.00085				,   beta_end=0.012				,   beta_schedule="scaled_linear"				,   clip_sample=__SCREAMING_SNAKE_CASE				,   set_alpha_to_one=__SCREAMING_SNAKE_CASE				,   )
        torch.manual_seed(0       )
        snake_case__      							=			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       )
        snake_case__      							=			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				,   )
        snake_case__      							=			CLIPTextModel(__SCREAMING_SNAKE_CASE       )
        snake_case__      							=			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       )

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

       def 						_SCREAMING_SNAKE_CASE(     self				,   __SCREAMING_SNAKE_CASE				,   __SCREAMING_SNAKE_CASE=0       ):






        '''simple docstring'''

        if str(__SCREAMING_SNAKE_CASE       ).startswith("mps"       ):
         snake_case__      							=			torch.manual_seed(__SCREAMING_SNAKE_CASE       )
        else:
         snake_case__      							=			torch.Generator(device=__SCREAMING_SNAKE_CASE       ).manual_seed(__SCREAMING_SNAKE_CASE       )
        snake_case__      							=			{
            "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 						_SCREAMING_SNAKE_CASE(     self       ):






        '''simple docstring'''

        snake_case__      							=			"cpu"  # ensure determinism for the device-dependent torch.Generator
        snake_case__      							=			self.get_dummy_components()
        snake_case__      							=			StableDiffusionParadigmsPipeline(**__SCREAMING_SNAKE_CASE       )
        snake_case__      							=			sd_pipe.to(__SCREAMING_SNAKE_CASE       )
        sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE       )

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

        snake_case__      							=			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 						_SCREAMING_SNAKE_CASE(     self       ):






        '''simple docstring'''

        snake_case__      							=			"cpu"  # ensure determinism for the device-dependent torch.Generator
        snake_case__      							=			self.get_dummy_components()
        torch.manual_seed(0       )
        snake_case__      							=			DDPMParallelScheduler()
        torch.manual_seed(0       )
        snake_case__      							=			StableDiffusionParadigmsPipeline(**__SCREAMING_SNAKE_CASE       )
        snake_case__      							=			sd_pipe.to(__SCREAMING_SNAKE_CASE       )
        sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE       )

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

        snake_case__      							=			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 						_SCREAMING_SNAKE_CASE(     self       ):






        '''simple docstring'''

        super().test_inference_batch_consistent(batch_sizes=[1, 2]       )

       def 						_SCREAMING_SNAKE_CASE(     self       ):






        '''simple docstring'''

        super().test_inference_batch_single_identical(batch_size=2				,   expected_max_diff=3e-3       )

       def 						_SCREAMING_SNAKE_CASE(     self       ):






        '''simple docstring'''

        snake_case__      							=			"cpu"  # ensure determinism for the device-dependent torch.Generator
        snake_case__      							=			self.get_dummy_components()
        snake_case__      							=			StableDiffusionParadigmsPipeline(**__SCREAMING_SNAKE_CASE       )
        snake_case__      							=			sd_pipe.to(__SCREAMING_SNAKE_CASE       )
        sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE       )

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

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

        snake_case__      							=			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        lowercase_			(       unittest.TestCase							):

       def 						_SCREAMING_SNAKE_CASE(     self       ):






        '''simple docstring'''

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

       def 						_SCREAMING_SNAKE_CASE(     self				,   __SCREAMING_SNAKE_CASE=0       ):






        '''simple docstring'''

        snake_case__      							=			torch.Generator(device=__SCREAMING_SNAKE_CASE       ).manual_seed(__SCREAMING_SNAKE_CASE       )
        snake_case__      							=			{
            "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 						_SCREAMING_SNAKE_CASE(     self       ):






        '''simple docstring'''

        snake_case__      							=			"stabilityai/stable-diffusion-2-base"
        snake_case__      							=			DDIMParallelScheduler.from_pretrained(__SCREAMING_SNAKE_CASE				,   subfolder="scheduler"       )
        snake_case__      							=			StableDiffusionParadigmsPipeline.from_pretrained(__SCREAMING_SNAKE_CASE				,   scheduler=__SCREAMING_SNAKE_CASE				,   safety_checker=__SCREAMING_SNAKE_CASE       )
        pipe.to(__SCREAMING_SNAKE_CASE       )
        pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE       )
        pipe.enable_attention_slicing()

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

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

        snake_case__      							=			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 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							):

       def __init__(     self				,   __SCREAMING_SNAKE_CASE				,   __SCREAMING_SNAKE_CASE				,   __SCREAMING_SNAKE_CASE				,   __SCREAMING_SNAKE_CASE				,   __SCREAMING_SNAKE_CASE				,   ):






        '''simple docstring'''

        super().__init__()

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

       def 						_SCREAMING_SNAKE_CASE(     self				,   __SCREAMING_SNAKE_CASE       ):






        '''simple docstring'''

        with torch.no_grad():
         snake_case__      							=			self.vae.encode(tfms.ToTensor()(__SCREAMING_SNAKE_CASE       ).unsqueeze(0       ).to(self.device       ) * 2 - 1       )
         snake_case__      							=			0.18215 * latent.latent_dist.sample()
        return latent

       def 						_SCREAMING_SNAKE_CASE(     self				,   __SCREAMING_SNAKE_CASE       ):






        '''simple docstring'''

        snake_case__      							=			(1 / 0.18215) * latent
        with torch.no_grad():
         snake_case__      							=			self.vae.decode(__SCREAMING_SNAKE_CASE       ).sample
        snake_case__      							=			(img / 2 + 0.5).clamp(0				,   1       )
        snake_case__      							=			img.detach().cpu().permute(0				,   2				,   3				,   1       ).numpy()
        snake_case__      							=			(img * 255).round().astype("uint8"       )
        return Image.fromarray(img[0]       )

       def 						_SCREAMING_SNAKE_CASE(     self				,   __SCREAMING_SNAKE_CASE       ):






        '''simple docstring'''

        snake_case__      							=			self.tokenizer(
            __SCREAMING_SNAKE_CASE				,   padding="max_length"				,   max_length=self.tokenizer.model_max_length				,   truncation=__SCREAMING_SNAKE_CASE				,   return_tensors="pt"				,   )

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

        snake_case__      							=			self.tokenizer(
            ""				,   padding="max_length"				,   max_length=self.tokenizer.model_max_length				,   truncation=__SCREAMING_SNAKE_CASE				,   return_tensors="pt"				,   )

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

        return torch.cat([uncond_embedding, text_embedding]       )

       def __call__(     self				,   __SCREAMING_SNAKE_CASE				,   __SCREAMING_SNAKE_CASE				,   __SCREAMING_SNAKE_CASE = 0.3				,   __SCREAMING_SNAKE_CASE = 0.6				,   __SCREAMING_SNAKE_CASE = 0.5				,   __SCREAMING_SNAKE_CASE = 42				,   __SCREAMING_SNAKE_CASE = 50				,   __SCREAMING_SNAKE_CASE = 7.5				,   ):






        '''simple docstring'''

        snake_case__      							=			steps - int(kmin * steps       )
        snake_case__      							=			steps - int(kmax * steps       )

        snake_case__      							=			self.prep_text(__SCREAMING_SNAKE_CASE       )

        self.scheduler.set_timesteps(__SCREAMING_SNAKE_CASE       )

        snake_case__							,				snake_case__      							=			img.size
        snake_case__      							=			self.encode(__SCREAMING_SNAKE_CASE       )

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

        snake_case__      							=			self.scheduler.add_noise(
            __SCREAMING_SNAKE_CASE				,   __SCREAMING_SNAKE_CASE				,   timesteps=self.scheduler.timesteps[tmax]				,   )

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

        snake_case__      							=			self.scheduler.scale_model_input(__SCREAMING_SNAKE_CASE				,   self.scheduler.timesteps[tmax]       )

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

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

        snake_case__      							=			self.scheduler.step(__SCREAMING_SNAKE_CASE				,   self.scheduler.timesteps[tmax]				,   __SCREAMING_SNAKE_CASE       ).prev_sample

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

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

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

          snake_case__      							=			self.scheduler.scale_model_input(__SCREAMING_SNAKE_CASE				,   __SCREAMING_SNAKE_CASE       )

          with torch.no_grad():
           snake_case__      							=			self.unet(
               __SCREAMING_SNAKE_CASE				,   __SCREAMING_SNAKE_CASE				,   encoder_hidden_states=__SCREAMING_SNAKE_CASE				,   ).sample

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

          snake_case__      							=			self.scheduler.step(__SCREAMING_SNAKE_CASE				,   __SCREAMING_SNAKE_CASE				,   __SCREAMING_SNAKE_CASE       ).prev_sample

        return self.decode(__SCREAMING_SNAKE_CASE       )

