

'''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 					lowerCAmelCase       ( A				,     A				,     unittest.TestCase				):
       lowerCAmelCase_				=	StableDiffusionParadigmsPipeline
       lowerCAmelCase_				=	TEXT_TO_IMAGE_PARAMS
       lowerCAmelCase_				=	TEXT_TO_IMAGE_BATCH_PARAMS
       lowerCAmelCase_				=	TEXT_TO_IMAGE_IMAGE_PARAMS
       lowerCAmelCase_				=	TEXT_TO_IMAGE_IMAGE_PARAMS







       def      snake_case				(							self     :	Any ):







            """simple docstring"""







            torch.manual_seed(0 )
            __lowercase     						=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=__lowercase  , )
            __lowercase     						=DDIMParallelScheduler(
                beta_start=0.0_0_0_8_5  , beta_end=0.0_1_2  , beta_schedule='scaled_linear'  , clip_sample=__lowercase  , set_alpha_to_one=__lowercase  , )
            torch.manual_seed(0 )
            __lowercase     						=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 )
            __lowercase     						=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  , )
            __lowercase     						=CLIPTextModel(__lowercase )
            __lowercase     						=CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

            __lowercase     						={
                'unet': unet,
                'scheduler': scheduler,
                'vae': vae,
                'text_encoder': text_encoder,
                'tokenizer': tokenizer,
                'safety_checker': None,
                'feature_extractor': None,
            }
            return components







       def      snake_case				(							self     :	Optional[int]  , __lowercase     :	Union[str, Any]  , __lowercase     :	List[Any]=0 ):







            """simple docstring"""







            if str(__lowercase ).startswith('mps' ):
                 __lowercase     						=torch.manual_seed(__lowercase )
            else:
                 __lowercase     						=torch.Generator(device=__lowercase ).manual_seed(__lowercase )
            __lowercase     						={
                '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      snake_case				(							self     :	Tuple ):







            """simple docstring"""







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

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

            __lowercase     						=np.array([0.4_7_7_3, 0.5_4_1_7, 0.4_7_2_3, 0.4_9_2_5, 0.5_6_3_1, 0.4_7_5_2, 0.5_2_4_0, 0.4_9_3_5, 0.5_0_2_3] )

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







       def      snake_case				(							self     :	Optional[Any] ):







            """simple docstring"""







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

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

            __lowercase     						=np.array([0.3_5_7_3, 0.4_4_2_0, 0.4_9_6_0, 0.4_7_9_9, 0.3_7_9_6, 0.3_8_7_9, 0.4_8_1_9, 0.4_3_6_5, 0.4_4_6_8] )

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







       def      snake_case				(							self     :	Union[str, Any] ):







            """simple docstring"""







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







       def      snake_case				(							self     :	Tuple ):







            """simple docstring"""







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







       def      snake_case				(							self     :	Tuple ):







            """simple docstring"""







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

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

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

            __lowercase     						=np.array([0.4_7_7_1, 0.5_4_2_0, 0.4_6_8_3, 0.4_9_1_8, 0.5_6_3_6, 0.4_7_2_5, 0.5_2_3_0, 0.4_9_2_3, 0.5_0_1_5] )

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






@slow
@require_torch_gpu
class 					lowerCAmelCase       ( unittest.TestCase				):







       def      snake_case				(							self     :	Union[str, Any] ):







            """simple docstring"""







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







       def      snake_case				(							self     :	Any  , __lowercase     :	List[Any]=0 ):







            """simple docstring"""







            __lowercase     						=torch.Generator(device=__lowercase ).manual_seed(__lowercase )
            __lowercase     						={
                '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      snake_case				(							self     :	int ):







            """simple docstring"""







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

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

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

            __lowercase     						=np.array([0.9_6_2_2, 0.9_6_0_2, 0.9_7_4_8, 0.9_5_9_1, 0.9_6_3_0, 0.9_6_9_1, 0.9_6_6_1, 0.9_6_3_1, 0.9_7_4_1] )

            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 					lowerCAmelCase       ( A				):







       def __init__(							self     :	Union[str, Any]  , __lowercase     :	AutoencoderKL  , __lowercase     :	CLIPTextModel  , __lowercase     :	CLIPTokenizer  , __lowercase     :	UNetaDConditionModel  , __lowercase     :	Union[PNDMScheduler, LMSDiscreteScheduler, DDIMScheduler]  , ):







            """simple docstring"""







            super().__init__()

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







       def      snake_case				(							self     :	str  , __lowercase     :	Optional[int] ):







            """simple docstring"""







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







       def      snake_case				(							self     :	List[str]  , __lowercase     :	Tuple ):







            """simple docstring"""







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







       def      snake_case				(							self     :	int  , __lowercase     :	Optional[int] ):







            """simple docstring"""







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

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

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

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

            return torch.cat([uncond_embedding, text_embedding] )







       def __call__(							self     :	Any  , __lowercase     :	Image.Image  , __lowercase     :	str  , __lowercase     :	float = 0.3  , __lowercase     :	float = 0.6  , __lowercase     :	float = 0.5  , __lowercase     :	int = 42  , __lowercase     :	int = 50  , __lowercase     :	float = 7.5  , ):







            """simple docstring"""







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

            __lowercase     						=self.prep_text(__lowercase )

            self.scheduler.set_timesteps(__lowercase )

            __lowercase						,       __lowercase     						=img.size
            __lowercase     						=self.encode(__lowercase )

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

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

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

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

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

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

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

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

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

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

                      __lowercase     						=self.scheduler.scale_model_input(__lowercase  , __lowercase )

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

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

                      __lowercase     						=self.scheduler.step(__lowercase  , __lowercase  , __lowercase ).prev_sample

            return self.decode(__lowercase )


