




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__			(		A__			:   List[Any]      ):





       '''simple docstring'''



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

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

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

       return image





class        lowerCamelCase__(   __lowerCamelCase):
      def __init__(							self:					Optional[Any]       ,			UpperCamelCase_:					OnnxRuntimeModel       ,			UpperCamelCase_:					OnnxRuntimeModel       ,			UpperCamelCase_:					Any       ,			UpperCamelCase_:					OnnxRuntimeModel       ,			UpperCamelCase_:					DDPMScheduler       ,			UpperCamelCase_:					Any       ,			UpperCamelCase_:					int = 3_50       ,			):
             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:					List[Any]       ,			UpperCamelCase_:					Union[str, List[str]]       ,			UpperCamelCase_:					Union[torch.FloatTensor, PIL.Image.Image, List[PIL.Image.Image]]       ,			UpperCamelCase_:					int = 75       ,			UpperCamelCase_:					float = 9.0       ,			UpperCamelCase_:					int = 20       ,			UpperCamelCase_:					Optional[Union[str, List[str]]] = None       ,			UpperCamelCase_:					Optional[int] = 1       ,			UpperCamelCase_:					float = 0.0       ,			UpperCamelCase_:					Optional[Union[torch.Generator, List[torch.Generator]]] = None       ,			UpperCamelCase_:					Optional[torch.FloatTensor] = None       ,			UpperCamelCase_:					Optional[np.ndarray] = None       ,			UpperCamelCase_:					Optional[np.ndarray] = None       ,			UpperCamelCase_:					Optional[str] = "pil"       ,			UpperCamelCase_:					bool = True       ,			UpperCamelCase_:					Optional[Callable[[int, int, torch.FloatTensor], None]] = None       ,			UpperCamelCase_:					Optional[int] = 1       ,			):
             self.check_inputs(UpperCamelCase_       ,			UpperCamelCase_       ,			UpperCamelCase_       ,			UpperCamelCase_  )

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

             __lowerCamelCase  =		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.
             __lowerCamelCase  =		guidance_scale > 1.0

             # 3. Encode input prompt
             __lowerCamelCase  =		self._encode_prompt(
                 UpperCamelCase_       ,			UpperCamelCase_       ,			UpperCamelCase_       ,			UpperCamelCase_       ,			UpperCamelCase_       ,			prompt_embeds=UpperCamelCase_       ,			negative_prompt_embeds=UpperCamelCase_       ,			)

             __lowerCamelCase  =		ORT_TO_PT_TYPE[str(text_embeddings.dtype  )]

             # 4. Preprocess image
             __lowerCamelCase  =		preprocess(UpperCamelCase_  )
             __lowerCamelCase  =		image.cpu()

             # 5. set timesteps
             self.scheduler.set_timesteps(UpperCamelCase_       ,			device=UpperCamelCase_  )
             __lowerCamelCase  =		self.scheduler.timesteps

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

             __lowerCamelCase  =		2 if do_classifier_free_guidance else 1
             __lowerCamelCase  =		np.concatenate([image] * batch_multiplier * num_images_per_prompt  )
             __lowerCamelCase  =		np.concatenate([noise_level] * image.shape[0]  )

             # 6. Prepare latent variables
             __lowerCamelCase,  __lowerCamelCase  =		image.shape[2:]
             __lowerCamelCase  =		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
             __lowerCamelCase  =		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
             __lowerCamelCase  =		self.prepare_extra_step_kwargs(UpperCamelCase_       ,			UpperCamelCase_  )

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

             # 9. Denoising loop
             __lowerCamelCase  =		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
                           __lowerCamelCase  =		np.concatenate([latents] * 2  ) if do_classifier_free_guidance else latents

                           # concat latents, mask, masked_image_latents in the channel dimension
                           __lowerCamelCase  =		self.scheduler.scale_model_input(UpperCamelCase_       ,			UpperCamelCase_  )
                           __lowerCamelCase  =		np.concatenate([latent_model_input, image]       ,			axis=1  )

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

                           # predict the noise residual
                           __lowerCamelCase  =		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:
                                  __lowerCamelCase,  __lowerCamelCase  =		np.split(UpperCamelCase_       ,			2  )
                           __lowerCamelCase  =		noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

                           # compute the previous noisy sample x_t -> x_t-1
                           __lowerCamelCase  =		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
             __lowerCamelCase  =		self.decode_latents(latents.float()  )

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

             if not return_dict:
                    return (image,)

             return ImagePipelineOutput(images=UpperCamelCase_  )
      def      lowerCAmelCase__		(							self:					int       ,			UpperCamelCase_:					Tuple  ):
             __lowerCamelCase  =		1 / 0.0_8333 * latents
             __lowerCamelCase  =		self.vae(latent_sample=UpperCamelCase_  )[0]
             __lowerCamelCase  =		np.clip(image / 2 + 0.5       ,			0       ,			1  )
             __lowerCamelCase  =		image.transpose((0, 2, 3, 1)  )
             return image




      def      lowerCAmelCase__		(							self:					List[Any]       ,			UpperCamelCase_:					Union[str, List[str]]       ,			UpperCamelCase_:					Optional[Any]       ,			UpperCamelCase_:					Optional[int]       ,			UpperCamelCase_:					bool       ,			UpperCamelCase_:					Optional[str]       ,			UpperCamelCase_:					Optional[torch.FloatTensor] = None       ,			UpperCamelCase_:					Optional[torch.FloatTensor] = None       ,			):
             if prompt is not None and isinstance(UpperCamelCase_       ,			UpperCamelCase_  ):
                    __lowerCamelCase  =		1
             elif prompt is not None and isinstance(UpperCamelCase_       ,			UpperCamelCase_  ):
                    __lowerCamelCase  =		len(UpperCamelCase_  )
             else:
                    __lowerCamelCase  =		prompt_embeds.shape[0]

             if prompt_embeds is None:
                    __lowerCamelCase  =		self.tokenizer(
                        UpperCamelCase_       ,			padding="""max_length"""       ,			max_length=self.tokenizer.model_max_length       ,			truncation=UpperCamelCase_       ,			return_tensors="""pt"""       ,			)
                    __lowerCamelCase  =		text_inputs.input_ids
                    __lowerCamelCase  =		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_  ):
                           __lowerCamelCase  =		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
                    __lowerCamelCase  =		self.text_encoder(
                        input_ids=text_input_ids.int().to(UpperCamelCase_  )       ,			)
                    __lowerCamelCase  =		prompt_embeds[0]

             __lowerCamelCase,  __lowerCamelCase,  __lowerCamelCase  =		prompt_embeds.shape
             # duplicate text embeddings for each generation per prompt, using mps friendly method
             __lowerCamelCase  =		prompt_embeds.repeat(1       ,			UpperCamelCase_  )
             __lowerCamelCase  =		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:
                    __lowerCamelCase  =		42
                    if negative_prompt is None:
                           __lowerCamelCase  =		[""""""] * 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_  ):
                           __lowerCamelCase  =		[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:
                           __lowerCamelCase  =		negative_prompt

                    __lowerCamelCase  =		text_input_ids.shape[-1]
                    __lowerCamelCase  =		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

                    __lowerCamelCase  =		self.text_encoder(
                        input_ids=uncond_input.input_ids.int().to(UpperCamelCase_  )       ,			)
                    __lowerCamelCase  =		uncond_embeddings[0]

             if do_classifier_free_guidance:
                    __lowerCamelCase  =		uncond_embeddings.shape[1]
                    # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
                    __lowerCamelCase  =		uncond_embeddings.repeat(1       ,			UpperCamelCase_  )
                    __lowerCamelCase  =		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
                    __lowerCamelCase  =		np.concatenate([uncond_embeddings, prompt_embeds]  )

             return prompt_embeds







from __future__ import annotations

from collections.abc import Iterable





class        lowerCamelCase__:
      def __init__(							self:					Dict       ,			UpperCamelCase_:					int       ,			UpperCamelCase_:					str  ):
             __lowerCamelCase,  __lowerCamelCase  =		float(UpperCamelCase_  ), float(UpperCamelCase_  )
      def __eq__(							self:					List[Any]       ,			UpperCamelCase_:					List[str]  ):
             return self.x == other.x and self.y == other.y
      def __ne__(							self:					Tuple       ,			UpperCamelCase_:					Dict  ):
             return not self == other
      def __gt__(							self:					List[Any]       ,			UpperCamelCase_:					List[Any]  ):
             if self.x > other.x:
                    return True
             elif self.x == other.x:
                    return self.y > other.y
             return False
      def __lt__(							self:					str       ,			UpperCamelCase_:					int  ):
             return not self > other
      def __ge__(							self:					Dict       ,			UpperCamelCase_:					Any  ):
             if self.x > other.x:
                    return True
             elif self.x == other.x:
                    return self.y >= other.y
             return False
      def __le__(							self:					Union[str, Any]       ,			UpperCamelCase_:					Optional[Any]  ):
             if self.x < other.x:
                    return True
             elif self.x == other.x:
                    return self.y <= other.y
             return False
      def __repr__(							self:					Union[str, Any]  ):
             return F'({self.x}, {self.y})'




      def __hash__(							self:					List[Any]  ):
             return hash(self.x  )




def    lowerCamelCase__			(		A__			:   list[Point] | list[list[float]] | Iterable[list[float]]    ,     ):





       '''simple docstring'''



       __lowerCamelCase  =		[]
       if list_of_tuples:
              for p in list_of_tuples:
                     if isinstance(A__    ,     A__      ):
                            points.append(A__      )
                     else:
                            try:
                                   points.append(Point(p[0]    ,     p[1]      )      )
                            except (IndexError, TypeError):
                                   print(
                                       f'Ignoring deformed point {p}. All points'
                                       """ must have at least 2 coordinates."""      )
       return points




def    lowerCamelCase__			(		A__			:   list[Point] | list[list[float]]      ):





       '''simple docstring'''



       if not hasattr(A__    ,     """__iter__"""      ):
              __lowerCamelCase  =		f'Expecting an iterable object but got an non-iterable type {points}'
              raise ValueError(A__      )

       if not points:
              __lowerCamelCase  =		f'Expecting a list of points but got {points}'
              raise ValueError(A__      )

       return _construct_points(A__      )




def    lowerCamelCase__			(		A__			:   Point    ,     A__			:   Point    ,     A__			:   Point      ):





       '''simple docstring'''



       __lowerCamelCase  =		(a.x * b.y + b.x * c.y + c.x * a.y) - (a.y * b.x + b.y * c.x + c.y * a.x)
       return det




def    lowerCamelCase__			(		A__			:   list[Point]      ):





       '''simple docstring'''



       __lowerCamelCase  =		sorted(_validate_input(A__      )      )
       __lowerCamelCase  =		len(A__      )
       __lowerCamelCase  =		set()

       for i in range(n - 1      ):
              for j in range(i + 1    ,     A__      ):
                     __lowerCamelCase  =		__lowerCamelCase  =		False
                     __lowerCamelCase  =		True
                     for k in range(A__      ):
                            if k != i and k != j:
                                   __lowerCamelCase  =		_det(points[i]    ,     points[j]    ,     points[k]      )

                                   if det_k > 0:
                                          __lowerCamelCase  =		True
                                   elif det_k < 0:
                                          __lowerCamelCase  =		True
                                   else:
                                          # point[i], point[j], point[k] all lie on a straight line
                                          # if point[k] is to the left of point[i] or it's to the
                                          # right of point[j], then point[i], point[j] cannot be
                                          # part of the convex hull of A
                                          if points[k] < points[i] or points[k] > points[j]:
                                                 __lowerCamelCase  =		False
                                                 break

                            if points_left_of_ij and points_right_of_ij:
                                   __lowerCamelCase  =		False
                                   break

                     if ij_part_of_convex_hull:
                            convex_set.update([points[i], points[j]]      )

       return sorted(A__      )




def    lowerCamelCase__			(		A__			:   list[Point]      ):





       '''simple docstring'''



       __lowerCamelCase  =		sorted(_validate_input(A__      )      )
       __lowerCamelCase  =		len(A__      )

       # divide all the points into an upper hull and a lower hull
       # the left most point and the right most point are definitely
       # members of the convex hull by definition.
       # use these two anchors to divide all the points into two hulls,
       # an upper hull and a lower hull.

       # all points to the left (above) the line joining the extreme points belong to the
       # upper hull
       # all points to the right (below) the line joining the extreme points below to the
       # lower hull
       # ignore all points on the line joining the extreme points since they cannot be
       # part of the convex hull

       __lowerCamelCase  =		points[0]
       __lowerCamelCase  =		points[n - 1]

       __lowerCamelCase  =		{left_most_point, right_most_point}
       __lowerCamelCase  =		[]
       __lowerCamelCase  =		[]

       for i in range(1    ,     n - 1      ):
              __lowerCamelCase  =		_det(A__    ,     A__    ,     points[i]      )

              if det > 0:
                     upper_hull.append(points[i]      )
              elif det < 0:
                     lower_hull.append(points[i]      )

       _construct_hull(A__    ,     A__    ,     A__    ,     A__      )
       _construct_hull(A__    ,     A__    ,     A__    ,     A__      )

       return sorted(A__      )




def    lowerCamelCase__			(		A__			:   list[Point]    ,     A__			:   Point    ,     A__			:   Point    ,     A__			:   set[Point]      ):





       '''simple docstring'''



       if points:
              __lowerCamelCase  =		None
              __lowerCamelCase  =		float("""-inf"""      )
              __lowerCamelCase  =		[]

              for p in points:
                     __lowerCamelCase  =		_det(A__    ,     A__    ,     A__      )

                     if det > 0:
                            candidate_points.append(A__      )

                            if det > extreme_point_distance:
                                   __lowerCamelCase  =		det
                                   __lowerCamelCase  =		p

              if extreme_point:
                     _construct_hull(A__    ,     A__    ,     A__    ,     A__      )
                     convex_set.add(A__      )
                     _construct_hull(A__    ,     A__    ,     A__    ,     A__      )




def    lowerCamelCase__			(		A__			:   list[Point]      ):





       '''simple docstring'''



       __lowerCamelCase  =		sorted(_validate_input(A__      )      )
       __lowerCamelCase  =		len(A__      )

       __lowerCamelCase  =		points[:2]
       for i in range(2    ,     A__      ):
              __lowerCamelCase  =		_det(convex_hull[1]    ,     convex_hull[0]    ,     points[i]      )
              if det > 0:
                     convex_hull.insert(0    ,     points[i]      )
                     break
              elif det < 0:
                     convex_hull.append(points[i]      )
                     break
              else:
                     __lowerCamelCase  =		points[i]
       i += 1

       for j in range(A__    ,     A__      ):
              if (
                  _det(convex_hull[0]    ,     convex_hull[-1]    ,     points[j]      ) > 0
                  and _det(convex_hull[-1]    ,     convex_hull[0]    ,     points[1]      ) < 0
              ):
                     # The point lies within the convex hull
                     continue

              convex_hull.insert(0    ,     points[j]      )
              convex_hull.append(points[j]      )
              while _det(convex_hull[0]    ,     convex_hull[1]    ,     convex_hull[2]      ) >= 0:
                     del convex_hull[1]
              while _det(convex_hull[-1]    ,     convex_hull[-2]    ,     convex_hull[-3]      ) <= 0:
                     del convex_hull[-2]

    # `convex_hull` is contains the convex hull in circular order
       return sorted(convex_hull[1:] if len(A__      ) > 3 else convex_hull      )




def    lowerCamelCase__			(		):





       '''simple docstring'''



       __lowerCamelCase  =		[
           (0, 3),
           (2, 2),
           (1, 1),
           (2, 1),
           (3, 0),
           (0, 0),
           (3, 3),
           (2, -1),
           (2, -4),
           (1, -3),
       ]
       # the convex set of points is
       # [(0, 0), (0, 3), (1, -3), (2, -4), (3, 0), (3, 3)]
       __lowerCamelCase  =		convex_hull_bf(A__      )

       __lowerCamelCase  =		convex_hull_recursive(A__      )
       assert results_bf == results_recursive

       __lowerCamelCase  =		convex_hull_melkman(A__      )
       assert results_bf == results_melkman

       print(A__      )


if __name__ == "__main__":
       main()

