





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


UpperCamelCase__          =getLogger(__name__)


UpperCamelCase__          =4
UpperCamelCase__          =7

UpperCamelCase__          ={
    'float16': torch.floataa,
    'float32': torch.floataa,
}






def 			lowerCamelCase__	(__lowerCamelCase		):
     if isinstance(__lowerCamelCase,				torch.Tensor		):
          return image
     elif isinstance(__lowerCamelCase,				PIL.Image.Image		):
          _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							[image]

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

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

     return image





class 	lowerCAmelCase__(       __lowercase			):






       '''simple docstring'''


       def __init__(					self			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase = 3_5_0			,  )	->   Any:
            super().__init__(
                vae=__lowerCamelCase			,  text_encoder=__lowerCamelCase			,  tokenizer=__lowerCamelCase			,  unet=__lowerCamelCase			,  low_res_scheduler=__lowerCamelCase			,  scheduler=__lowerCamelCase			,  safety_checker=__lowerCamelCase			,  feature_extractor=__lowerCamelCase			,  watermarker=__lowerCamelCase			,  max_noise_level=__lowerCamelCase			,  )




       def __call__(					self			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase = 7_5			,  __lowerCamelCase = 9.0			,  __lowerCamelCase = 2_0			,  __lowerCamelCase = None			,  __lowerCamelCase = 1			,  __lowerCamelCase = 0.0			,  __lowerCamelCase = None			,  __lowerCamelCase = None			,  __lowerCamelCase = None			,  __lowerCamelCase = None			,  __lowerCamelCase = "pil"			,  __lowerCamelCase = True			,  __lowerCamelCase = None			,  __lowerCamelCase = 1			,  )	->   List[Any]:
            self.check_inputs(__lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase							)

            # 2. Define call parameters
            if prompt is not None and isinstance(__lowerCamelCase			,  __lowerCamelCase							):
                 _SCREAMING_SNAKE_CASE				:		Union[str, Any]  			=							1
            elif prompt is not None and isinstance(__lowerCamelCase			,  __lowerCamelCase							):
                 _SCREAMING_SNAKE_CASE				:		List[Any]  			=							len(__lowerCamelCase							)
            else:
                 _SCREAMING_SNAKE_CASE				:		Optional[Any]  			=							prompt_embeds.shape[0]

            _SCREAMING_SNAKE_CASE				:		Any  			=							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.
            _SCREAMING_SNAKE_CASE				:		Any  			=							guidance_scale > 1.0

            # 3. Encode input prompt
            _SCREAMING_SNAKE_CASE				:		Any  			=							self._encode_prompt(
                __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase			,  prompt_embeds=__lowerCamelCase			,  negative_prompt_embeds=__lowerCamelCase			,  )

            _SCREAMING_SNAKE_CASE				:		Tuple  			=							ORT_TO_PT_TYPE[str(text_embeddings.dtype							)]

            # 4. Preprocess image
            _SCREAMING_SNAKE_CASE				:		int  			=							preprocess(__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE				:		List[Any]  			=							image.cpu()

            # 5. set timesteps
            self.scheduler.set_timesteps(__lowerCamelCase			,  device=__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							self.scheduler.timesteps

            # 5. Add noise to image
            _SCREAMING_SNAKE_CASE				:		int  			=							torch.tensor([noise_level]			,  dtype=torch.long			,  device=__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE				:		int  			=							torch.randn(image.shape			,  generator=__lowerCamelCase			,  device=__lowerCamelCase			,  dtype=__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE				:		Dict  			=							self.low_res_scheduler.add_noise(__lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase							)

            _SCREAMING_SNAKE_CASE				:		Dict  			=							2 if do_classifier_free_guidance else 1
            _SCREAMING_SNAKE_CASE				:		Dict  			=							np.concatenate([image] * batch_multiplier * num_images_per_prompt							)
            _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							np.concatenate([noise_level] * image.shape[0]							)

            # 6. Prepare latent variables
            _SCREAMING_SNAKE_CASE			,	_SCREAMING_SNAKE_CASE				:		Any  			=							image.shape[2:]
            _SCREAMING_SNAKE_CASE				:		List[Any]  			=							self.prepare_latents(
                batch_size * num_images_per_prompt			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase			,  )

            # 7. Check that sizes of image and latents match
            _SCREAMING_SNAKE_CASE				:		List[Any]  			=							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
            _SCREAMING_SNAKE_CASE				:		int  			=							self.prepare_extra_step_kwargs(__lowerCamelCase			,  __lowerCamelCase							)

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

            # 9. Denoising loop
            _SCREAMING_SNAKE_CASE				:		Tuple  			=							len(__lowerCamelCase							) - num_inference_steps * self.scheduler.order
            with self.progress_bar(total=__lowerCamelCase							) as progress_bar:
                 for i, t in enumerate(__lowerCamelCase							):
                      # expand the latents if we are doing classifier free guidance
                      _SCREAMING_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
                      _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							self.scheduler.scale_model_input(__lowerCamelCase			,  __lowerCamelCase							)
                      _SCREAMING_SNAKE_CASE				:		Dict  			=							np.concatenate([latent_model_input, image]			,  axis=1							)

                      # timestep to tensor
                      _SCREAMING_SNAKE_CASE				:		Tuple  			=							np.array([t]			,  dtype=__lowerCamelCase							)

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

                      # perform guidance
                      if do_classifier_free_guidance:
                           _SCREAMING_SNAKE_CASE			,	_SCREAMING_SNAKE_CASE				:		Dict  			=							np.split(__lowerCamelCase			,  2							)
                      _SCREAMING_SNAKE_CASE				:		Tuple  			=							noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

                      # compute the previous noisy sample x_t -> x_t-1
                      _SCREAMING_SNAKE_CASE				:		Dict  			=							self.scheduler.step(
                          torch.from_numpy(__lowerCamelCase							)			,  __lowerCamelCase			,  __lowerCamelCase			,  **__lowerCamelCase							).prev_sample

                      # call the callback, if provided
                      if i == len(__lowerCamelCase							) - 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(__lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase							)

        # 10. Post-processing
            _SCREAMING_SNAKE_CASE				:		Any  			=							self.decode_latents(latents.float()							)

            # 11. Convert to PIL
            if output_type == "pil":
                 _SCREAMING_SNAKE_CASE				:		Any  			=							self.numpy_to_pil(__lowerCamelCase							)

            if not return_dict:
                 return (image,)

            return ImagePipelineOutput(images=__lowerCamelCase							)




       def 			UpperCamelCase_	(					self			,  __lowerCamelCase							)	->   Union[str, Any]:
            _SCREAMING_SNAKE_CASE				:		Any  			=							1 / 0.0_8333 * latents
            _SCREAMING_SNAKE_CASE				:		int  			=							self.vae(latent_sample=__lowerCamelCase							)[0]
            _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							np.clip(image / 2 + 0.5			,  0			,  1							)
            _SCREAMING_SNAKE_CASE				:		Optional[Any]  			=							image.transpose((0, 2, 3, 1)							)
            return image







       def 			UpperCamelCase_	(					self			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase = None			,  __lowerCamelCase = None			,  )	->   Optional[Any]:
            if prompt is not None and isinstance(__lowerCamelCase			,  __lowerCamelCase							):
                 _SCREAMING_SNAKE_CASE				:		Optional[Any]  			=							1
            elif prompt is not None and isinstance(__lowerCamelCase			,  __lowerCamelCase							):
                 _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							len(__lowerCamelCase							)
            else:
                 _SCREAMING_SNAKE_CASE				:		Tuple  			=							prompt_embeds.shape[0]

            if prompt_embeds is None:
                 _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							self.tokenizer(
                     __lowerCamelCase			,  padding="max_length"			,  max_length=self.tokenizer.model_max_length			,  truncation=__lowerCamelCase			,  return_tensors="pt"			,  )
                 _SCREAMING_SNAKE_CASE				:		List[str]  			=							text_inputs.input_ids
                 _SCREAMING_SNAKE_CASE				:		Any  			=							self.tokenizer(__lowerCamelCase			,  padding="longest"			,  return_tensors="pt"							).input_ids

                 if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
                     __lowerCamelCase			,  __lowerCamelCase							):
                      _SCREAMING_SNAKE_CASE				:		Optional[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
                 _SCREAMING_SNAKE_CASE				:		Any  			=							self.text_encoder(
                     input_ids=text_input_ids.int().to(__lowerCamelCase							)			,  )
                 _SCREAMING_SNAKE_CASE				:		List[str]  			=							prompt_embeds[0]

            _SCREAMING_SNAKE_CASE			,	_SCREAMING_SNAKE_CASE			,	_SCREAMING_SNAKE_CASE				:		List[str]  			=							prompt_embeds.shape
            # duplicate text embeddings for each generation per prompt, using mps friendly method
            _SCREAMING_SNAKE_CASE				:		int  			=							prompt_embeds.repeat(1			,  __lowerCamelCase							)
            _SCREAMING_SNAKE_CASE				:		str  			=							prompt_embeds.reshape(bs_embed * num_images_per_prompt			,  __lowerCamelCase			,  -1							)

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

                 _SCREAMING_SNAKE_CASE				:		List[str]  			=							text_input_ids.shape[-1]
                 _SCREAMING_SNAKE_CASE				:		Optional[Any]  			=							self.tokenizer(
                     __lowerCamelCase			,  padding="max_length"			,  max_length=__lowerCamelCase			,  truncation=__lowerCamelCase			,  return_tensors="pt"			,  )

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

                 _SCREAMING_SNAKE_CASE				:		List[str]  			=							self.text_encoder(
                     input_ids=uncond_input.input_ids.int().to(__lowerCamelCase							)			,  )
                 _SCREAMING_SNAKE_CASE				:		Tuple  			=							uncond_embeddings[0]

            if do_classifier_free_guidance:
                 _SCREAMING_SNAKE_CASE				:		int  			=							uncond_embeddings.shape[1]
                 # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
                 _SCREAMING_SNAKE_CASE				:		Dict  			=							uncond_embeddings.repeat(1			,  __lowerCamelCase							)
                 _SCREAMING_SNAKE_CASE				:		List[str]  			=							uncond_embeddings.reshape(batch_size * num_images_per_prompt			,  __lowerCamelCase			,  -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
                 _SCREAMING_SNAKE_CASE				:		str  			=							np.concatenate([uncond_embeddings, prompt_embeds]							)

            return prompt_embeds






import warnings
from dataclasses import dataclass, field
from typing import Optional, Tuple

from .training_args import TrainingArguments
from .utils import cached_property, is_tf_available, logging, requires_backends


UpperCamelCase__          =logging.get_logger(__name__)

if is_tf_available():
       import tensorflow as tf





@dataclass
class 	lowerCAmelCase__(       __lowercase			):






       '''simple docstring'''


       __snake_case       	=			'tf'
       __snake_case       	=			field(
           default=__lowercase     ,   metadata={'help': 'Name of TPU'}     ,   )

       __snake_case       	=			field(
           default=__lowercase     ,   metadata={'help': 'Zone of TPU'}     ,   )

       __snake_case       	=			field(
           default=__lowercase     ,   metadata={'help': 'Name of Cloud TPU-enabled project'}     ,   )

       __snake_case       	=			field(
           default=1.0     ,   metadata={'help': 'Power for the Polynomial decay LR scheduler.'}     ,   )

       __snake_case       	=			field(default=__lowercase     ,   metadata={'help': 'Whether to activate the XLA compilation or not'}			)




       @cached_property
       def 			UpperCamelCase_	(					self							)	->   Tuple["tf.distribute.Strategy", int]:
            requires_backends(self			,  ["tf"]							)
            logger.info("Tensorflow: setting up strategy"							)

            _SCREAMING_SNAKE_CASE				:		List[Any]  			=							tf.config.list_physical_devices("GPU"							)

            # Set to float16 at first
            if self.fpaa:
                 tf.keras.mixed_precision.set_global_policy("mixed_float16"							)

            if self.no_cuda:
                 _SCREAMING_SNAKE_CASE				:		List[Any]  			=							tf.distribute.OneDeviceStrategy(device="/cpu:0"							)
            else:
                 try:
                      if self.tpu_name:
                           _SCREAMING_SNAKE_CASE				:		Any  			=							tf.distribute.cluster_resolver.TPUClusterResolver(
                               self.tpu_name			,  zone=self.tpu_zone			,  project=self.gcp_project							)
                      else:
                           _SCREAMING_SNAKE_CASE				:		str  			=							tf.distribute.cluster_resolver.TPUClusterResolver()
                 except ValueError:
                      if self.tpu_name:
                           raise RuntimeError(F"""Couldn't connect to TPU {self.tpu_name}!"""							)
                      else:
                           _SCREAMING_SNAKE_CASE				:		Any  			=							None

                 if tpu:
                      # Set to bfloat16 in case of TPU
                      if self.fpaa:
                           tf.keras.mixed_precision.set_global_policy("mixed_bfloat16"							)

                      tf.config.experimental_connect_to_cluster(__lowerCamelCase							)
                      tf.tpu.experimental.initialize_tpu_system(__lowerCamelCase							)

                      _SCREAMING_SNAKE_CASE				:		List[str]  			=							tf.distribute.TPUStrategy(__lowerCamelCase							)

                 elif len(__lowerCamelCase							) == 0:
                      _SCREAMING_SNAKE_CASE				:		Dict  			=							tf.distribute.OneDeviceStrategy(device="/cpu:0"							)
                 elif len(__lowerCamelCase							) == 1:
                      _SCREAMING_SNAKE_CASE				:		Dict  			=							tf.distribute.OneDeviceStrategy(device="/gpu:0"							)
                 elif len(__lowerCamelCase							) > 1:
                      # If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0`
                      _SCREAMING_SNAKE_CASE				:		Any  			=							tf.distribute.MirroredStrategy()
                 else:
                      raise ValueError("Cannot find the proper strategy, please check your environment properties."							)

            return strategy




       @property
       def 			UpperCamelCase_	(					self							)	->   "tf.distribute.Strategy":
            requires_backends(self			,  ["tf"]							)
            return self._setup_strategy




       @property
       def 			UpperCamelCase_	(					self							)	->   int:
            requires_backends(self			,  ["tf"]							)
            return self._setup_strategy.num_replicas_in_sync




       @property
       def 			UpperCamelCase_	(					self							)	->   List[Any]:
            return False  # TF Logging is handled by Keras not the Trainer




       @property
       def 			UpperCamelCase_	(					self							)	->   int:
            if self.per_gpu_train_batch_size:
                 logger.warning(
                     "Using deprecated `--per_gpu_train_batch_size` argument which will be removed in a future "
                     "version. Using `--per_device_train_batch_size` is preferred."							)
            _SCREAMING_SNAKE_CASE				:		Tuple  			=							self.per_gpu_train_batch_size or self.per_device_train_batch_size
            return per_device_batch_size * self.n_replicas




       @property
       def 			UpperCamelCase_	(					self							)	->   int:
            if self.per_gpu_eval_batch_size:
                 logger.warning(
                     "Using deprecated `--per_gpu_eval_batch_size` argument which will be removed in a future "
                     "version. Using `--per_device_eval_batch_size` is preferred."							)
            _SCREAMING_SNAKE_CASE				:		int  			=							self.per_gpu_eval_batch_size or self.per_device_eval_batch_size
            return per_device_batch_size * self.n_replicas







       @property
       def 			UpperCamelCase_	(					self							)	->   int:
            requires_backends(self			,  ["tf"]							)
            warnings.warn(
                "The n_gpu argument is deprecated and will be removed in a future version, use n_replicas instead."			,  __lowerCamelCase			,  )
            return self._setup_strategy.num_replicas_in_sync