





"""simple docstring"""





from __future__ import annotations

import unittest

from transformers import CTRLConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import tensorflow as tf

    from transformers.models.ctrl.modeling_tf_ctrl import (
        TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFCTRLForSequenceClassification,
        TFCTRLLMHeadModel,
        TFCTRLModel,
    )





class       snake_case		(       __UpperCAmelCase       ):

      def __init__(      self       :Optional[int]						,							_lowerCamelCase       :Optional[int]						,							):
       __SCREAMING_SNAKE_CASE			:      Tuple 							=	parent
       __SCREAMING_SNAKE_CASE			:      str 							=	1_3
       __SCREAMING_SNAKE_CASE			:      List[str] 							=	7
       __SCREAMING_SNAKE_CASE			:      Optional[Any] 							=	True
       __SCREAMING_SNAKE_CASE			:      List[Any] 							=	True
       __SCREAMING_SNAKE_CASE			:      Optional[int] 							=	True
       __SCREAMING_SNAKE_CASE			:      Any 							=	True
       __SCREAMING_SNAKE_CASE			:      int 							=	True
       __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	9_9
       __SCREAMING_SNAKE_CASE			:      Optional[Any] 							=	3_2
       __SCREAMING_SNAKE_CASE			:      List[Any] 							=	2
       __SCREAMING_SNAKE_CASE			:      int 							=	4
       __SCREAMING_SNAKE_CASE			:      List[str] 							=	3_7
       __SCREAMING_SNAKE_CASE			:      Tuple 							=	'''gelu'''
       __SCREAMING_SNAKE_CASE			:      Tuple 							=	0.1
       __SCREAMING_SNAKE_CASE			:      List[Any] 							=	0.1
       __SCREAMING_SNAKE_CASE			:      Tuple 							=	5_1_2
       __SCREAMING_SNAKE_CASE			:      Any 							=	1_6
       __SCREAMING_SNAKE_CASE			:      Optional[int] 							=	2
       __SCREAMING_SNAKE_CASE			:      List[Any] 							=	0.0_2
       __SCREAMING_SNAKE_CASE			:      Dict 							=	3
       __SCREAMING_SNAKE_CASE			:      List[Any] 							=	4
       __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	None
       __SCREAMING_SNAKE_CASE			:      Any 							=	self.vocab_size - 1

      def 						SCREAMING_SNAKE_CASE_      (      self       :str						):
       __SCREAMING_SNAKE_CASE			:      Any 							=	ids_tensor([self.batch_size, self.seq_length]						,							self.vocab_size						)

       __SCREAMING_SNAKE_CASE			:      Dict 							=	None
       if self.use_input_mask:
        __SCREAMING_SNAKE_CASE			:      Any 							=	random_attention_mask([self.batch_size, self.seq_length]						)

       __SCREAMING_SNAKE_CASE			:      Optional[Any] 							=	None
       if self.use_token_type_ids:
        __SCREAMING_SNAKE_CASE			:      List[str] 							=	ids_tensor([self.batch_size, self.seq_length]						,							self.type_vocab_size						)

       __SCREAMING_SNAKE_CASE			:      Any 							=	None
       if self.use_mc_token_ids:
        __SCREAMING_SNAKE_CASE			:      int 							=	ids_tensor([self.batch_size, self.num_choices]						,							self.seq_length						)

       __SCREAMING_SNAKE_CASE			:      Any 							=	None
       __SCREAMING_SNAKE_CASE			:      Dict 							=	None
       __SCREAMING_SNAKE_CASE			:      Optional[int] 							=	None
       if self.use_labels:
        __SCREAMING_SNAKE_CASE			:      Tuple 							=	ids_tensor([self.batch_size]						,							self.type_sequence_label_size						)
        __SCREAMING_SNAKE_CASE			:      List[Any] 							=	ids_tensor([self.batch_size, self.seq_length]						,							self.num_labels						)
        __SCREAMING_SNAKE_CASE			:      List[Any] 							=	ids_tensor([self.batch_size]						,							self.num_choices						)

       __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	CTRLConfig(
           vocab_size=self.vocab_size						,							n_embd=self.hidden_size						,							n_layer=self.num_hidden_layers						,							n_head=self.num_attention_heads						,							n_positions=self.max_position_embeddings						,							pad_token_id=self.pad_token_id						,							)

       __SCREAMING_SNAKE_CASE			:      str 							=	ids_tensor([self.num_hidden_layers, self.num_attention_heads]						,							2						)

       return (
           config,
           input_ids,
           input_mask,
           head_mask,
           token_type_ids,
           mc_token_ids,
           sequence_labels,
           token_labels,
           choice_labels,
       )

      def 						SCREAMING_SNAKE_CASE_      (      self       :List[str]						,							_lowerCamelCase       :Dict						,							_lowerCamelCase       :Any						,							_lowerCamelCase       :int						,							_lowerCamelCase       :Any						,							_lowerCamelCase       :Optional[Any]						,							*_lowerCamelCase       :str						):
       __SCREAMING_SNAKE_CASE			:      List[str] 							=	TFCTRLModel(config=_lowerCamelCase						)
       __SCREAMING_SNAKE_CASE			:      List[Any] 							=	{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
       __SCREAMING_SNAKE_CASE			:      Optional[int] 							=	model(_lowerCamelCase						)

       __SCREAMING_SNAKE_CASE			:      int 							=	[input_ids, None, input_mask]  # None is the input for 'past'
       __SCREAMING_SNAKE_CASE			:      Optional[int] 							=	model(_lowerCamelCase						)

       __SCREAMING_SNAKE_CASE			:      List[Any] 							=	model(_lowerCamelCase						)

       self.parent.assertEqual(result.last_hidden_state.shape						,							(self.batch_size, self.seq_length, self.hidden_size)						)

      def 						SCREAMING_SNAKE_CASE_      (      self       :Dict						,							_lowerCamelCase       :List[str]						,							_lowerCamelCase       :Optional[Any]						,							_lowerCamelCase       :Union[str, Any]						,							_lowerCamelCase       :int						,							_lowerCamelCase       :Union[str, Any]						,							*_lowerCamelCase       :Union[str, Any]						):
       __SCREAMING_SNAKE_CASE			:      Any 							=	TFCTRLLMHeadModel(config=_lowerCamelCase						)
       __SCREAMING_SNAKE_CASE			:      int 							=	{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
       __SCREAMING_SNAKE_CASE			:      str 							=	model(_lowerCamelCase						)
       self.parent.assertEqual(result.logits.shape						,							(self.batch_size, self.seq_length, self.vocab_size)						)

      def 						SCREAMING_SNAKE_CASE_      (      self       :List[str]						,							_lowerCamelCase       :Tuple						,							_lowerCamelCase       :List[str]						,							_lowerCamelCase       :Optional[Any]						,							_lowerCamelCase       :List[Any]						,							_lowerCamelCase       :Optional[int]						,							*_lowerCamelCase       :Dict						):
       __SCREAMING_SNAKE_CASE			:      List[str] 							=	self.num_labels
       __SCREAMING_SNAKE_CASE			:      List[Any] 							=	ids_tensor([self.batch_size]						,							self.type_sequence_label_size						)
       __SCREAMING_SNAKE_CASE			:      Dict 							=	{
           '''input_ids''': input_ids,
           '''token_type_ids''': token_type_ids,
           '''labels''': sequence_labels,
       }
       __SCREAMING_SNAKE_CASE			:      Any 							=	TFCTRLForSequenceClassification(_lowerCamelCase						)
       __SCREAMING_SNAKE_CASE			:      Any 							=	model(_lowerCamelCase						)
       self.parent.assertEqual(result.logits.shape						,							(self.batch_size, self.num_labels)						)



      def 						SCREAMING_SNAKE_CASE_      (      self       :List[str]						):
       __SCREAMING_SNAKE_CASE			:      Optional[Any] 							=	self.prepare_config_and_inputs()

       (
           (
           __SCREAMING_SNAKE_CASE
       )							,				(
           __SCREAMING_SNAKE_CASE
       )							,				(
           __SCREAMING_SNAKE_CASE
       )							,				(
           __SCREAMING_SNAKE_CASE
       )							,				(
           __SCREAMING_SNAKE_CASE
       )							,				(
           __SCREAMING_SNAKE_CASE
       )							,				(
           __SCREAMING_SNAKE_CASE
       )							,				(
           __SCREAMING_SNAKE_CASE
       )							,				(
           __SCREAMING_SNAKE_CASE
       )							,				
       )			:      Optional[int] 							=	config_and_inputs

       __SCREAMING_SNAKE_CASE			:      Dict 							=	{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
       return config, inputs_dict







@require_tf
class       snake_case		(       __UpperCAmelCase			,					__UpperCAmelCase			,					unittest.TestCase       ):
      lowerCamelCase__      =			(TFCTRLModel, TFCTRLLMHeadModel, TFCTRLForSequenceClassification) if is_tf_available() else ()
      lowerCamelCase__      =			(TFCTRLLMHeadModel,) if is_tf_available() else ()
      lowerCamelCase__      =			(
          {
              '''feature-extraction''': TFCTRLModel,
              '''text-classification''': TFCTRLForSequenceClassification,
              '''text-generation''': TFCTRLLMHeadModel,
              '''zero-shot''': TFCTRLForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      lowerCamelCase__      =			False
      lowerCamelCase__      =			False

      def 						SCREAMING_SNAKE_CASE_      (      self       :List[Any]						,							_lowerCamelCase       :Any						,							_lowerCamelCase       :Union[str, Any]						,							_lowerCamelCase       :Optional[Any]						,							_lowerCamelCase       :Optional[Any]						,							_lowerCamelCase       :int						):
       if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
        # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
        # `CTRLConfig` was never used in pipeline tests, either because of a missing checkpoint or because a tiny
        # config could not be created.
        return True

       return False

      def 						SCREAMING_SNAKE_CASE_      (      self       :Any						):
       __SCREAMING_SNAKE_CASE			:      Optional[int] 							=	TFCTRLModelTester(self						)
       __SCREAMING_SNAKE_CASE			:      Dict 							=	ConfigTester(self						,							config_class=_lowerCamelCase						,							n_embd=3_7						)

      def 						SCREAMING_SNAKE_CASE_      (      self       :int						):
       self.config_tester.run_common_tests()

      def 						SCREAMING_SNAKE_CASE_      (      self       :str						):
       __SCREAMING_SNAKE_CASE			:      str 							=	self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_ctrl_model(*_lowerCamelCase						)

      def 						SCREAMING_SNAKE_CASE_      (      self       :Optional[Any]						):
       __SCREAMING_SNAKE_CASE			:      str 							=	self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_ctrl_lm_head(*_lowerCamelCase						)

      def 						SCREAMING_SNAKE_CASE_      (      self       :str						):
       __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_ctrl_for_sequence_classification(*_lowerCamelCase						)

      def 						SCREAMING_SNAKE_CASE_      (      self       :Optional[Any]						):
       __SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE			:      Dict 							=	self.model_tester.prepare_config_and_inputs_for_common()
       __SCREAMING_SNAKE_CASE			:      List[str] 							=	[TFCTRLLMHeadModel]
       __SCREAMING_SNAKE_CASE			:      int 							=	[TFCTRLForSequenceClassification]

       for model_class in self.all_model_classes:
        __SCREAMING_SNAKE_CASE			:      int 							=	model_class(_lowerCamelCase						)
        model.build()  # may be needed for the get_bias() call below
        assert isinstance(model.get_input_embeddings()						,							tf.keras.layers.Layer						)

        if model_class in list_lm_models:
         __SCREAMING_SNAKE_CASE			:      List[Any] 							=	model.get_output_embeddings()
         assert isinstance(_lowerCamelCase						,							tf.keras.layers.Layer						)
         __SCREAMING_SNAKE_CASE			:      Optional[Any] 							=	model.get_bias()
         assert isinstance(_lowerCamelCase						,							_lowerCamelCase						)
         for k, v in name.items():
          assert isinstance(_lowerCamelCase						,							tf.Variable						)
        elif model_class in list_other_models_with_output_ebd:
         __SCREAMING_SNAKE_CASE			:      Any 							=	model.get_output_embeddings()
         assert isinstance(_lowerCamelCase						,							tf.keras.layers.Layer						)
         __SCREAMING_SNAKE_CASE			:      List[Any] 							=	model.get_bias()
         assert name is None
        else:
         __SCREAMING_SNAKE_CASE			:      Any 							=	model.get_output_embeddings()
         assert x is None
         __SCREAMING_SNAKE_CASE			:      int 							=	model.get_bias()
         assert name is None



      @slow
      def 						SCREAMING_SNAKE_CASE_      (      self       :Tuple						):
       for model_name in TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        __SCREAMING_SNAKE_CASE			:      int 							=	TFCTRLModel.from_pretrained(_lowerCamelCase						)
        self.assertIsNotNone(_lowerCamelCase						)







@require_tf
class       snake_case		(       unittest.TestCase       ):

      @slow
      def 						SCREAMING_SNAKE_CASE_      (      self       :Optional[int]						):
       __SCREAMING_SNAKE_CASE			:      Tuple 							=	TFCTRLLMHeadModel.from_pretrained('''ctrl'''						)
       __SCREAMING_SNAKE_CASE			:      List[Any] 							=	tf.convert_to_tensor([[1_1_8_5_9, 0, 1_6_1_1, 8]]						,							dtype=tf.intaa						)  # Legal the president is
       __SCREAMING_SNAKE_CASE			:      str 							=	[
           1_1_8_5_9,
           0,
           1_6_1_1,
           8,
           5,
           1_5_0,
           2_6_4_4_9,
           2,
           1_9,
           3_4_8,
           4_6_9,
           3,
           2_5_9_5,
           4_8,
           2_0_7_4_0,
           2_4_6_5_3_3,
           2_4_6_5_3_3,
           1_9,
           3_0,
           5,
       ]  # Legal the president is a good guy and I don't want to lose my job. \n \n I have a

       __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	model.generate(_lowerCamelCase						,							do_sample=_lowerCamelCase						)
       self.assertListEqual(output_ids[0].numpy().tolist()						,							_lowerCamelCase						)








"""simple docstring"""





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,
)


_lowerCamelCase  					=     logging.get_logger(__name__)  # pylint: disable=invalid-name

_lowerCamelCase  					=     '''
    Examples:
        ```py
        >>> import torch
        >>> import numpy as np

        >>> from diffusers import KandinskyV22PriorEmb2EmbPipeline, KandinskyV22ControlnetImg2ImgPipeline
        >>> from transformers import pipeline
        >>> from diffusers.utils import load_image


        >>> def make_hint(image, depth_estimator):
        ...     image = depth_estimator(image)["depth"]
        ...     image = np.array(image)
        ...     image = image[:, :, None]
        ...     image = np.concatenate([image, image, image], axis=2)
        ...     detected_map = torch.from_numpy(image).float() / 255.0
        ...     hint = detected_map.permute(2, 0, 1)
        ...     return hint


        >>> depth_estimator = pipeline("depth-estimation")

        >>> pipe_prior = KandinskyV22PriorEmb2EmbPipeline.from_pretrained(
        ...     "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16
        ... )
        >>> pipe_prior = pipe_prior.to("cuda")

        >>> pipe = KandinskyV22ControlnetImg2ImgPipeline.from_pretrained(
        ...     "kandinsky-community/kandinsky-2-2-controlnet-depth", torch_dtype=torch.float16
        ... )
        >>> pipe = pipe.to("cuda")

        >>> img = load_image(
        ...     "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
        ...     "/kandinsky/cat.png"
        ... ).resize((768, 768))


        >>> hint = make_hint(img, depth_estimator).unsqueeze(0).half().to("cuda")

        >>> prompt = "A robot, 4k photo"
        >>> negative_prior_prompt = "lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature"

        >>> generator = torch.Generator(device="cuda").manual_seed(43)

        >>> img_emb = pipe_prior(prompt=prompt, image=img, strength=0.85, generator=generator)
        >>> negative_emb = pipe_prior(prompt=negative_prior_prompt, image=img, strength=1, generator=generator)

        >>> images = pipe(
        ...     image=img,
        ...     strength=0.5,
        ...     image_embeds=img_emb.image_embeds,
        ...     negative_image_embeds=negative_emb.image_embeds,
        ...     hint=hint,
        ...     num_inference_steps=50,
        ...     generator=generator,
        ...     height=768,
        ...     width=768,
        ... ).images

        >>> images[0].save("robot_cat.png")
        ```
'''




def 		lowerCAmelCase_     (  lowercase_       :							Optional[int]			,				lowercase_       :							Optional[Any]			,				lowercase_       :							Optional[int]=8		):

 '''simple docstring'''






 __SCREAMING_SNAKE_CASE			:      Tuple 							=	height // scale_factor**2
 if height % scale_factor**2 != 0:
  new_height += 1
 __SCREAMING_SNAKE_CASE			:      Optional[Any] 							=	width // scale_factor**2
 if width % scale_factor**2 != 0:
  new_width += 1
 return new_height * scale_factor, new_width * scale_factor




def 		lowerCAmelCase_     (  lowercase_       :							str			,				lowercase_       :							Tuple=512			,				lowercase_       :							Any=512		):

 '''simple docstring'''






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






class       snake_case		(       __UpperCAmelCase       ):

      def __init__(      self       :Dict						,							_lowerCamelCase       :UNetaDConditionModel						,							_lowerCamelCase       :DDPMScheduler						,							_lowerCamelCase       :VQModel						,							):
       super().__init__()

       self.register_modules(
           unet=_lowerCamelCase						,							scheduler=_lowerCamelCase						,							movq=_lowerCamelCase						,							)
       __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	2 ** (len(self.movq.config.block_out_channels						) - 1)

      def 						SCREAMING_SNAKE_CASE_      (      self       :Union[str, Any]						,							_lowerCamelCase       :Union[str, Any]						,							_lowerCamelCase       :List[Any]						,							_lowerCamelCase       :List[Any]						):
       # get the original timestep using init_timestep
       __SCREAMING_SNAKE_CASE			:      str 							=	min(int(num_inference_steps * strength						)						,							_lowerCamelCase						)

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

       return timesteps, num_inference_steps - t_start

      def 						SCREAMING_SNAKE_CASE_      (      self       :Any						,							_lowerCamelCase       :str						,							_lowerCamelCase       :int						,							_lowerCamelCase       :str						,							_lowerCamelCase       :Optional[int]						,							_lowerCamelCase       :Optional[Any]						,							_lowerCamelCase       :Union[str, Any]						,							_lowerCamelCase       :Dict=None						):
       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						)}'''						)

       __SCREAMING_SNAKE_CASE			:      Optional[Any] 							=	image.to(device=_lowerCamelCase						,							dtype=_lowerCamelCase						)

       __SCREAMING_SNAKE_CASE			:      List[Any] 							=	batch_size * num_images_per_prompt

       if image.shape[1] == 4:
        __SCREAMING_SNAKE_CASE			:      Optional[Any] 							=	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						):
         __SCREAMING_SNAKE_CASE			:      Optional[int] 							=	[
             self.movq.encode(image[i : i + 1]						).latent_dist.sample(generator[i]						) for i in range(_lowerCamelCase						)
         ]
         __SCREAMING_SNAKE_CASE			:      List[Any] 							=	torch.cat(_lowerCamelCase						,							dim=0						)
        else:
         __SCREAMING_SNAKE_CASE			:      Dict 							=	self.movq.encode(_lowerCamelCase						).latent_dist.sample(_lowerCamelCase						)

        __SCREAMING_SNAKE_CASE			:      Optional[Any] 							=	self.movq.config.scaling_factor * init_latents

       __SCREAMING_SNAKE_CASE			:      List[Any] 							=	torch.cat([init_latents]						,							dim=0						)

       __SCREAMING_SNAKE_CASE			:      int 							=	init_latents.shape
       __SCREAMING_SNAKE_CASE			:      int 							=	randn_tensor(_lowerCamelCase						,							generator=_lowerCamelCase						,							device=_lowerCamelCase						,							dtype=_lowerCamelCase						)

       # get latents
       __SCREAMING_SNAKE_CASE			:      Optional[Any] 							=	self.scheduler.add_noise(_lowerCamelCase						,							_lowerCamelCase						,							_lowerCamelCase						)

       __SCREAMING_SNAKE_CASE			:      List[str] 							=	init_latents

       return latents

      def 						SCREAMING_SNAKE_CASE_      (      self       :str						,							_lowerCamelCase       :List[Any]=0						):
       if is_accelerate_available():
        from accelerate import cpu_offload
       else:
        raise ImportError('''Please install accelerate via `pip install accelerate`'''						)

       __SCREAMING_SNAKE_CASE			:      List[str] 							=	torch.device(f'''cuda:{gpu_id}'''						)

       __SCREAMING_SNAKE_CASE			:      Tuple 							=	[
           self.unet,
           self.movq,
       ]
       for cpu_offloaded_model in models:
        if cpu_offloaded_model is not None:
         cpu_offload(_lowerCamelCase						,							_lowerCamelCase						)

      def 						SCREAMING_SNAKE_CASE_      (      self       :Optional[Any]						,							_lowerCamelCase       :Optional[Any]=0						):
       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.'''						)

       __SCREAMING_SNAKE_CASE			:      Dict 							=	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)

       __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	None
       for cpu_offloaded_model in [self.unet, self.movq]:
        __SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE			:      Optional[int] 							=	cpu_offload_with_hook(_lowerCamelCase						,							_lowerCamelCase						,							prev_module_hook=_lowerCamelCase						)

       # We'll offload the last model manually.
       __SCREAMING_SNAKE_CASE			:      Any 							=	hook

      @property
      # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
      def 						SCREAMING_SNAKE_CASE_      (      self       :Optional[int]						):
       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       :List[str]						,							_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       :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						,							):
       __SCREAMING_SNAKE_CASE			:      List[str] 							=	self._execution_device

       __SCREAMING_SNAKE_CASE			:      Optional[int] 							=	guidance_scale > 1.0

       if isinstance(_lowerCamelCase						,							_lowerCamelCase						):
        __SCREAMING_SNAKE_CASE			:      Optional[int] 							=	torch.cat(_lowerCamelCase						,							dim=0						)
       if isinstance(_lowerCamelCase						,							_lowerCamelCase						):
        __SCREAMING_SNAKE_CASE			:      List[str] 							=	torch.cat(_lowerCamelCase						,							dim=0						)
       if isinstance(_lowerCamelCase						,							_lowerCamelCase						):
        __SCREAMING_SNAKE_CASE			:      Dict 							=	torch.cat(_lowerCamelCase						,							dim=0						)

       __SCREAMING_SNAKE_CASE			:      Tuple 							=	image_embeds.shape[0]

       if do_classifier_free_guidance:
        __SCREAMING_SNAKE_CASE			:      Tuple 							=	image_embeds.repeat_interleave(_lowerCamelCase						,							dim=0						)
        __SCREAMING_SNAKE_CASE			:      Any 							=	negative_image_embeds.repeat_interleave(_lowerCamelCase						,							dim=0						)
        __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	hint.repeat_interleave(_lowerCamelCase						,							dim=0						)

       __SCREAMING_SNAKE_CASE			:      Tuple 							=	torch.cat([negative_image_embeds, image_embeds]						,							dim=0						).to(dtype=self.unet.dtype						,							device=_lowerCamelCase						)
       __SCREAMING_SNAKE_CASE			:      Optional[int] 							=	torch.cat([hint, hint]						,							dim=0						).to(dtype=self.unet.dtype						,							device=_lowerCamelCase						)

       if not isinstance(_lowerCamelCase						,							_lowerCamelCase						):
        __SCREAMING_SNAKE_CASE			:      Optional[Any] 							=	[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'''						)

       __SCREAMING_SNAKE_CASE			:      Optional[Any] 							=	torch.cat([prepare_image(_lowerCamelCase						,							_lowerCamelCase						,							_lowerCamelCase						) for i in image]						,							dim=0						)
       __SCREAMING_SNAKE_CASE			:      Dict 							=	image.to(dtype=image_embeds.dtype						,							device=_lowerCamelCase						)

       __SCREAMING_SNAKE_CASE			:      List[Any] 							=	self.movq.encode(_lowerCamelCase						)['''latents''']
       __SCREAMING_SNAKE_CASE			:      Tuple 							=	latents.repeat_interleave(_lowerCamelCase						,							dim=0						)
       self.scheduler.set_timesteps(_lowerCamelCase						,							device=_lowerCamelCase						)
       __SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE			:      Optional[int] 							=	self.get_timesteps(_lowerCamelCase						,							_lowerCamelCase						,							_lowerCamelCase						)
       __SCREAMING_SNAKE_CASE			:      Optional[int] 							=	timesteps[:1].repeat(batch_size * num_images_per_prompt						)
       __SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE			:      Dict 							=	downscale_height_and_width(_lowerCamelCase						,							_lowerCamelCase						,							self.movq_scale_factor						)
       __SCREAMING_SNAKE_CASE			:      Union[str, 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
        __SCREAMING_SNAKE_CASE			:      int 							=	torch.cat([latents] * 2						) if do_classifier_free_guidance else latents

        __SCREAMING_SNAKE_CASE			:      Any 							=	{'''image_embeds''': image_embeds, '''hint''': hint}
        __SCREAMING_SNAKE_CASE			:      int 							=	self.unet(
            sample=_lowerCamelCase						,							timestep=_lowerCamelCase						,							encoder_hidden_states=_lowerCamelCase						,							added_cond_kwargs=_lowerCamelCase						,							return_dict=_lowerCamelCase						,							)[0]

        if do_classifier_free_guidance:
         __SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE			:      Tuple 							=	noise_pred.split(latents.shape[1]						,							dim=1						)
         __SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE			:      Optional[int] 							=	noise_pred.chunk(2						)
         __SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE			:      Any 							=	variance_pred.chunk(2						)
         __SCREAMING_SNAKE_CASE			:      int 							=	noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
         __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	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"]
        ):
         __SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE			:      Any 							=	noise_pred.split(latents.shape[1]						,							dim=1						)

        # compute the previous noisy sample x_t -> x_t-1

        __SCREAMING_SNAKE_CASE			:      Optional[int] 							=	self.scheduler.step(
            _lowerCamelCase						,							_lowerCamelCase						,							_lowerCamelCase						,							generator=_lowerCamelCase						,							)[0]

       # post-processing
       __SCREAMING_SNAKE_CASE			:      List[str] 							=	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"]:
        __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	image * 0.5 + 0.5
        __SCREAMING_SNAKE_CASE			:      str 							=	image.clamp(0						,							1						)
        __SCREAMING_SNAKE_CASE			:      int 							=	image.cpu().permute(0						,							2						,							3						,							1						).float().numpy()

       if output_type == "pil":
        __SCREAMING_SNAKE_CASE			:      Optional[int] 							=	self.numpy_to_pil(_lowerCamelCase						)

       if not return_dict:
        return (image,)

       return ImagePipelineOutput(images=_lowerCamelCase						)

