from logging import getLogger
from typing import Any, Callable, List, Optional, Union

import numpy as np
import PIL
import torch

from ...schedulers import DDPMScheduler
from ..onnx_utils import ORT_TO_NP_TYPE, OnnxRuntimeModel
from ..pipeline_utils import ImagePipelineOutput
from . import StableDiffusionUpscalePipeline


lowerCAmelCase__	:  int						=  getLogger(__name__)


lowerCAmelCase__	:  Any						=  4
lowerCAmelCase__	:  Any						=  7

lowerCAmelCase__	:  int						=  {
    '''float16''': torch.floataa,
    '''float32''': torch.floataa,
}
def 					UpperCamelCase__		(							A__       )  ->   Tuple:
				if isinstance(A__     ,					torch.Tensor       ):
								return image
				elif isinstance(A__     ,					PIL.Image.Image       ):
								snake_case__   :      List[str]   				= [image]

				if isinstance(image[0]     ,					PIL.Image.Image       ):
								snake_case__				,    snake_case__   :      Optional[Any]   				= image[0].size
								snake_case__				,    snake_case__   :      Dict   				= (x - x % 64 for x in (w, h))  # resize to integer multiple of 32

								snake_case__   :      Tuple   				= [np.array(i.resize((w, h)       )       )[None, :] for i in image]
								snake_case__   :      List[str]   				= np.concatenate(A__     ,					axis=0       )
								snake_case__   :      List[Any]   				= np.array(A__       ).astype(np.floataa       ) / 2_5_5.0
								snake_case__   :      Tuple   				= image.transpose(0     ,					3     ,					1     ,					2       )
								snake_case__   :      Optional[int]   				= 2.0 * image - 1.0
								snake_case__   :      Any   				= torch.from_numpy(A__       )
				elif isinstance(image[0]     ,					torch.Tensor       ):
								snake_case__   :      Any   				= torch.cat(A__     ,					dim=0       )

				return image





class 					__snake_case    (     _lowerCamelCase      ):




					def __init__(  self       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase = 350       ,			)	->    Optional[int]:



									'''simple docstring'''







									super().__init__(
									    vae=__UpperCamelCase       ,			text_encoder=__UpperCamelCase       ,			tokenizer=__UpperCamelCase       ,			unet=__UpperCamelCase       ,			low_res_scheduler=__UpperCamelCase       ,			scheduler=__UpperCamelCase       ,			safety_checker=__UpperCamelCase       ,			feature_extractor=__UpperCamelCase       ,			watermarker=__UpperCamelCase       ,			max_noise_level=__UpperCamelCase       ,			)




					def __call__(  self       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase = 75       ,			__UpperCamelCase = 9.0       ,			__UpperCamelCase = 20       ,			__UpperCamelCase = None       ,			__UpperCamelCase = 1       ,			__UpperCamelCase = 0.0       ,			__UpperCamelCase = None       ,			__UpperCamelCase = None       ,			__UpperCamelCase = None       ,			__UpperCamelCase = None       ,			__UpperCamelCase = "pil"       ,			__UpperCamelCase = True       ,			__UpperCamelCase = None       ,			__UpperCamelCase = 1       ,			)	->    Tuple:



									'''simple docstring'''







									self.check_inputs(__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase						)

									# 2. Define call parameters
									if prompt is not None and isinstance(__UpperCamelCase       ,			__UpperCamelCase						):
													snake_case__   :      Dict   				= 1
									elif prompt is not None and isinstance(__UpperCamelCase       ,			__UpperCamelCase						):
													snake_case__   :      List[str]   				= len(__UpperCamelCase						)
									else:
													snake_case__   :      List[str]   				= prompt_embeds.shape[0]

									snake_case__   :      List[str]   				= self._execution_device
									# 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__   :      str   				= guidance_scale > 1.0

									# 3. Encode input prompt
									snake_case__   :      Optional[int]   				= self._encode_prompt(
									    __UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			prompt_embeds=__UpperCamelCase       ,			negative_prompt_embeds=__UpperCamelCase       ,			)

									snake_case__   :      str   				= ORT_TO_PT_TYPE[str(text_embeddings.dtype						)]

									# 4. Preprocess image
									snake_case__   :      str   				= preprocess(__UpperCamelCase						)
									snake_case__   :      Any   				= image.cpu()

									# 5. set timesteps
									self.scheduler.set_timesteps(__UpperCamelCase       ,			device=__UpperCamelCase						)
									snake_case__   :      List[str]   				= self.scheduler.timesteps

									# 5. Add noise to image
									snake_case__   :      Union[str, Any]   				= torch.tensor([noise_level]       ,			dtype=torch.long       ,			device=__UpperCamelCase						)
									snake_case__   :      Optional[Any]   				= torch.randn(image.shape       ,			generator=__UpperCamelCase       ,			device=__UpperCamelCase       ,			dtype=__UpperCamelCase						)
									snake_case__   :      Any   				= self.low_res_scheduler.add_noise(__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase						)

									snake_case__   :      List[Any]   				= 2 if do_classifier_free_guidance else 1
									snake_case__   :      Tuple   				= np.concatenate([image] * batch_multiplier * num_images_per_prompt						)
									snake_case__   :      Tuple   				= np.concatenate([noise_level] * image.shape[0]						)

									# 6. Prepare latent variables
									snake_case__				,    snake_case__   :      Optional[Any]   				= image.shape[2:]
									snake_case__   :      List[str]   				= self.prepare_latents(
									    batch_size * num_images_per_prompt       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			)

									# 7. Check that sizes of image and latents match
									snake_case__   :      Tuple   				= image.shape[1]
									if NUM_LATENT_CHANNELS + num_channels_image != NUM_UNET_INPUT_CHANNELS:
													raise ValueError(
													    'Incorrect configuration settings! The config of `pipeline.unet` expects'
													    F""" {NUM_UNET_INPUT_CHANNELS} but received `num_channels_latents`: {NUM_LATENT_CHANNELS} +"""
													    F""" `num_channels_image`: {num_channels_image} """
													    F""" = {NUM_LATENT_CHANNELS+num_channels_image}. Please verify the config of"""
													    ' `pipeline.unet` or your `image` input.'						)

									# 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
									snake_case__   :      List[Any]   				= self.prepare_extra_step_kwargs(__UpperCamelCase       ,			__UpperCamelCase						)

									snake_case__   :      int   				= next(
									    (input.type for input in self.unet.model.get_inputs() if input.name == 'timestep')       ,			'tensor(float)'						)
									snake_case__   :      List[str]   				= ORT_TO_NP_TYPE[timestep_dtype]

									# 9. Denoising loop
									snake_case__   :      str   				= len(__UpperCamelCase						) - num_inference_steps * self.scheduler.order
									with self.progress_bar(total=__UpperCamelCase						) as progress_bar:
													for i, t in enumerate(__UpperCamelCase						):
																	# expand the latents if we are doing classifier free guidance
																	snake_case__   :      List[str]   				= np.concatenate([latents] * 2						) if do_classifier_free_guidance else latents

																	# concat latents, mask, masked_image_latents in the channel dimension
																	snake_case__   :      Dict   				= self.scheduler.scale_model_input(__UpperCamelCase       ,			__UpperCamelCase						)
																	snake_case__   :      str   				= np.concatenate([latent_model_input, image]       ,			axis=1						)

																	# timestep to tensor
																	snake_case__   :      str   				= np.array([t]       ,			dtype=__UpperCamelCase						)

																	# predict the noise residual
																	snake_case__   :      Union[str, Any]   				= self.unet(
																	    sample=__UpperCamelCase       ,			timestep=__UpperCamelCase       ,			encoder_hidden_states=__UpperCamelCase       ,			class_labels=noise_level.astype(np.intaa						)       ,			)[0]

																	# perform guidance
																	if do_classifier_free_guidance:
																					snake_case__				,    snake_case__   :      Optional[int]   				= np.split(__UpperCamelCase       ,			2						)
																	snake_case__   :      List[Any]   				= noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

																	# compute the previous noisy sample x_t -> x_t-1
																	snake_case__   :      Optional[Any]   				= self.scheduler.step(
																	    torch.from_numpy(__UpperCamelCase						)       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			**__UpperCamelCase						).prev_sample

																	# call the callback, if provided
																	if i == len(__UpperCamelCase						) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
																					progress_bar.update()
																					if callback is not None and i % callback_steps == 0:
																									callback(__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase						)

        # 10. Post-processing
									snake_case__   :      Optional[int]   				= self.decode_latents(latents.float()						)

									# 11. Convert to PIL
									if output_type == "pil":
													snake_case__   :      Tuple   				= self.numpy_to_pil(__UpperCamelCase						)

									if not return_dict:
													return (image,)

									return ImagePipelineOutput(images=__UpperCamelCase						)




					def      __a     (  self       ,			__UpperCamelCase						)	->    Optional[Any]:



									'''simple docstring'''







									snake_case__   :      List[str]   				= 1 / 0.0_8_3_3_3 * latents
									snake_case__   :      Any   				= self.vae(latent_sample=__UpperCamelCase						)[0]
									snake_case__   :      Union[str, Any]   				= np.clip(image / 2 + 0.5       ,			0       ,			1						)
									snake_case__   :      str   				= image.transpose((0, 2, 3, 1)						)
									return image




					def      __a     (  self       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase = None       ,			__UpperCamelCase = None       ,			)	->    Optional[int]:



									'''simple docstring'''







									if prompt is not None and isinstance(__UpperCamelCase       ,			__UpperCamelCase						):
													snake_case__   :      List[Any]   				= 1
									elif prompt is not None and isinstance(__UpperCamelCase       ,			__UpperCamelCase						):
													snake_case__   :      Any   				= len(__UpperCamelCase						)
									else:
													snake_case__   :      Any   				= prompt_embeds.shape[0]

									if prompt_embeds is None:
													snake_case__   :      List[Any]   				= self.tokenizer(
													    __UpperCamelCase       ,			padding='max_length'       ,			max_length=self.tokenizer.model_max_length       ,			truncation=__UpperCamelCase       ,			return_tensors='pt'       ,			)
													snake_case__   :      Any   				= text_inputs.input_ids
													snake_case__   :      str   				= self.tokenizer(__UpperCamelCase       ,			padding='longest'       ,			return_tensors='pt'						).input_ids

													if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
													    __UpperCamelCase       ,			__UpperCamelCase						):
																	snake_case__   :      Any   				= self.tokenizer.batch_decode(
																	    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]						)
																	logger.warning(
																	    'The following part of your input was truncated because CLIP can only handle sequences up to'
																	    F""" {self.tokenizer.model_max_length} tokens: {removed_text}"""						)

													# no positional arguments to text_encoder
													snake_case__   :      List[str]   				= self.text_encoder(
													    input_ids=text_input_ids.int().to(__UpperCamelCase						)       ,			)
													snake_case__   :      Optional[int]   				= prompt_embeds[0]

									snake_case__				,    snake_case__				,    snake_case__   :      int   				= prompt_embeds.shape
									# duplicate text embeddings for each generation per prompt, using mps friendly method
									snake_case__   :      Optional[int]   				= prompt_embeds.repeat(1       ,			__UpperCamelCase						)
									snake_case__   :      Optional[Any]   				= prompt_embeds.reshape(bs_embed * num_images_per_prompt       ,			__UpperCamelCase       ,			-1						)

									# get unconditional embeddings for classifier free guidance
									if do_classifier_free_guidance and negative_prompt_embeds is None:
													snake_case__   :      List[str]
													if negative_prompt is None:
																	snake_case__   :      Union[str, Any]   				= [''] * batch_size
													elif type(__UpperCamelCase						) is not type(__UpperCamelCase						):
																	raise TypeError(
																	    F"""`negative_prompt` should be the same type to `prompt`, but got {type(__UpperCamelCase						)} !="""
																	    F""" {type(__UpperCamelCase						)}."""						)
													elif isinstance(__UpperCamelCase       ,			__UpperCamelCase						):
																	snake_case__   :      str   				= [negative_prompt]
													elif batch_size != len(__UpperCamelCase						):
																	raise ValueError(
																	    F"""`negative_prompt`: {negative_prompt} has batch size {len(__UpperCamelCase						)}, but `prompt`:"""
																	    F""" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"""
																	    ' the batch size of `prompt`.'						)
													else:
																	snake_case__   :      Optional[int]   				= negative_prompt

													snake_case__   :      Dict   				= text_input_ids.shape[-1]
													snake_case__   :      Dict   				= self.tokenizer(
													    __UpperCamelCase       ,			padding='max_length'       ,			max_length=__UpperCamelCase       ,			truncation=__UpperCamelCase       ,			return_tensors='pt'       ,			)

													# if hasattr(uncond_input, "attention_mask"):
													#     attention_mask = uncond_input.attention_mask.to(device)
													# else:
													#     attention_mask = None

													snake_case__   :      List[str]   				= self.text_encoder(
													    input_ids=uncond_input.input_ids.int().to(__UpperCamelCase						)       ,			)
													snake_case__   :      Optional[int]   				= uncond_embeddings[0]

									if do_classifier_free_guidance:
													snake_case__   :      str   				= uncond_embeddings.shape[1]
													# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
													snake_case__   :      Any   				= uncond_embeddings.repeat(1       ,			__UpperCamelCase						)
													snake_case__   :      Optional[Any]   				= uncond_embeddings.reshape(batch_size * num_images_per_prompt       ,			__UpperCamelCase       ,			-1						)

													# 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__   :      Dict   				= np.concatenate([uncond_embeddings, prompt_embeds]						)

									return prompt_embeds



from ...utils import is_transformers_available
from .pipeline_latent_diffusion_superresolution import LDMSuperResolutionPipeline


if is_transformers_available():
				from .pipeline_latent_diffusion import LDMBertModel, LDMTextToImagePipeline


