
import inspect
import unittest

from transformers import SegformerConfig, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import (
       MODEL_MAPPING,
       SegformerForImageClassification,
       SegformerForSemanticSegmentation,
       SegformerModel,
   )
   from transformers.models.segformer.modeling_segformer import SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
   from PIL import Image

   from transformers import SegformerImageProcessor






class UpperCAmelCase   (						SCREAMING_SNAKE_CASE__   ):





   '''simple docstring'''


   def  __magic_name__				(					self     :							Any     ):

      """simple docstring"""
      _A:  Any					=       self.config_class(**self.inputs_dict     )
      self.parent.assertTrue(hasattr(lowerCAmelCase_			,			'''hidden_sizes'''     )     )
      self.parent.assertTrue(hasattr(lowerCAmelCase_			,			'''num_attention_heads'''     )     )
      self.parent.assertTrue(hasattr(lowerCAmelCase_			,			'''num_encoder_blocks'''     )     )






class UpperCAmelCase   :





   '''simple docstring'''


   def __init__(					self     :							Union[str, Any]			,			lowerCAmelCase_     :							Optional[int]			,			lowerCAmelCase_     :							List[str]=1_3			,			lowerCAmelCase_     :							Union[str, Any]=6_4			,			lowerCAmelCase_     :							List[str]=3			,			lowerCAmelCase_     :							Optional[int]=4			,			lowerCAmelCase_     :							List[Any]=[2, 2, 2, 2]			,			lowerCAmelCase_     :							Tuple=[8, 4, 2, 1]			,			lowerCAmelCase_     :							Any=[1_6, 3_2, 6_4, 1_2_8]			,			lowerCAmelCase_     :							str=[1, 4, 8, 1_6]			,			lowerCAmelCase_     :							Any=[1, 2, 4, 8]			,			lowerCAmelCase_     :							Union[str, Any]=True			,			lowerCAmelCase_     :							Union[str, Any]=True			,			lowerCAmelCase_     :							int="gelu"			,			lowerCAmelCase_     :							int=0.1			,			lowerCAmelCase_     :							int=0.1			,			lowerCAmelCase_     :							Union[str, Any]=0.02			,			lowerCAmelCase_     :							List[Any]=3			,			lowerCAmelCase_     :							str=None			,			):

      """simple docstring"""
      _A:  Union[str, Any]					=       parent
      _A:  Tuple					=       batch_size
      _A:  Union[str, Any]					=       image_size
      _A:  Union[str, Any]					=       num_channels
      _A:  int					=       num_encoder_blocks
      _A:  int					=       sr_ratios
      _A:  List[Any]					=       depths
      _A:  str					=       hidden_sizes
      _A:  int					=       downsampling_rates
      _A:  Optional[int]					=       num_attention_heads
      _A:  Dict					=       is_training
      _A:  Union[str, Any]					=       use_labels
      _A:  List[Any]					=       hidden_act
      _A:  List[str]					=       hidden_dropout_prob
      _A:  int					=       attention_probs_dropout_prob
      _A:  Union[str, Any]					=       initializer_range
      _A:  List[Any]					=       num_labels
      _A:  str					=       scope


   def  __magic_name__				(					self     :							Tuple     ):

      """simple docstring"""
      _A:  Optional[int]					=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

      _A:  List[Any]					=       None
      if self.use_labels:
         _A:  Dict					=       ids_tensor([self.batch_size, self.image_size, self.image_size]			,			self.num_labels     )

      _A:  Union[str, Any]					=       self.get_config()
      return config, pixel_values, labels


   def  __magic_name__				(					self     :							str     ):

      """simple docstring"""
      return SegformerConfig(
          image_size=self.image_size			,			num_channels=self.num_channels			,			num_encoder_blocks=self.num_encoder_blocks			,			depths=self.depths			,			hidden_sizes=self.hidden_sizes			,			num_attention_heads=self.num_attention_heads			,			hidden_act=self.hidden_act			,			hidden_dropout_prob=self.hidden_dropout_prob			,			attention_probs_dropout_prob=self.attention_probs_dropout_prob			,			initializer_range=self.initializer_range			,			)


   def  __magic_name__				(					self     :							Dict			,			lowerCAmelCase_     :							List[Any]			,			lowerCAmelCase_     :							Optional[Any]			,			lowerCAmelCase_     :							Dict     ):

      """simple docstring"""
      _A:  Optional[int]					=       SegformerModel(config=lowerCAmelCase_     )
      model.to(lowerCAmelCase_     )
      model.eval()
      _A:  Any					=       model(lowerCAmelCase_     )
      _A:  List[Any]					=       self.image_size // (self.downsampling_rates[-1] * 2)
      self.parent.assertEqual(
          result.last_hidden_state.shape			,			(self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)     )


   def  __magic_name__				(					self     :							Union[str, Any]			,			lowerCAmelCase_     :							List[Any]			,			lowerCAmelCase_     :							Dict			,			lowerCAmelCase_     :							str     ):

      """simple docstring"""
      _A:  int					=       self.num_labels
      _A:  Union[str, Any]					=       SegformerForSemanticSegmentation(lowerCAmelCase_     )
      model.to(lowerCAmelCase_     )
      model.eval()
      _A:  Union[str, Any]					=       model(lowerCAmelCase_     )
      self.parent.assertEqual(
          result.logits.shape			,			(self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)     )
      _A:  Tuple					=       model(lowerCAmelCase_			,			labels=lowerCAmelCase_     )
      self.parent.assertEqual(
          result.logits.shape			,			(self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)     )
      self.parent.assertGreater(result.loss			,			0.0     )


   def  __magic_name__				(					self     :							Union[str, Any]			,			lowerCAmelCase_     :							Union[str, Any]			,			lowerCAmelCase_     :							Any			,			lowerCAmelCase_     :							Dict     ):

      """simple docstring"""
      _A:  Dict					=       1
      _A:  List[Any]					=       SegformerForSemanticSegmentation(config=lowerCAmelCase_     )
      model.to(lowerCAmelCase_     )
      model.eval()
      _A:  Optional[Any]					=       torch.randint(0			,			1			,			(self.batch_size, self.image_size, self.image_size)     ).to(lowerCAmelCase_     )
      _A:  str					=       model(lowerCAmelCase_			,			labels=lowerCAmelCase_     )
      self.parent.assertGreater(result.loss			,			0.0     )


   def  __magic_name__				(					self     :							Dict     ):

      """simple docstring"""
      _A:  int					=       self.prepare_config_and_inputs()
      _A			,  _A			,  _A:  Dict					=       config_and_inputs
      _A:  List[Any]					=       {'''pixel_values''': pixel_values}
      return config, inputs_dict






@require_torch
class UpperCAmelCase   (						SCREAMING_SNAKE_CASE__							,							SCREAMING_SNAKE_CASE__							,							unittest.TestCase   ):





   '''simple docstring'''
   __UpperCamelCase		:    Union[str, Any] 				=    (
       (
           SegformerModel,
           SegformerForSemanticSegmentation,
           SegformerForImageClassification,
       )
       if is_torch_available()
       else ()
   )
   __UpperCamelCase		:    Union[str, Any] 				=    (
       {
           '''feature-extraction''': SegformerModel,
           '''image-classification''': SegformerForImageClassification,
           '''image-segmentation''': SegformerForSemanticSegmentation,
       }
       if is_torch_available()
       else {}
   )

   __UpperCamelCase		:    Optional[int] 				=    True
   __UpperCamelCase		:    List[Any] 				=    False
   __UpperCamelCase		:    List[Any] 				=    False
   __UpperCamelCase		:    Optional[Any] 				=    False


   def  __magic_name__				(					self     :							List[str]     ):

      """simple docstring"""
      _A:  str					=       SegformerModelTester(self     )
      _A:  int					=       SegformerConfigTester(self			,			config_class=lowerCAmelCase_     )


   def  __magic_name__				(					self     :							List[Any]     ):

      """simple docstring"""
      self.config_tester.run_common_tests()


   def  __magic_name__				(					self     :							List[str]     ):

      """simple docstring"""
      _A:  Union[str, Any]					=       self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*lowerCAmelCase_     )


   def  __magic_name__				(					self     :							List[Any]     ):

      """simple docstring"""
      _A:  Union[str, Any]					=       self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_binary_image_segmentation(*lowerCAmelCase_     )


   def  __magic_name__				(					self     :							Optional[Any]     ):

      """simple docstring"""
      _A:  Optional[int]					=       self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_segmentation(*lowerCAmelCase_     )


   @unittest.skip('''SegFormer does not use inputs_embeds'''     )
   def  __magic_name__				(					self     :							Any     ):

      """simple docstring"""
      pass


   @unittest.skip('''SegFormer does not have get_input_embeddings method and get_output_embeddings methods'''     )
   def  __magic_name__				(					self     :							int     ):

      """simple docstring"""
      pass


   def  __magic_name__				(					self     :							List[Any]     ):

      """simple docstring"""
      _A			,  _A:  Optional[Any]					=       self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
         _A:  str					=       model_class(lowerCAmelCase_     )
         _A:  Dict					=       inspect.signature(model.forward     )
         # signature.parameters is an OrderedDict => so arg_names order is deterministic
         _A:  str					=       [*signature.parameters.keys()]

         _A:  int					=       ['''pixel_values''']
         self.assertListEqual(arg_names[:1]			,			lowerCAmelCase_     )


   def  __magic_name__				(					self     :							Optional[Any]     ):

      """simple docstring"""
      _A			,  _A:  List[Any]					=       self.model_tester.prepare_config_and_inputs_for_common()
      _A:  Optional[int]					=       True

      for model_class in self.all_model_classes:
         _A:  Union[str, Any]					=       True
         _A:  Optional[int]					=       False
         _A:  Any					=       True
         _A:  int					=       model_class(lowerCAmelCase_     )
         model.to(lowerCAmelCase_     )
         model.eval()
         with torch.no_grad():
            _A:  Dict					=       model(**self._prepare_for_class(lowerCAmelCase_			,			lowerCAmelCase_     )     )
         _A:  Any					=       outputs.attentions

         _A:  Tuple					=       sum(self.model_tester.depths     )
         self.assertEqual(len(lowerCAmelCase_     )			,			lowerCAmelCase_     )

         # check that output_attentions also work using config
         del inputs_dict["output_attentions"]
         _A:  List[Any]					=       True
         _A:  List[Any]					=       model_class(lowerCAmelCase_     )
         model.to(lowerCAmelCase_     )
         model.eval()
         with torch.no_grad():
            _A:  Optional[int]					=       model(**self._prepare_for_class(lowerCAmelCase_			,			lowerCAmelCase_     )     )
         _A:  Dict					=       outputs.attentions

         self.assertEqual(len(lowerCAmelCase_     )			,			lowerCAmelCase_     )

         # verify the first attentions (first block, first layer)
         _A:  Union[str, Any]					=       (self.model_tester.image_size // 4) ** 2
         _A:  List[str]					=       (self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
         self.assertListEqual(
             list(attentions[0].shape[-3:]     )			,			[self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]			,			)

         # verify the last attentions (last block, last layer)
         _A:  int					=       (self.model_tester.image_size // 3_2) ** 2
         _A:  Union[str, Any]					=       (self.model_tester.image_size // (3_2 * self.model_tester.sr_ratios[-1])) ** 2
         self.assertListEqual(
             list(attentions[-1].shape[-3:]     )			,			[self.model_tester.num_attention_heads[-1], expected_seq_len, expected_reduced_seq_len]			,			)
         _A:  Tuple					=       len(lowerCAmelCase_     )

         # Check attention is always last and order is fine
         _A:  Union[str, Any]					=       True
         _A:  Tuple					=       True
         _A:  Union[str, Any]					=       model_class(lowerCAmelCase_     )
         model.to(lowerCAmelCase_     )
         model.eval()
         with torch.no_grad():
            _A:  Optional[int]					=       model(**self._prepare_for_class(lowerCAmelCase_			,			lowerCAmelCase_     )     )

         self.assertEqual(out_len + 1			,			len(lowerCAmelCase_     )     )

         _A:  List[Any]					=       outputs.attentions

         self.assertEqual(len(lowerCAmelCase_     )			,			lowerCAmelCase_     )
         # verify the first attentions (first block, first layer)
         _A:  str					=       (self.model_tester.image_size // 4) ** 2
         _A:  List[Any]					=       (self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
         self.assertListEqual(
             list(self_attentions[0].shape[-3:]     )			,			[self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]			,			)


   def  __magic_name__				(					self     :							int     ):

      """simple docstring"""
      def check_hidden_states_output(lowerCAmelCase_     :							Optional[int]			,			lowerCAmelCase_     :							List[str]			,			lowerCAmelCase_     :							List[Any]     ):
         _A:  Dict					=       model_class(lowerCAmelCase_     )
         model.to(lowerCAmelCase_     )
         model.eval()

         with torch.no_grad():
            _A:  str					=       model(**self._prepare_for_class(lowerCAmelCase_			,			lowerCAmelCase_     )     )

         _A:  List[Any]					=       outputs.hidden_states

         _A:  Optional[Any]					=       self.model_tester.num_encoder_blocks
         self.assertEqual(len(lowerCAmelCase_     )			,			lowerCAmelCase_     )

         # verify the first hidden states (first block)
         self.assertListEqual(
             list(hidden_states[0].shape[-3:]     )			,			[
                 self.model_tester.hidden_sizes[0],
                 self.model_tester.image_size // 4,
                 self.model_tester.image_size // 4,
             ]			,			)

      _A			,  _A:  Union[str, Any]					=       self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
         _A:  Any					=       True
         check_hidden_states_output(lowerCAmelCase_			,			lowerCAmelCase_			,			lowerCAmelCase_     )

         # check that output_hidden_states also work using config
         del inputs_dict["output_hidden_states"]
         _A:  Union[str, Any]					=       True

         check_hidden_states_output(lowerCAmelCase_			,			lowerCAmelCase_			,			lowerCAmelCase_     )


   def  __magic_name__				(					self     :							Optional[Any]     ):

      """simple docstring"""
      if not self.model_tester.is_training:
         return

      _A			,  _A:  int					=       self.model_tester.prepare_config_and_inputs_for_common()
      _A:  str					=       True

      for model_class in self.all_model_classes:
         if model_class in get_values(lowerCAmelCase_     ):
            continue

         _A:  List[str]					=       model_class(lowerCAmelCase_     )
         model.to(lowerCAmelCase_     )
         model.train()
         _A:  Optional[Any]					=       self._prepare_for_class(lowerCAmelCase_			,			lowerCAmelCase_			,			return_labels=lowerCAmelCase_     )
         _A:  Any					=       model(**lowerCAmelCase_     ).loss
         loss.backward()


   @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''     )
   def  __magic_name__				(					self     :							int     ):

      """simple docstring"""
      pass


   @slow
   def  __magic_name__				(					self     :							Dict     ):

      """simple docstring"""
      for model_name in SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         _A:  Any					=       SegformerModel.from_pretrained(lowerCAmelCase_     )
         self.assertIsNotNone(lowerCAmelCase_     )






def   lowerCamelCase__							(					)      ->   Dict:
   _A:  List[str]					=       Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''      )
   return image






@require_torch
class UpperCAmelCase   (						unittest.TestCase   ):





   '''simple docstring'''


   @slow
   def  __magic_name__				(					self     :							List[str]     ):

      """simple docstring"""
      # only resize + normalize
      _A:  List[Any]					=       SegformerImageProcessor(
          image_scale=(5_1_2, 5_1_2)			,			keep_ratio=lowerCAmelCase_			,			align=lowerCAmelCase_			,			do_random_crop=lowerCAmelCase_     )
      _A:  Tuple					=       SegformerForSemanticSegmentation.from_pretrained('''nvidia/segformer-b0-finetuned-ade-512-512'''     ).to(
          lowerCAmelCase_     )

      _A:  Tuple					=       prepare_img()
      _A:  List[Any]					=       image_processor(images=lowerCAmelCase_			,			return_tensors='''pt'''     )
      _A:  Any					=       encoded_inputs.pixel_values.to(lowerCAmelCase_     )

      with torch.no_grad():
         _A:  str					=       model(lowerCAmelCase_     )

      _A:  Optional[Any]					=       torch.Size((1, model.config.num_labels, 1_2_8, 1_2_8)     )
      self.assertEqual(outputs.logits.shape			,			lowerCAmelCase_     )

      _A:  str					=       torch.tensor(
          [
              [[-4.6310, -5.5232, -6.2356], [-5.1921, -6.1444, -6.5996], [-5.4424, -6.2790, -6.7574]],
              [[-12.1391, -13.3122, -13.9554], [-12.8732, -13.9352, -14.3563], [-12.9438, -13.8226, -14.2513]],
              [[-12.5134, -13.4686, -14.4915], [-12.8669, -14.4343, -14.7758], [-13.2523, -14.5819, -15.0694]],
          ]     ).to(lowerCAmelCase_     )
      self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]			,			lowerCAmelCase_			,			atol=1e-4     )     )


   @slow
   def  __magic_name__				(					self     :							int     ):

      """simple docstring"""
      # only resize + normalize
      _A:  List[str]					=       SegformerImageProcessor(
          image_scale=(5_1_2, 5_1_2)			,			keep_ratio=lowerCAmelCase_			,			align=lowerCAmelCase_			,			do_random_crop=lowerCAmelCase_     )
      _A:  Any					=       SegformerForSemanticSegmentation.from_pretrained(
          '''nvidia/segformer-b1-finetuned-cityscapes-1024-1024'''     ).to(lowerCAmelCase_     )

      _A:  Tuple					=       prepare_img()
      _A:  Optional[int]					=       image_processor(images=lowerCAmelCase_			,			return_tensors='''pt'''     )
      _A:  Union[str, Any]					=       encoded_inputs.pixel_values.to(lowerCAmelCase_     )

      with torch.no_grad():
         _A:  List[Any]					=       model(lowerCAmelCase_     )

      _A:  Union[str, Any]					=       torch.Size((1, model.config.num_labels, 1_2_8, 1_2_8)     )
      self.assertEqual(outputs.logits.shape			,			lowerCAmelCase_     )

      _A:  Any					=       torch.tensor(
          [
              [[-13.5748, -13.9111, -12.6500], [-14.3500, -15.3683, -14.2328], [-14.7532, -16.0424, -15.6087]],
              [[-17.1651, -15.8725, -12.9653], [-17.2580, -17.3718, -14.8223], [-16.6058, -16.8783, -16.7452]],
              [[-3.6456, -3.0209, -1.4203], [-3.0797, -3.1959, -2.0000], [-1.8757, -1.9217, -1.6997]],
          ]     ).to(lowerCAmelCase_     )
      self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]			,			lowerCAmelCase_			,			atol=1e-1     )     )


   @slow
   def  __magic_name__				(					self     :							List[Any]     ):

      """simple docstring"""
      # only resize + normalize
      _A:  List[str]					=       SegformerImageProcessor(
          image_scale=(5_1_2, 5_1_2)			,			keep_ratio=lowerCAmelCase_			,			align=lowerCAmelCase_			,			do_random_crop=lowerCAmelCase_     )
      _A:  Optional[Any]					=       SegformerForSemanticSegmentation.from_pretrained('''nvidia/segformer-b0-finetuned-ade-512-512'''     ).to(
          lowerCAmelCase_     )

      _A:  Dict					=       prepare_img()
      _A:  Dict					=       image_processor(images=lowerCAmelCase_			,			return_tensors='''pt'''     )
      _A:  Optional[int]					=       encoded_inputs.pixel_values.to(lowerCAmelCase_     )

      with torch.no_grad():
         _A:  List[Any]					=       model(lowerCAmelCase_     )

      _A:  Optional[Any]					=       outputs.logits.detach().cpu()

      _A:  Any					=       image_processor.post_process_semantic_segmentation(outputs=lowerCAmelCase_			,			target_sizes=[(5_0_0, 3_0_0)]     )
      _A:  Any					=       torch.Size((5_0_0, 3_0_0)     )
      self.assertEqual(segmentation[0].shape			,			lowerCAmelCase_     )

      _A:  str					=       image_processor.post_process_semantic_segmentation(outputs=lowerCAmelCase_     )
      _A:  str					=       torch.Size((1_2_8, 1_2_8)     )
      self.assertEqual(segmentation[0].shape			,			lowerCAmelCase_     )





from typing import List, Optional, Union

import numpy as np
import PIL
import torch
from PIL import Image

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


UpperCAmelCase__		:     Optional[int]     			=      logging.get_logger(__name__)  # pylint: disable=invalid-name

UpperCAmelCase__		:     Dict     			=      '\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyV22Img2ImgPipeline, KandinskyV22PriorPipeline\n        >>> from diffusers.utils import load_image\n        >>> import torch\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(\n        ...     "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16\n        ... )\n        >>> pipe_prior.to("cuda")\n\n        >>> prompt = "A red cartoon frog, 4k"\n        >>> image_emb, zero_image_emb = pipe_prior(prompt, return_dict=False)\n\n        >>> pipe = KandinskyV22Img2ImgPipeline.from_pretrained(\n        ...     "kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16\n        ... )\n        >>> pipe.to("cuda")\n\n        >>> init_image = load_image(\n        ...     "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"\n        ...     "/kandinsky/frog.png"\n        ... )\n\n        >>> image = pipe(\n        ...     image=init_image,\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=100,\n        ...     strength=0.2,\n        ... ).images\n\n        >>> image[0].save("red_frog.png")\n        ```\n'






def   lowerCamelCase__							(					a		,		a		,		a=8      )      ->   List[Any]:
   _A:  int					=       height // scale_factor**2
   if height % scale_factor**2 != 0:
      new_height += 1
   _A:  str					=       width // scale_factor**2
   if width % scale_factor**2 != 0:
      new_width += 1
   return new_height * scale_factor, new_width * scale_factor






def   lowerCamelCase__							(					a		,		a=5_12		,		a=5_12      )      ->   Dict:
   _A:  Union[str, Any]					=       pil_image.resize((w, h)		,		resample=Image.BICUBIC		,		reducing_gap=1      )
   _A:  Tuple					=       np.array(pil_image.convert('''RGB'''      )      )
   _A:  List[str]					=       arr.astype(np.floataa      ) / 127.5 - 1
   _A:  Tuple					=       np.transpose(a		,		[2, 0, 1]      )
   _A:  Any					=       torch.from_numpy(a      ).unsqueeze(0      )
   return image






class UpperCAmelCase   (						SCREAMING_SNAKE_CASE__   ):





   '''simple docstring'''


   def __init__(					self     :							int			,			lowerCAmelCase_     :							UNetaDConditionModel			,			lowerCAmelCase_     :							DDPMScheduler			,			lowerCAmelCase_     :							VQModel			,			):

      """simple docstring"""
      super().__init__()

      self.register_modules(
          unet=lowerCAmelCase_			,			scheduler=lowerCAmelCase_			,			movq=lowerCAmelCase_			,			)
      _A:  List[Any]					=       2 ** (len(self.movq.config.block_out_channels     ) - 1)


   def  __magic_name__				(					self     :							List[str]			,			lowerCAmelCase_     :							Union[str, Any]			,			lowerCAmelCase_     :							List[Any]			,			lowerCAmelCase_     :							Union[str, Any]     ):

      """simple docstring"""
      # get the original timestep using init_timestep
      _A:  Union[str, Any]					=       min(int(num_inference_steps * strength     )			,			lowerCAmelCase_     )

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

      return timesteps, num_inference_steps - t_start


   def  __magic_name__				(					self     :							List[str]			,			lowerCAmelCase_     :							List[str]			,			lowerCAmelCase_     :							str			,			lowerCAmelCase_     :							List[Any]			,			lowerCAmelCase_     :							List[Any]			,			lowerCAmelCase_     :							Tuple			,			lowerCAmelCase_     :							Optional[int]			,			lowerCAmelCase_     :							Optional[int]=None     ):

      """simple docstring"""
      if not isinstance(lowerCAmelCase_			,			(torch.Tensor, PIL.Image.Image, list)     ):
         raise ValueError(
             F"""`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(lowerCAmelCase_     )}"""     )

      _A:  Optional[int]					=       image.to(device=lowerCAmelCase_			,			dtype=lowerCAmelCase_     )

      _A:  Union[str, Any]					=       batch_size * num_images_per_prompt

      if image.shape[1] == 4:
         _A:  Optional[int]					=       image

      else:
         if isinstance(lowerCAmelCase_			,			lowerCAmelCase_     ) and len(lowerCAmelCase_     ) != batch_size:
            raise ValueError(
                F"""You have passed a list of generators of length {len(lowerCAmelCase_     )}, but requested an effective batch"""
                F""" size of {batch_size}. Make sure the batch size matches the length of the generators."""     )

         elif isinstance(lowerCAmelCase_			,			lowerCAmelCase_     ):
            _A:  List[Any]					=       [
                self.movq.encode(image[i : i + 1]     ).latent_dist.sample(generator[i]     ) for i in range(lowerCAmelCase_     )
            ]
            _A:  Optional[Any]					=       torch.cat(lowerCAmelCase_			,			dim=0     )
         else:
            _A:  Optional[int]					=       self.movq.encode(lowerCAmelCase_     ).latent_dist.sample(lowerCAmelCase_     )

         _A:  int					=       self.movq.config.scaling_factor * init_latents

      _A:  Optional[Any]					=       torch.cat([init_latents]			,			dim=0     )

      _A:  Any					=       init_latents.shape
      _A:  Optional[Any]					=       randn_tensor(lowerCAmelCase_			,			generator=lowerCAmelCase_			,			device=lowerCAmelCase_			,			dtype=lowerCAmelCase_     )

      # get latents
      _A:  Union[str, Any]					=       self.scheduler.add_noise(lowerCAmelCase_			,			lowerCAmelCase_			,			lowerCAmelCase_     )

      _A:  List[str]					=       init_latents

      return latents


   def  __magic_name__				(					self     :							Optional[int]			,			lowerCAmelCase_     :							Optional[int]=0     ):

      """simple docstring"""
      if is_accelerate_available():
         from accelerate import cpu_offload
      else:
         raise ImportError('''Please install accelerate via `pip install accelerate`'''     )

      _A:  Any					=       torch.device(F"""cuda:{gpu_id}"""     )

      _A:  int					=       [
          self.unet,
          self.movq,
      ]
      for cpu_offloaded_model in models:
         if cpu_offloaded_model is not None:
            cpu_offload(lowerCAmelCase_			,			lowerCAmelCase_     )


   def  __magic_name__				(					self     :							Any			,			lowerCAmelCase_     :							Any=0     ):

      """simple docstring"""
      if is_accelerate_available() and is_accelerate_version('''>='''			,			'''0.17.0.dev0'''     ):
         from accelerate import cpu_offload_with_hook
      else:
         raise ImportError('''`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'''     )

      _A:  Any					=       torch.device(F"""cuda:{gpu_id}"""     )

      if self.device.type != "cpu":
         self.to('''cpu'''			,			silence_dtype_warnings=lowerCAmelCase_     )
         torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

      _A:  int					=       None
      for cpu_offloaded_model in [self.unet, self.movq]:
         _A			,  _A:  List[Any]					=       cpu_offload_with_hook(lowerCAmelCase_			,			lowerCAmelCase_			,			prev_module_hook=lowerCAmelCase_     )

      # We'll offload the last model manually.
      _A:  Tuple					=       hook


   @property
   # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
   def  __magic_name__				(					self     :							List[Any]     ):

      """simple docstring"""
      if not hasattr(self.unet			,			'''_hf_hook'''     ):
         return self.device
      for module in self.unet.modules():
         if (
             hasattr(lowerCAmelCase_			,			'''_hf_hook'''     )
             and hasattr(module._hf_hook			,			'''execution_device'''     )
             and module._hf_hook.execution_device is not None
         ):
            return torch.device(module._hf_hook.execution_device     )
      return self.device


   @torch.no_grad()
   @replace_example_docstring(lowerCAmelCase_     )
   def __call__(					self     :							Optional[Any]			,			lowerCAmelCase_     :							Union[torch.FloatTensor, List[torch.FloatTensor]]			,			lowerCAmelCase_     :							Union[torch.FloatTensor, PIL.Image.Image, List[torch.FloatTensor], List[PIL.Image.Image]]			,			lowerCAmelCase_     :							Union[torch.FloatTensor, List[torch.FloatTensor]]			,			lowerCAmelCase_     :							int = 5_1_2			,			lowerCAmelCase_     :							int = 5_1_2			,			lowerCAmelCase_     :							int = 1_0_0			,			lowerCAmelCase_     :							float = 4.0			,			lowerCAmelCase_     :							float = 0.3			,			lowerCAmelCase_     :							int = 1			,			lowerCAmelCase_     :							Optional[Union[torch.Generator, List[torch.Generator]]] = None			,			lowerCAmelCase_     :							Optional[str] = "pil"			,			lowerCAmelCase_     :							bool = True			,			):

      """simple docstring"""
      _A:  Any					=       self._execution_device

      _A:  Any					=       guidance_scale > 1.0

      if isinstance(lowerCAmelCase_			,			lowerCAmelCase_     ):
         _A:  Any					=       torch.cat(lowerCAmelCase_			,			dim=0     )
      _A:  int					=       image_embeds.shape[0]
      if isinstance(lowerCAmelCase_			,			lowerCAmelCase_     ):
         _A:  Dict					=       torch.cat(lowerCAmelCase_			,			dim=0     )

      if do_classifier_free_guidance:
         _A:  Any					=       image_embeds.repeat_interleave(lowerCAmelCase_			,			dim=0     )
         _A:  str					=       negative_image_embeds.repeat_interleave(lowerCAmelCase_			,			dim=0     )

      _A:  Dict					=       torch.cat([negative_image_embeds, image_embeds]			,			dim=0     ).to(dtype=self.unet.dtype			,			device=lowerCAmelCase_     )

      if not isinstance(lowerCAmelCase_			,			lowerCAmelCase_     ):
         _A:  List[str]					=       [image]
      if not all(isinstance(lowerCAmelCase_			,			(PIL.Image.Image, torch.Tensor)     ) for i in image     ):
         raise ValueError(
             F"""Input is in incorrect format: {[type(lowerCAmelCase_     ) for i in image]}. Currently, we only support  PIL image and pytorch tensor"""     )

      _A:  List[str]					=       torch.cat([prepare_image(lowerCAmelCase_			,			lowerCAmelCase_			,			lowerCAmelCase_     ) for i in image]			,			dim=0     )
      _A:  Tuple					=       image.to(dtype=image_embeds.dtype			,			device=lowerCAmelCase_     )

      _A:  Optional[Any]					=       self.movq.encode(lowerCAmelCase_     )['''latents''']
      _A:  Optional[int]					=       latents.repeat_interleave(lowerCAmelCase_			,			dim=0     )
      self.scheduler.set_timesteps(lowerCAmelCase_			,			device=lowerCAmelCase_     )
      _A			,  _A:  List[Any]					=       self.get_timesteps(lowerCAmelCase_			,			lowerCAmelCase_			,			lowerCAmelCase_     )
      _A:  Dict					=       timesteps[:1].repeat(batch_size * num_images_per_prompt     )
      _A			,  _A:  Optional[int]					=       downscale_height_and_width(lowerCAmelCase_			,			lowerCAmelCase_			,			self.movq_scale_factor     )
      _A:  Any					=       self.prepare_latents(
          lowerCAmelCase_			,			lowerCAmelCase_			,			lowerCAmelCase_			,			lowerCAmelCase_			,			image_embeds.dtype			,			lowerCAmelCase_			,			lowerCAmelCase_     )
      for i, t in enumerate(self.progress_bar(lowerCAmelCase_     )     ):
         # expand the latents if we are doing classifier free guidance
         _A:  Dict					=       torch.cat([latents] * 2     ) if do_classifier_free_guidance else latents

         _A:  str					=       {'''image_embeds''': image_embeds}
         _A:  Optional[int]					=       self.unet(
             sample=lowerCAmelCase_			,			timestep=lowerCAmelCase_			,			encoder_hidden_states=lowerCAmelCase_			,			added_cond_kwargs=lowerCAmelCase_			,			return_dict=lowerCAmelCase_			,			)[0]

         if do_classifier_free_guidance:
            _A			,  _A:  str					=       noise_pred.split(latents.shape[1]			,			dim=1     )
            _A			,  _A:  int					=       noise_pred.chunk(2     )
            _A			,  _A:  int					=       variance_pred.chunk(2     )
            _A:  Dict					=       noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
            _A:  List[str]					=       torch.cat([noise_pred, variance_pred_text]			,			dim=1     )

         if not (
             hasattr(self.scheduler.config			,			'''variance_type'''     )
             and self.scheduler.config.variance_type in ["learned", "learned_range"]
         ):
            _A			,  _A:  Optional[Any]					=       noise_pred.split(latents.shape[1]			,			dim=1     )

         # compute the previous noisy sample x_t -> x_t-1
         _A:  Any					=       self.scheduler.step(
             lowerCAmelCase_			,			lowerCAmelCase_			,			lowerCAmelCase_			,			generator=lowerCAmelCase_			,			)[0]

      # post-processing
      _A:  Tuple					=       self.movq.decode(lowerCAmelCase_			,			force_not_quantize=lowerCAmelCase_     )['''sample''']

      if output_type not in ["pt", "np", "pil"]:
         raise ValueError(F"""Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}"""     )

      if output_type in ["np", "pil"]:
         _A:  int					=       image * 0.5 + 0.5
         _A:  Any					=       image.clamp(0			,			1     )
         _A:  Any					=       image.cpu().permute(0			,			2			,			3			,			1     ).float().numpy()

      if output_type == "pil":
         _A:  Union[str, Any]					=       self.numpy_to_pil(lowerCAmelCase_     )

      if not return_dict:
         return (image,)

      return ImagePipelineOutput(images=lowerCAmelCase_     )



