




'''simple docstring'''



import inspect
from typing import Optional, Union

import numpy as np
import PIL
import torch
from torch.nn import functional as F
from torchvision import transforms
from transformers import CLIPFeatureExtractor, CLIPModel, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    DPMSolverMultistepScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput
from diffusers.utils import (
    PIL_INTERPOLATION,
    randn_tensor,
)







def       _UpperCamelCase			(_lowerCamelCase				:  List[Any] ,						_lowerCamelCase				:  List[Any] ,						_lowerCamelCase				:  Union[str, Any]						)-> Optional[Any]:




     '''simple docstring'''





     if isinstance(_lowerCamelCase ,						torch.Tensor						):
          return image
     elif isinstance(_lowerCamelCase ,						PIL.Image.Image						):
          __snake_case            =	[image]

     if isinstance(image[0] ,						PIL.Image.Image						):
          __snake_case            =	[np.array(i.resize((w, h) ,						resample=PIL_INTERPOLATION['''lanczos''']						)						)[None, :] for i in image]
          __snake_case            =	np.concatenate(_lowerCamelCase ,						axis=0						)
          __snake_case            =	np.array(_lowerCamelCase						).astype(np.floataa						) / 255.0
          __snake_case            =	image.transpose(0 ,						3 ,						1 ,						2						)
          __snake_case            =	2.0 * image - 1.0
          __snake_case            =	torch.from_numpy(_lowerCamelCase						)
     elif isinstance(image[0] ,						torch.Tensor						):
          __snake_case            =	torch.cat(_lowerCamelCase ,						dim=0						)
     return image







def       _UpperCamelCase			(_lowerCamelCase				:  Dict ,						_lowerCamelCase				:  Optional[int] ,						_lowerCamelCase				:  int ,						_lowerCamelCase				:  Optional[Any]=0.9995						)-> Union[str, Any]:




     '''simple docstring'''





     if not isinstance(_lowerCamelCase ,						np.ndarray						):
          __snake_case            =	True
          __snake_case            =	va.device
          __snake_case            =	va.cpu().numpy()
          __snake_case            =	va.cpu().numpy()

     __snake_case            =	np.sum(va * va / (np.linalg.norm(_lowerCamelCase						) * np.linalg.norm(_lowerCamelCase						))						)
     if np.abs(_lowerCamelCase						) > DOT_THRESHOLD:
          __snake_case            =	(1 - t) * va + t * va
     else:
          __snake_case            =	np.arccos(_lowerCamelCase						)
          __snake_case            =	np.sin(_lowerCamelCase						)
          __snake_case            =	theta_a * t
          __snake_case            =	np.sin(_lowerCamelCase						)
          __snake_case            =	np.sin(theta_a - theta_t						) / sin_theta_a
          __snake_case            =	sin_theta_t / sin_theta_a
          __snake_case            =	sa * va + sa * va

     if inputs_are_torch:
          __snake_case            =	torch.from_numpy(_lowerCamelCase						).to(_lowerCamelCase						)

     return va







def       _UpperCamelCase			(_lowerCamelCase				:  str ,						_lowerCamelCase				:  Dict						)-> Union[str, Any]:




     '''simple docstring'''





     __snake_case            =	F.normalize(_lowerCamelCase ,						dim=-1						)
     __snake_case            =	F.normalize(_lowerCamelCase ,						dim=-1						)
     return (x - y).norm(dim=-1						).div(2						).arcsin().pow(2						).mul(2						)







def       _UpperCamelCase			(_lowerCamelCase				:  List[str] ,						_lowerCamelCase				:  List[str]						)-> int:




     '''simple docstring'''





     for param in model.parameters():
          __snake_case            =	value







class      lowerCAmelCase (			__lowerCAmelCase):


       def __init__( self     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE=None     ,	__SCREAMING_SNAKE_CASE=None     ,	__SCREAMING_SNAKE_CASE=None     ,	)   ->	str:



            '''simple docstring'''




            super().__init__()
            self.register_modules(
                vae=__SCREAMING_SNAKE_CASE     ,	text_encoder=__SCREAMING_SNAKE_CASE     ,	clip_model=__SCREAMING_SNAKE_CASE     ,	tokenizer=__SCREAMING_SNAKE_CASE     ,	unet=__SCREAMING_SNAKE_CASE     ,	scheduler=__SCREAMING_SNAKE_CASE     ,	feature_extractor=__SCREAMING_SNAKE_CASE     ,	coca_model=__SCREAMING_SNAKE_CASE     ,	coca_tokenizer=__SCREAMING_SNAKE_CASE     ,	coca_transform=__SCREAMING_SNAKE_CASE     ,	)
            __snake_case            =	(
                feature_extractor.size
                if isinstance(feature_extractor.size     ,	__SCREAMING_SNAKE_CASE	)
                else feature_extractor.size['''shortest_edge''']
            )
            __snake_case            =	transforms.Normalize(mean=feature_extractor.image_mean     ,	std=feature_extractor.image_std	)
            set_requires_grad(self.text_encoder     ,	__SCREAMING_SNAKE_CASE	)
            set_requires_grad(self.clip_model     ,	__SCREAMING_SNAKE_CASE	)


       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE = "auto"	)   ->	List[Any]:



            '''simple docstring'''




            if slice_size == "auto":
                 # half the attention head size is usually a good trade-off between
                 # speed and memory
                 __snake_case            =	self.unet.config.attention_head_dim // 2
            self.unet.set_attention_slice(__SCREAMING_SNAKE_CASE	)


       def    lowerCAmelCase   ( self	)   ->	str:



            '''simple docstring'''




            self.enable_attention_slicing(__SCREAMING_SNAKE_CASE	)


       def    lowerCAmelCase   ( self	)   ->	Tuple:



            '''simple docstring'''




            set_requires_grad(self.vae     ,	__SCREAMING_SNAKE_CASE	)


       def    lowerCAmelCase   ( self	)   ->	Optional[int]:



            '''simple docstring'''




            set_requires_grad(self.vae     ,	__SCREAMING_SNAKE_CASE	)


       def    lowerCAmelCase   ( self	)   ->	str:



            '''simple docstring'''




            set_requires_grad(self.unet     ,	__SCREAMING_SNAKE_CASE	)


       def    lowerCAmelCase   ( self	)   ->	List[str]:



            '''simple docstring'''




            set_requires_grad(self.unet     ,	__SCREAMING_SNAKE_CASE	)


       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)   ->	Any:



            '''simple docstring'''




            __snake_case            =	min(int(num_inference_steps * strength	)     ,	__SCREAMING_SNAKE_CASE	)

            __snake_case            =	max(num_inference_steps - init_timestep     ,	0	)
            __snake_case            =	self.scheduler.timesteps[t_start:]

            return timesteps, num_inference_steps - t_start


       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE=None	)   ->	str:



            '''simple docstring'''




            if not isinstance(__SCREAMING_SNAKE_CASE     ,	torch.Tensor	):
                 raise ValueError(F'''`image` has to be of type `torch.Tensor` but is {type(__SCREAMING_SNAKE_CASE	)}'''	)

            __snake_case            =	image.to(device=__SCREAMING_SNAKE_CASE     ,	dtype=__SCREAMING_SNAKE_CASE	)

            if isinstance(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	):
                 __snake_case            =	[
                     self.vae.encode(image[i : i + 1]	).latent_dist.sample(generator[i]	) for i in range(__SCREAMING_SNAKE_CASE	)
                 ]
                 __snake_case            =	torch.cat(__SCREAMING_SNAKE_CASE     ,	dim=0	)
            else:
                 __snake_case            =	self.vae.encode(__SCREAMING_SNAKE_CASE	).latent_dist.sample(__SCREAMING_SNAKE_CASE	)

            # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
            __snake_case            =	0.18_215 * init_latents
            __snake_case            =	init_latents.repeat_interleave(__SCREAMING_SNAKE_CASE     ,	dim=0	)

            __snake_case            =	randn_tensor(init_latents.shape     ,	generator=__SCREAMING_SNAKE_CASE     ,	device=__SCREAMING_SNAKE_CASE     ,	dtype=__SCREAMING_SNAKE_CASE	)

            # get latents
            __snake_case            =	self.scheduler.add_noise(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)
            __snake_case            =	init_latents

            return latents


       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE	)   ->	str:



            '''simple docstring'''




            __snake_case            =	self.coca_transform(__SCREAMING_SNAKE_CASE	).unsqueeze(0	)
            with torch.no_grad(), torch.cuda.amp.autocast():
                 __snake_case            =	self.coca_model.generate(transformed_image.to(device=self.device     ,	dtype=self.coca_model.dtype	)	)
            __snake_case            =	self.coca_tokenizer.decode(generated[0].cpu().numpy()	)
            return generated.split('''<end_of_text>'''	)[0].replace('''<start_of_text>'''     ,	''''''	).rstrip(''' .,'''	)


       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)   ->	int:



            '''simple docstring'''




            __snake_case            =	self.feature_extractor.preprocess(__SCREAMING_SNAKE_CASE	)
            __snake_case            =	torch.from_numpy(clip_image_input['''pixel_values'''][0]	).unsqueeze(0	).to(self.device	).half()
            __snake_case            =	self.clip_model.get_image_features(__SCREAMING_SNAKE_CASE	)
            __snake_case            =	image_embeddings_clip / image_embeddings_clip.norm(p=2     ,	dim=-1     ,	keepdim=__SCREAMING_SNAKE_CASE	)
            __snake_case            =	image_embeddings_clip.repeat_interleave(__SCREAMING_SNAKE_CASE     ,	dim=0	)
            return image_embeddings_clip


       @torch.enable_grad()
       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	)   ->	int:



            '''simple docstring'''




            __snake_case            =	latents.detach().requires_grad_()

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

            # predict the noise residual
            __snake_case            =	self.unet(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	encoder_hidden_states=__SCREAMING_SNAKE_CASE	).sample

            if isinstance(self.scheduler     ,	(PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)	):
                 __snake_case            =	self.scheduler.alphas_cumprod[timestep]
                 __snake_case            =	1 - alpha_prod_t
                 # compute predicted original sample from predicted noise also called
                 # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
                 __snake_case            =	(latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5

                 __snake_case            =	torch.sqrt(__SCREAMING_SNAKE_CASE	)
                 __snake_case            =	pred_original_sample * (fac) + latents * (1 - fac)
            elif isinstance(self.scheduler     ,	__SCREAMING_SNAKE_CASE	):
                 __snake_case            =	self.scheduler.sigmas[index]
                 __snake_case            =	latents - sigma * noise_pred
            else:
                 raise ValueError(F'''scheduler type {type(self.scheduler	)} not supported'''	)

            # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
            __snake_case            =	1 / 0.18_215 * sample
            __snake_case            =	self.vae.decode(__SCREAMING_SNAKE_CASE	).sample
            __snake_case            =	(image / 2 + 0.5).clamp(0     ,	1	)

            __snake_case            =	transforms.Resize(self.feature_extractor_size	)(__SCREAMING_SNAKE_CASE	)
            __snake_case            =	self.normalize(__SCREAMING_SNAKE_CASE	).to(latents.dtype	)

            __snake_case            =	self.clip_model.get_image_features(__SCREAMING_SNAKE_CASE	)
            __snake_case            =	image_embeddings_clip / image_embeddings_clip.norm(p=2     ,	dim=-1     ,	keepdim=__SCREAMING_SNAKE_CASE	)

            __snake_case            =	spherical_dist_loss(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	).mean() * clip_guidance_scale

            __snake_case            =	-torch.autograd.grad(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)[0]

            if isinstance(self.scheduler     ,	__SCREAMING_SNAKE_CASE	):
                 __snake_case            =	latents.detach() + grads * (sigma**2)
                 __snake_case            =	noise_pred_original
            else:
                 __snake_case            =	noise_pred_original - torch.sqrt(__SCREAMING_SNAKE_CASE	) * grads
            return noise_pred, latents




       @torch.no_grad()
       def __call__( self     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE = None     ,	__SCREAMING_SNAKE_CASE = None     ,	__SCREAMING_SNAKE_CASE = 512     ,	__SCREAMING_SNAKE_CASE = 512     ,	__SCREAMING_SNAKE_CASE = 0.6     ,	__SCREAMING_SNAKE_CASE = 50     ,	__SCREAMING_SNAKE_CASE = 7.5     ,	__SCREAMING_SNAKE_CASE = 1     ,	__SCREAMING_SNAKE_CASE = 0.0     ,	__SCREAMING_SNAKE_CASE = 100     ,	__SCREAMING_SNAKE_CASE = None     ,	__SCREAMING_SNAKE_CASE = "pil"     ,	__SCREAMING_SNAKE_CASE = True     ,	__SCREAMING_SNAKE_CASE = 0.8     ,	__SCREAMING_SNAKE_CASE = 0.1     ,	__SCREAMING_SNAKE_CASE = 0.1     ,	)   ->	Any:



            '''simple docstring'''




            if isinstance(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	) and len(__SCREAMING_SNAKE_CASE	) != batch_size:
                 raise ValueError(F'''You have passed {batch_size} batch_size, but only {len(__SCREAMING_SNAKE_CASE	)} generators.'''	)

            if height % 8 != 0 or width % 8 != 0:
                 raise ValueError(F'''`height` and `width` have to be divisible by 8 but are {height} and {width}.'''	)

            if isinstance(__SCREAMING_SNAKE_CASE     ,	torch.Generator	) and batch_size > 1:
                 __snake_case            =	[generator] + [None] * (batch_size - 1)

            __snake_case            =	[
                ('''model''', self.coca_model is None),
                ('''tokenizer''', self.coca_tokenizer is None),
                ('''transform''', self.coca_transform is None),
            ]
            __snake_case            =	[x[0] for x in coca_is_none if x[1]]
            __snake_case            =	''', '''.join(__SCREAMING_SNAKE_CASE	)
            # generate prompts with coca model if prompt is None
            if content_prompt is None:
                 if len(__SCREAMING_SNAKE_CASE	):
                      raise ValueError(
                          F'''Content prompt is None and CoCa [{coca_is_none_str}] is None.'''
                          F'''Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline.'''	)
                 __snake_case            =	self.get_image_description(__SCREAMING_SNAKE_CASE	)
            if style_prompt is None:
                 if len(__SCREAMING_SNAKE_CASE	):
                      raise ValueError(
                          F'''Style prompt is None and CoCa [{coca_is_none_str}] is None.'''
                          F''' Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline.'''	)
                 __snake_case            =	self.get_image_description(__SCREAMING_SNAKE_CASE	)

            # get prompt text embeddings for content and style
            __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(content_text_input.input_ids.to(self.device	)	)[0]

            __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(style_text_input.input_ids.to(self.device	)	)[0]

            __snake_case            =	slerp(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)

            # duplicate text embeddings for each generation per prompt
            __snake_case            =	text_embeddings.repeat_interleave(__SCREAMING_SNAKE_CASE     ,	dim=0	)

            # set timesteps
            __snake_case            =	'''offset''' in set(inspect.signature(self.scheduler.set_timesteps	).parameters.keys()	)
            __snake_case            =	{}
            if accepts_offset:
                 __snake_case            =	1

            self.scheduler.set_timesteps(__SCREAMING_SNAKE_CASE     ,	**__SCREAMING_SNAKE_CASE	)
            # Some schedulers like PNDM have timesteps as arrays
            # It's more optimized to move all timesteps to correct device beforehand
            self.scheduler.timesteps.to(self.device	)

            __snake_case       ,     __snake_case            =	self.get_timesteps(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	self.device	)
            __snake_case            =	timesteps[:1].repeat(__SCREAMING_SNAKE_CASE	)

            # Preprocess image
            __snake_case            =	preprocess(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)
            __snake_case            =	self.prepare_latents(
                __SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	text_embeddings.dtype     ,	self.device     ,	__SCREAMING_SNAKE_CASE	)

            __snake_case            =	preprocess(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)
            __snake_case            =	self.prepare_latents(
                __SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	text_embeddings.dtype     ,	self.device     ,	__SCREAMING_SNAKE_CASE	)

            __snake_case            =	slerp(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)

            if clip_guidance_scale > 0:
                 __snake_case            =	self.get_clip_image_embeddings(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)
                 __snake_case            =	self.get_clip_image_embeddings(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)
                 __snake_case            =	slerp(
                     __SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)

            # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
            # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
            # corresponds to doing no classifier free guidance.
            __snake_case            =	guidance_scale > 1.0
            # get unconditional embeddings for classifier free guidance
            if do_classifier_free_guidance:
                 __snake_case            =	content_text_input.input_ids.shape[-1]
                 __snake_case            =	self.tokenizer(['''''']     ,	padding='''max_length'''     ,	max_length=__SCREAMING_SNAKE_CASE     ,	return_tensors='''pt'''	)
                 __snake_case            =	self.text_encoder(uncond_input.input_ids.to(self.device	)	)[0]
                 # duplicate unconditional embeddings for each generation per prompt
                 __snake_case            =	uncond_embeddings.repeat_interleave(__SCREAMING_SNAKE_CASE     ,	dim=0	)

                 # For classifier free guidance, we need to do two forward passes.
                 # Here we concatenate the unconditional and text embeddings into a single batch
                 # to avoid doing two forward passes
                 __snake_case            =	torch.cat([uncond_embeddings, text_embeddings]	)

            # get the initial random noise unless the user supplied it

            # Unlike in other pipelines, latents need to be generated in the target device
            # for 1-to-1 results reproducibility with the CompVis implementation.
            # However this currently doesn't work in `mps`.
            __snake_case            =	(batch_size, self.unet.config.in_channels, height // 8, width // 8)
            __snake_case            =	text_embeddings.dtype
            if latents is None:
                 if self.device.type == "mps":
                      # randn does not work reproducibly on mps
                      __snake_case            =	torch.randn(__SCREAMING_SNAKE_CASE     ,	generator=__SCREAMING_SNAKE_CASE     ,	device='''cpu'''     ,	dtype=__SCREAMING_SNAKE_CASE	).to(
                          self.device	)
                 else:
                      __snake_case            =	torch.randn(__SCREAMING_SNAKE_CASE     ,	generator=__SCREAMING_SNAKE_CASE     ,	device=self.device     ,	dtype=__SCREAMING_SNAKE_CASE	)
            else:
                 if latents.shape != latents_shape:
                      raise ValueError(F'''Unexpected latents shape, got {latents.shape}, expected {latents_shape}'''	)
                 __snake_case            =	latents.to(self.device	)

            # scale the initial noise by the standard deviation required by the scheduler
            __snake_case            =	latents * self.scheduler.init_noise_sigma

            # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
            # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
            # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
            # and should be between [0, 1]
            __snake_case            =	'''eta''' in set(inspect.signature(self.scheduler.step	).parameters.keys()	)
            __snake_case            =	{}
            if accepts_eta:
                 __snake_case            =	eta

            # check if the scheduler accepts generator
            __snake_case            =	'''generator''' in set(inspect.signature(self.scheduler.step	).parameters.keys()	)
            if accepts_generator:
                 __snake_case            =	generator

            with self.progress_bar(total=__SCREAMING_SNAKE_CASE	):
                 for i, t in enumerate(__SCREAMING_SNAKE_CASE	):
                      # expand the latents if we are doing classifier free guidance
                      __snake_case            =	torch.cat([latents] * 2	) if do_classifier_free_guidance else latents
                      __snake_case            =	self.scheduler.scale_model_input(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)

                      # predict the noise residual
                      __snake_case            =	self.unet(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	encoder_hidden_states=__SCREAMING_SNAKE_CASE	).sample

                      # perform classifier free guidance
                      if do_classifier_free_guidance:
                           __snake_case       ,     __snake_case            =	noise_pred.chunk(2	)
                           __snake_case            =	noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

                      # perform clip guidance
                      if clip_guidance_scale > 0:
                           __snake_case            =	(
                               text_embeddings.chunk(2	)[1] if do_classifier_free_guidance else text_embeddings
                           )
                           __snake_case       ,     __snake_case            =	self.cond_fn(
                               __SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	)

                      # compute the previous noisy sample x_t -> x_t-1
                      __snake_case            =	self.scheduler.step(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	**__SCREAMING_SNAKE_CASE	).prev_sample

        # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
            __snake_case            =	1 / 0.18_215 * latents
            __snake_case            =	self.vae.decode(__SCREAMING_SNAKE_CASE	).sample

            __snake_case            =	(image / 2 + 0.5).clamp(0     ,	1	)
            __snake_case            =	image.cpu().permute(0     ,	2     ,	3     ,	1	).numpy()

            if output_type == "pil":
                 __snake_case            =	self.numpy_to_pil(__SCREAMING_SNAKE_CASE	)

            if not return_dict:
                 return (image, None)

            return StableDiffusionPipelineOutput(images=__SCREAMING_SNAKE_CASE     ,	nsfw_content_detected=__SCREAMING_SNAKE_CASE	)










'''simple docstring'''
import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, CycleDiffusionPipeline, DDIMScheduler, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






class        lowerCAmelCase_  (     UpperCAmelCase      ,							UpperCAmelCase      ,							unittest.TestCase	):
       __UpperCAmelCase					=CycleDiffusionPipeline
       __UpperCAmelCase					=TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
           'negative_prompt',
           'height',
           'width',
           'negative_prompt_embeds',
       }
       __UpperCAmelCase					=PipelineTesterMixin.required_optional_params - {'latents'}
       __UpperCAmelCase					=TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({'source_prompt'}	)
       __UpperCAmelCase					=IMAGE_TO_IMAGE_IMAGE_PARAMS
       __UpperCAmelCase					=IMAGE_TO_IMAGE_IMAGE_PARAMS






       def   UpperCamelCase						( self   )->      Dict:
            torch.manual_seed(0   )
            _A				=							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			,     )
            _A				=							DDIMScheduler(
                beta_start=0.0_0085			,     beta_end=0.012			,     beta_schedule='scaled_linear'			,     num_train_timesteps=1000			,     clip_sample=_UpperCamelCase			,     set_alpha_to_one=_UpperCamelCase			,     )
            torch.manual_seed(0   )
            _A				=							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			,     )
            torch.manual_seed(0   )
            _A				=							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			,     )
            _A				=							CLIPTextModel(_UpperCamelCase   )
            _A				=							CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'   )

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






       def   UpperCamelCase						( self			,     _UpperCamelCase			,     _UpperCamelCase=0   )->      List[Any]:
            _A				=							floats_tensor((1, 3, 32, 32)			,     rng=random.Random(_UpperCamelCase   )   ).to(_UpperCamelCase   )
            _A				=							image / 2 + 0.5
            if str(_UpperCamelCase   ).startswith('mps'   ):
                 _A				=							torch.manual_seed(_UpperCamelCase   )
            else:
                 _A				=							torch.Generator(device=_UpperCamelCase   ).manual_seed(_UpperCamelCase   )
            _A				=							{
                'prompt': 'An astronaut riding an elephant',
                'source_prompt': 'An astronaut riding a horse',
                'image': image,
                'generator': generator,
                'num_inference_steps': 2,
                'eta': 0.1,
                'strength': 0.8,
                'guidance_scale': 3,
                'source_guidance_scale': 1,
                'output_type': 'numpy',
            }
            return inputs






       def   UpperCamelCase						( self   )->      Any:
            _A				=							'cpu'  # ensure determinism for the device-dependent torch.Generator

            _A				=							self.get_dummy_components()
            _A				=							CycleDiffusionPipeline(**_UpperCamelCase   )
            _A				=							pipe.to(_UpperCamelCase   )
            pipe.set_progress_bar_config(disable=_UpperCamelCase   )

            _A				=							self.get_dummy_inputs(_UpperCamelCase   )
            _A				=							pipe(**_UpperCamelCase   )
            _A				=							output.images

            _A				=							images[0, -3:, -3:, -1]

            assert images.shape == (1, 32, 32, 3)
            _A				=							np.array([0.4459, 0.4943, 0.4544, 0.6643, 0.5474, 0.4327, 0.5701, 0.5959, 0.5179]   )

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






       @unittest.skipIf(torch_device != 'cuda'			,     'This test requires a GPU'   )
       def   UpperCamelCase						( self   )->      List[Any]:
            _A				=							self.get_dummy_components()
            for name, module in components.items():
                 if hasattr(_UpperCamelCase			,     'half'   ):
                      _A				=							module.half()
            _A				=							CycleDiffusionPipeline(**_UpperCamelCase   )
            _A				=							pipe.to(_UpperCamelCase   )
            pipe.set_progress_bar_config(disable=_UpperCamelCase   )

            _A				=							self.get_dummy_inputs(_UpperCamelCase   )
            _A				=							pipe(**_UpperCamelCase   )
            _A				=							output.images

            _A				=							images[0, -3:, -3:, -1]

            assert images.shape == (1, 32, 32, 3)
            _A				=							np.array([0.3506, 0.4543, 0.446, 0.4575, 0.5195, 0.4155, 0.5273, 0.518, 0.4116]   )

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






       @skip_mps
       def   UpperCamelCase						( self   )->      Any:
            return super().test_save_load_local()






       @unittest.skip('non-deterministic pipeline'   )
       def   UpperCamelCase						( self   )->      List[Any]:
            return super().test_inference_batch_single_identical()






       @skip_mps
       def   UpperCamelCase						( self   )->      int:
            return super().test_dict_tuple_outputs_equivalent()






       @skip_mps
       def   UpperCamelCase						( self   )->      Tuple:
            return super().test_save_load_optional_components()






       @skip_mps
       def   UpperCamelCase						( self   )->      Optional[Any]:
            return super().test_attention_slicing_forward_pass()






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






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






       def   UpperCamelCase						( self   )->      int:
            _A				=							load_image(
                'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
                '/cycle-diffusion/black_colored_car.png'   )
            _A				=							load_numpy(
                'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/cycle-diffusion/blue_colored_car_fp16.npy'   )
            _A				=							init_image.resize((512, 512)   )

            _A				=							'CompVis/stable-diffusion-v1-4'
            _A				=							DDIMScheduler.from_pretrained(_UpperCamelCase			,     subfolder='scheduler'   )
            _A				=							CycleDiffusionPipeline.from_pretrained(
                _UpperCamelCase			,     scheduler=_UpperCamelCase			,     safety_checker=_UpperCamelCase			,     torch_dtype=torch.floataa			,     revision='fp16'   )

            pipe.to(_UpperCamelCase   )
            pipe.set_progress_bar_config(disable=_UpperCamelCase   )
            pipe.enable_attention_slicing()

            _A				=							'A black colored car'
            _A				=							'A blue colored car'

            _A				=							torch.manual_seed(0   )
            _A				=							pipe(
                prompt=_UpperCamelCase			,     source_prompt=_UpperCamelCase			,     image=_UpperCamelCase			,     num_inference_steps=100			,     eta=0.1			,     strength=0.85			,     guidance_scale=3			,     source_guidance_scale=1			,     generator=_UpperCamelCase			,     output_type='np'			,     )
            _A				=							output.images

            # the values aren't exactly equal, but the images look the same visually
            assert np.abs(image - expected_image   ).max() < 5e-1






       def   UpperCamelCase						( self   )->      Union[str, Any]:
            _A				=							load_image(
                'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
                '/cycle-diffusion/black_colored_car.png'   )
            _A				=							load_numpy(
                'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/cycle-diffusion/blue_colored_car.npy'   )
            _A				=							init_image.resize((512, 512)   )

            _A				=							'CompVis/stable-diffusion-v1-4'
            _A				=							DDIMScheduler.from_pretrained(_UpperCamelCase			,     subfolder='scheduler'   )
            _A				=							CycleDiffusionPipeline.from_pretrained(_UpperCamelCase			,     scheduler=_UpperCamelCase			,     safety_checker=_UpperCamelCase   )

            pipe.to(_UpperCamelCase   )
            pipe.set_progress_bar_config(disable=_UpperCamelCase   )
            pipe.enable_attention_slicing()

            _A				=							'A black colored car'
            _A				=							'A blue colored car'

            _A				=							torch.manual_seed(0   )
            _A				=							pipe(
                prompt=_UpperCamelCase			,     source_prompt=_UpperCamelCase			,     image=_UpperCamelCase			,     num_inference_steps=100			,     eta=0.1			,     strength=0.85			,     guidance_scale=3			,     source_guidance_scale=1			,     generator=_UpperCamelCase			,     output_type='np'			,     )
            _A				=							output.images

            assert np.abs(image - expected_image   ).max() < 2e-2





