





'''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							,     lowerCAmelCase							,     lowerCAmelCase  ):





						"""simple docstring"""


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

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





						"""simple docstring"""


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

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

						if inputs_are_torch:
												_lowerCAmelCase												=							torch.from_numpy(lowerCAmelCase  ).to(lowerCAmelCase  )

						return va
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase  ):





						"""simple docstring"""


						_lowerCAmelCase												=							F.normalize(lowerCAmelCase							,     dim=-1  )
						_lowerCAmelCase												=							F.normalize(lowerCAmelCase							,     dim=-1  )
						return (x - y).norm(dim=-1  ).div(2  ).arcsin().pow(2  ).mul(2  )
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase  ):





						"""simple docstring"""


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





class 			UpperCAmelCase					(      snake_case_       ):




						def __init__(      self      :			Dict      ,      __snake_case      :			AutoencoderKL      ,      __snake_case      :			CLIPTextModel      ,      __snake_case      :			CLIPModel      ,      __snake_case      :			CLIPTokenizer      ,      __snake_case      :			UNetaDConditionModel      ,      __snake_case      :			Union[PNDMScheduler, LMSDiscreteScheduler, DDIMScheduler, DPMSolverMultistepScheduler]      ,      __snake_case      :			CLIPFeatureExtractor      ,      __snake_case      :			str=None      ,      __snake_case      :			List[str]=None      ,      __snake_case      :			str=None      ,      )    ->				Union[str, Any]:
												super().__init__()
												self.register_modules(
												    vae=__snake_case      ,      text_encoder=__snake_case      ,      clip_model=__snake_case      ,      tokenizer=__snake_case      ,      unet=__snake_case      ,      scheduler=__snake_case      ,      feature_extractor=__snake_case      ,      coca_model=__snake_case      ,      coca_tokenizer=__snake_case      ,      coca_transform=__snake_case      ,      )
												_lowerCAmelCase												=							(
												    feature_extractor.size
												    if isinstance(feature_extractor.size      ,      __snake_case  )
												    else feature_extractor.size["""shortest_edge"""]
												)
												_lowerCAmelCase												=							transforms.Normalize(mean=feature_extractor.image_mean      ,      std=feature_extractor.image_std  )
												set_requires_grad(self.text_encoder      ,      __snake_case  )
												set_requires_grad(self.clip_model      ,      __snake_case  )




						def     lowercase__   (      self      :			Dict      ,      __snake_case      :			Optional[Union[str, int]] = "auto"  )    ->				Tuple:
												if slice_size == "auto":
																		# half the attention head size is usually a good trade-off between
																		# speed and memory
																		_lowerCAmelCase												=							self.unet.config.attention_head_dim // 2
												self.unet.set_attention_slice(__snake_case  )




						def     lowercase__   (      self      :			Tuple  )    ->				Dict:
												self.enable_attention_slicing(__snake_case  )




						def     lowercase__   (      self      :			Dict  )    ->				List[Any]:
												set_requires_grad(self.vae      ,      __snake_case  )




						def     lowercase__   (      self      :			Dict  )    ->				Tuple:
												set_requires_grad(self.vae      ,      __snake_case  )




						def     lowercase__   (      self      :			Dict  )    ->				str:
												set_requires_grad(self.unet      ,      __snake_case  )




						def     lowercase__   (      self      :			Dict  )    ->				List[str]:
												set_requires_grad(self.unet      ,      __snake_case  )




						def     lowercase__   (      self      :			Any      ,      __snake_case      :			List[Any]      ,      __snake_case      :			Union[str, Any]      ,      __snake_case      :			Any  )    ->				List[Any]:
												# get the original timestep using init_timestep
												_lowerCAmelCase												=							min(int(num_inference_steps * strength  )      ,      __snake_case  )

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

												return timesteps, num_inference_steps - t_start




						def     lowercase__   (      self      :			str      ,      __snake_case      :			str      ,      __snake_case      :			Any      ,      __snake_case      :			List[Any]      ,      __snake_case      :			int      ,      __snake_case      :			str      ,      __snake_case      :			Union[str, Any]=None  )    ->				List[str]:
												if not isinstance(__snake_case      ,      torch.Tensor  ):
																		raise ValueError(f"`image` has to be of type `torch.Tensor` but is {type(__snake_case  )}"  )

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

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

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

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

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

												return latents




						def     lowercase__   (      self      :			int      ,      __snake_case      :			Union[str, Any]  )    ->				int:
												_lowerCAmelCase												=							self.coca_transform(__snake_case  ).unsqueeze(0  )
												with torch.no_grad(), torch.cuda.amp.autocast():
																		_lowerCAmelCase												=							self.coca_model.generate(transformed_image.to(device=self.device      ,      dtype=self.coca_model.dtype  )  )
												_lowerCAmelCase												=							self.coca_tokenizer.decode(generated[0].cpu().numpy()  )
												return generated.split("""<end_of_text>"""  )[0].replace("""<start_of_text>"""      ,      """"""  ).rstrip(""" .,"""  )




						def     lowercase__   (      self      :			List[Any]      ,      __snake_case      :			int      ,      __snake_case      :			int  )    ->				Optional[Any]:
												_lowerCAmelCase												=							self.feature_extractor.preprocess(__snake_case  )
												_lowerCAmelCase												=							torch.from_numpy(clip_image_input["""pixel_values"""][0]  ).unsqueeze(0  ).to(self.device  ).half()
												_lowerCAmelCase												=							self.clip_model.get_image_features(__snake_case  )
												_lowerCAmelCase												=							image_embeddings_clip / image_embeddings_clip.norm(p=2      ,      dim=-1      ,      keepdim=__snake_case  )
												_lowerCAmelCase												=							image_embeddings_clip.repeat_interleave(__snake_case      ,      dim=0  )
												return image_embeddings_clip




						@torch.enable_grad()
						def     lowercase__   (      self      :			Any      ,      __snake_case      :			List[Any]      ,      __snake_case      :			List[Any]      ,      __snake_case      :			Optional[Any]      ,      __snake_case      :			str      ,      __snake_case      :			Dict      ,      __snake_case      :			List[Any]      ,      __snake_case      :			List[Any]      ,      )    ->				List[str]:
												_lowerCAmelCase												=							latents.detach().requires_grad_()

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

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

												if isinstance(self.scheduler      ,      (PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)  ):
																		_lowerCAmelCase												=							self.scheduler.alphas_cumprod[timestep]
																		_lowerCAmelCase												=							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
																		_lowerCAmelCase												=							(latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5

																		_lowerCAmelCase												=							torch.sqrt(__snake_case  )
																		_lowerCAmelCase												=							pred_original_sample * (fac) + latents * (1 - fac)
												elif isinstance(self.scheduler      ,      __snake_case  ):
																		_lowerCAmelCase												=							self.scheduler.sigmas[index]
																		_lowerCAmelCase												=							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
												_lowerCAmelCase												=							1 / 0.1_82_15 * sample
												_lowerCAmelCase												=							self.vae.decode(__snake_case  ).sample
												_lowerCAmelCase												=							(image / 2 + 0.5).clamp(0      ,      1  )

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

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

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

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

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




						@torch.no_grad()
						def __call__(      self      :			Optional[Any]      ,      __snake_case      :			Union[torch.FloatTensor, PIL.Image.Image]      ,      __snake_case      :			Union[torch.FloatTensor, PIL.Image.Image]      ,      __snake_case      :			Optional[str] = None      ,      __snake_case      :			Optional[str] = None      ,      __snake_case      :			Optional[int] = 5_12      ,      __snake_case      :			Optional[int] = 5_12      ,      __snake_case      :			float = 0.6      ,      __snake_case      :			Optional[int] = 50      ,      __snake_case      :			Optional[float] = 7.5      ,      __snake_case      :			Optional[int] = 1      ,      __snake_case      :			float = 0.0      ,      __snake_case      :			Optional[float] = 1_00      ,      __snake_case      :			Optional[torch.Generator] = None      ,      __snake_case      :			Optional[str] = "pil"      ,      __snake_case      :			bool = True      ,      __snake_case      :			float = 0.8      ,      __snake_case      :			float = 0.1      ,      __snake_case      :			float = 0.1      ,      )    ->				Optional[Any]:
												if isinstance(__snake_case      ,      __snake_case  ) and len(__snake_case  ) != batch_size:
																		raise ValueError(f"You have passed {batch_size} batch_size, but only {len(__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(__snake_case      ,      torch.Generator  ) and batch_size > 1:
																		_lowerCAmelCase												=							[generator] + [None] * (batch_size - 1)

												_lowerCAmelCase												=							[
												    ("""model""", self.coca_model is None),
												    ("""tokenizer""", self.coca_tokenizer is None),
												    ("""transform""", self.coca_transform is None),
												]
												_lowerCAmelCase												=							[x[0] for x in coca_is_none if x[1]]
												_lowerCAmelCase												=							""", """.join(__snake_case  )
												# generate prompts with coca model if prompt is None
												if content_prompt is None:
																		if len(__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."  )
																		_lowerCAmelCase												=							self.get_image_description(__snake_case  )
												if style_prompt is None:
																		if len(__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."  )
																		_lowerCAmelCase												=							self.get_image_description(__snake_case  )

												# get prompt text embeddings for content and style
												_lowerCAmelCase												=							self.tokenizer(
												    __snake_case      ,      padding="""max_length"""      ,      max_length=self.tokenizer.model_max_length      ,      truncation=__snake_case      ,      return_tensors="""pt"""      ,      )
												_lowerCAmelCase												=							self.text_encoder(content_text_input.input_ids.to(self.device  )  )[0]

												_lowerCAmelCase												=							self.tokenizer(
												    __snake_case      ,      padding="""max_length"""      ,      max_length=self.tokenizer.model_max_length      ,      truncation=__snake_case      ,      return_tensors="""pt"""      ,      )
												_lowerCAmelCase												=							self.text_encoder(style_text_input.input_ids.to(self.device  )  )[0]

												_lowerCAmelCase												=							slerp(__snake_case      ,      __snake_case      ,      __snake_case  )

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

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

												self.scheduler.set_timesteps(__snake_case      ,      **__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  )

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

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

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

												_lowerCAmelCase												=							slerp(__snake_case      ,      __snake_case      ,      __snake_case  )

												if clip_guidance_scale > 0:
																		_lowerCAmelCase												=							self.get_clip_image_embeddings(__snake_case      ,      __snake_case  )
																		_lowerCAmelCase												=							self.get_clip_image_embeddings(__snake_case      ,      __snake_case  )
																		_lowerCAmelCase												=							slerp(
																		    __snake_case      ,      __snake_case      ,      __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.
												_lowerCAmelCase												=							guidance_scale > 1.0
												# get unconditional embeddings for classifier free guidance
												if do_classifier_free_guidance:
																		_lowerCAmelCase												=							content_text_input.input_ids.shape[-1]
																		_lowerCAmelCase												=							self.tokenizer([""""""]      ,      padding="""max_length"""      ,      max_length=__snake_case      ,      return_tensors="""pt"""  )
																		_lowerCAmelCase												=							self.text_encoder(uncond_input.input_ids.to(self.device  )  )[0]
																		# duplicate unconditional embeddings for each generation per prompt
																		_lowerCAmelCase												=							uncond_embeddings.repeat_interleave(__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
																		_lowerCAmelCase												=							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`.
												_lowerCAmelCase												=							(batch_size, self.unet.config.in_channels, height // 8, width // 8)
												_lowerCAmelCase												=							text_embeddings.dtype
												if latents is None:
																		if self.device.type == "mps":
																								# randn does not work reproducibly on mps
																								_lowerCAmelCase												=							torch.randn(__snake_case      ,      generator=__snake_case      ,      device="""cpu"""      ,      dtype=__snake_case  ).to(
																								    self.device  )
																		else:
																								_lowerCAmelCase												=							torch.randn(__snake_case      ,      generator=__snake_case      ,      device=self.device      ,      dtype=__snake_case  )
												else:
																		if latents.shape != latents_shape:
																								raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}"  )
																		_lowerCAmelCase												=							latents.to(self.device  )

												# scale the initial noise by the standard deviation required by the scheduler
												_lowerCAmelCase												=							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]
												_lowerCAmelCase												=							"""eta""" in set(inspect.signature(self.scheduler.step  ).parameters.keys()  )
												_lowerCAmelCase												=							{}
												if accepts_eta:
																		_lowerCAmelCase												=							eta

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

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

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

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

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

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

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

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

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

												if not return_dict:
																		return (image, None)

												return StableDiffusionPipelineOutput(images=__snake_case      ,      nsfw_content_detected=__snake_case  )













'''simple docstring'''






import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizer

from diffusers import (
    AltDiffusionImgaImgPipeline,
    AutoencoderKL,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
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


enable_full_determinism()




class 			UpperCAmelCase					(      unittest.TestCase       ):




      def     lowercase__   (      self      :			Optional[int]  )    ->				Dict:
            # clean up the VRAM after each test
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()




      @property
      def     lowercase__   (      self      :			List[Any]  )    ->				List[str]:
            _lowerCAmelCase												=							1
            _lowerCAmelCase												=							3
            _lowerCAmelCase												=							(32, 32)

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




      @property
      def     lowercase__   (      self      :			int  )    ->				Union[str, Any]:
            torch.manual_seed(0  )
            _lowerCAmelCase												=							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]  )    ->				List[str]:
            torch.manual_seed(0  )
            _lowerCAmelCase												=							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  )    ->				Optional[Any]:
            torch.manual_seed(0  )
            _lowerCAmelCase												=							RobertaSeriesConfig(
                hidden_size=32      ,      project_dim=32      ,      intermediate_size=37      ,      layer_norm_eps=1E-0_5      ,      num_attention_heads=4      ,      num_hidden_layers=5      ,      pad_token_id=1      ,      vocab_size=50_06      ,      )
            return RobertaSeriesModelWithTransformation(__snake_case  )




      @property
      def     lowercase__   (      self      :			Union[str, Any]  )    ->				str:
            def extract(*__snake_case      :			List[Any]      ,      **__snake_case      :			Any  ):




                  class 			UpperCAmelCase					:




                        def __init__(      self      :			Any  )    ->				Any:
                              _lowerCAmelCase												=							torch.ones([0]  )




                        def     lowercase__   (      self      :			Optional[Any]      ,      __snake_case      :			Tuple  )    ->				Dict:
                              self.pixel_values.to(__snake_case  )
                              return self




                  return Out()

            return extract




      def     lowercase__   (      self      :			List[str]  )    ->				Optional[int]:
            _lowerCAmelCase												=							"""cpu"""  # ensure determinism for the device-dependent torch.Generator
            _lowerCAmelCase												=							self.dummy_cond_unet
            _lowerCAmelCase												=							PNDMScheduler(skip_prk_steps=__snake_case  )
            _lowerCAmelCase												=							self.dummy_vae
            _lowerCAmelCase												=							self.dummy_text_encoder
            _lowerCAmelCase												=							XLMRobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-xlm-roberta"""  )
            _lowerCAmelCase												=							77

            _lowerCAmelCase												=							self.dummy_image.to(__snake_case  )
            _lowerCAmelCase												=							init_image / 2 + 0.5

            # make sure here that pndm scheduler skips prk
            _lowerCAmelCase												=							AltDiffusionImgaImgPipeline(
                unet=__snake_case      ,      scheduler=__snake_case      ,      vae=__snake_case      ,      text_encoder=__snake_case      ,      tokenizer=__snake_case      ,      safety_checker=__snake_case      ,      feature_extractor=self.dummy_extractor      ,      )
            _lowerCAmelCase												=							VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor      ,      do_normalize=__snake_case  )
            _lowerCAmelCase												=							alt_pipe.to(__snake_case  )
            alt_pipe.set_progress_bar_config(disable=__snake_case  )

            _lowerCAmelCase												=							"""A painting of a squirrel eating a burger"""
            _lowerCAmelCase												=							torch.Generator(device=__snake_case  ).manual_seed(0  )
            _lowerCAmelCase												=							alt_pipe(
                [prompt]      ,      generator=__snake_case      ,      guidance_scale=6.0      ,      num_inference_steps=2      ,      output_type="""np"""      ,      image=__snake_case      ,      )

            _lowerCAmelCase												=							output.images

            _lowerCAmelCase												=							torch.Generator(device=__snake_case  ).manual_seed(0  )
            _lowerCAmelCase												=							alt_pipe(
                [prompt]      ,      generator=__snake_case      ,      guidance_scale=6.0      ,      num_inference_steps=2      ,      output_type="""np"""      ,      image=__snake_case      ,      return_dict=__snake_case      ,      )[0]

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

            assert image.shape == (1, 32, 32, 3)
            _lowerCAmelCase												=							np.array([0.44_27, 0.37_31, 0.42_49, 0.49_41, 0.45_46, 0.41_48, 0.41_93, 0.46_66, 0.44_99]  )

            assert np.abs(image_slice.flatten() - expected_slice  ).max() < 5E-3
            assert np.abs(image_from_tuple_slice.flatten() - expected_slice  ).max() < 5E-3




      @unittest.skipIf(torch_device != """cuda"""      ,      """This test requires a GPU"""  )
      def     lowercase__   (      self      :			Tuple  )    ->				str:
            _lowerCAmelCase												=							self.dummy_cond_unet
            _lowerCAmelCase												=							PNDMScheduler(skip_prk_steps=__snake_case  )
            _lowerCAmelCase												=							self.dummy_vae
            _lowerCAmelCase												=							self.dummy_text_encoder
            _lowerCAmelCase												=							XLMRobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-xlm-roberta"""  )
            _lowerCAmelCase												=							77

            _lowerCAmelCase												=							self.dummy_image.to(__snake_case  )

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

            # make sure here that pndm scheduler skips prk
            _lowerCAmelCase												=							AltDiffusionImgaImgPipeline(
                unet=__snake_case      ,      scheduler=__snake_case      ,      vae=__snake_case      ,      text_encoder=__snake_case      ,      tokenizer=__snake_case      ,      safety_checker=__snake_case      ,      feature_extractor=self.dummy_extractor      ,      )
            _lowerCAmelCase												=							VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor      ,      do_normalize=__snake_case  )
            _lowerCAmelCase												=							alt_pipe.to(__snake_case  )
            alt_pipe.set_progress_bar_config(disable=__snake_case  )

            _lowerCAmelCase												=							"""A painting of a squirrel eating a burger"""
            _lowerCAmelCase												=							torch.manual_seed(0  )
            _lowerCAmelCase												=							alt_pipe(
                [prompt]      ,      generator=__snake_case      ,      num_inference_steps=2      ,      output_type="""np"""      ,      image=__snake_case      ,      ).images

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




      @unittest.skipIf(torch_device != """cuda"""      ,      """This test requires a GPU"""  )
      def     lowercase__   (      self      :			Optional[Any]  )    ->				Optional[Any]:
            _lowerCAmelCase												=							load_image(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
                """/img2img/sketch-mountains-input.jpg"""  )
            # resize to resolution that is divisible by 8 but not 16 or 32
            _lowerCAmelCase												=							init_image.resize((7_60, 5_04)  )

            _lowerCAmelCase												=							"""BAAI/AltDiffusion"""
            _lowerCAmelCase												=							AltDiffusionImgaImgPipeline.from_pretrained(
                __snake_case      ,      safety_checker=__snake_case      ,      )
            pipe.to(__snake_case  )
            pipe.set_progress_bar_config(disable=__snake_case  )
            pipe.enable_attention_slicing()

            _lowerCAmelCase												=							"""A fantasy landscape, trending on artstation"""

            _lowerCAmelCase												=							torch.manual_seed(0  )
            _lowerCAmelCase												=							pipe(
                prompt=__snake_case      ,      image=__snake_case      ,      strength=0.75      ,      guidance_scale=7.5      ,      generator=__snake_case      ,      output_type="""np"""      ,      )
            _lowerCAmelCase												=							output.images[0]

            _lowerCAmelCase												=							image[2_55:2_58, 3_83:3_86, -1]

            assert image.shape == (5_04, 7_60, 3)
            _lowerCAmelCase												=							np.array([0.93_58, 0.93_97, 0.95_99, 0.99_01, 1.00_00, 1.00_00, 0.98_82, 1.00_00, 1.00_00]  )

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




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




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




      def     lowercase__   (      self      :			Dict  )    ->				Tuple:
            _lowerCAmelCase												=							load_image(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
                """/img2img/sketch-mountains-input.jpg"""  )
            _lowerCAmelCase												=							init_image.resize((7_68, 5_12)  )
            _lowerCAmelCase												=							load_numpy(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy"""  )

            _lowerCAmelCase												=							"""BAAI/AltDiffusion"""
            _lowerCAmelCase												=							AltDiffusionImgaImgPipeline.from_pretrained(
                __snake_case      ,      safety_checker=__snake_case      ,      )
            pipe.to(__snake_case  )
            pipe.set_progress_bar_config(disable=__snake_case  )
            pipe.enable_attention_slicing()

            _lowerCAmelCase												=							"""A fantasy landscape, trending on artstation"""

            _lowerCAmelCase												=							torch.manual_seed(0  )
            _lowerCAmelCase												=							pipe(
                prompt=__snake_case      ,      image=__snake_case      ,      strength=0.75      ,      guidance_scale=7.5      ,      generator=__snake_case      ,      output_type="""np"""      ,      )
            _lowerCAmelCase												=							output.images[0]

            assert image.shape == (5_12, 7_68, 3)
            # img2img is flaky across GPUs even in fp32, so using MAE here
            assert np.abs(expected_image - image  ).max() < 1E-2






