import inspect
import unittest
import warnings
from math import ceil, floor

from transformers import LevitConfig
from transformers.file_utils import cached_property, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, 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_FOR_IMAGE_CLASSIFICATION_MAPPING,
					    MODEL_MAPPING,
					    LevitForImageClassification,
					    LevitForImageClassificationWithTeacher,
					    LevitModel,
					)
					from transformers.models.levit.modeling_levit import LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
					from PIL import Image

					from transformers import LevitImageProcessor

class 			a					(		__lowerCAmelCase			):


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						int							=      self.config_class(**self.inputs_dict		)
										self.parent.assertTrue(hasattr(lowerCamelCase__  ,   "hidden_sizes"		)		)
										self.parent.assertTrue(hasattr(lowerCamelCase__  ,   "num_attention_heads"		)		)







class 			a					:


						def __init__(    self  ,   A_  ,   A_=13  ,   A_=64  ,   A_=3  ,   A_=3  ,   A_=2  ,   A_=1  ,   A_=16  ,   A_=[128, 256, 384]  ,   A_=[4, 6, 8]  ,   A_=[2, 3, 4]  ,   A_=[16, 16, 16]  ,   A_=0  ,   A_=[2, 2, 2]  ,   A_=[2, 2, 2]  ,   A_=0.02  ,   A_=True  ,   A_=True  ,   A_=2  ,   ):







										'''simple docstring'''


										_UpperCAmelCase    :						Union[str, Any]							=      parent
										_UpperCAmelCase    :						Optional[Any]							=      batch_size
										_UpperCAmelCase    :						List[Any]							=      image_size
										_UpperCAmelCase    :						List[str]							=      num_channels
										_UpperCAmelCase    :						Tuple							=      kernel_size
										_UpperCAmelCase    :						str							=      stride
										_UpperCAmelCase    :						Optional[Any]							=      padding
										_UpperCAmelCase    :						List[str]							=      hidden_sizes
										_UpperCAmelCase    :						Union[str, Any]							=      num_attention_heads
										_UpperCAmelCase    :						Any							=      depths
										_UpperCAmelCase    :						int							=      key_dim
										_UpperCAmelCase    :						Union[str, Any]							=      drop_path_rate
										_UpperCAmelCase    :						Optional[Any]							=      patch_size
										_UpperCAmelCase    :						Optional[int]							=      attention_ratio
										_UpperCAmelCase    :						Any							=      mlp_ratio
										_UpperCAmelCase    :						Dict							=      initializer_range
										_UpperCAmelCase    :						Tuple							=      [
										    ['''Subsample''', key_dim[0], hidden_sizes[0] // key_dim[0], 4, 2, 2],
										    ['''Subsample''', key_dim[0], hidden_sizes[1] // key_dim[0], 4, 2, 2],
										]
										_UpperCAmelCase    :						Dict							=      is_training
										_UpperCAmelCase    :						Optional[Any]							=      use_labels
										_UpperCAmelCase    :						List[str]							=      num_labels
										_UpperCAmelCase    :						Tuple							=      initializer_range


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						Dict							=      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

										_UpperCAmelCase    :						Any							=      None
										if self.use_labels:
														_UpperCAmelCase    :						str							=      ids_tensor([self.batch_size]  ,   self.num_labels		)

										_UpperCAmelCase    :						List[str]							=      self.get_config()
										return config, pixel_values, labels


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										return LevitConfig(
										    image_size=self.image_size  ,   num_channels=self.num_channels  ,   kernel_size=self.kernel_size  ,   stride=self.stride  ,   padding=self.padding  ,   patch_size=self.patch_size  ,   hidden_sizes=self.hidden_sizes  ,   num_attention_heads=self.num_attention_heads  ,   depths=self.depths  ,   key_dim=self.key_dim  ,   drop_path_rate=self.drop_path_rate  ,   mlp_ratio=self.mlp_ratio  ,   attention_ratio=self.attention_ratio  ,   initializer_range=self.initializer_range  ,   down_ops=self.down_ops  ,   )


						def     _UpperCAmelCase			(    self  ,   A_  ,   A_  ,   A_		):







										'''simple docstring'''


										_UpperCAmelCase    :						Union[str, Any]							=      LevitModel(config=lowerCamelCase__		)
										model.to(lowerCamelCase__		)
										model.eval()
										_UpperCAmelCase    :						Any							=      model(lowerCamelCase__		)
										_UpperCAmelCase    :						int							=      (self.image_size, self.image_size)
										_UpperCAmelCase    :						Dict							=      image_size[0], image_size[1]
										for _ in range(4		):
														_UpperCAmelCase    :						Optional[int]							=      floor(((height + 2 * self.padding - self.kernel_size) / self.stride) + 1		)
														_UpperCAmelCase    :						Optional[int]							=      floor(((width + 2 * self.padding - self.kernel_size) / self.stride) + 1		)
										self.parent.assertEqual(
										    result.last_hidden_state.shape  ,   (self.batch_size, ceil(height / 4		) * ceil(width / 4		), self.hidden_sizes[-1])  ,   )


						def     _UpperCAmelCase			(    self  ,   A_  ,   A_  ,   A_		):







										'''simple docstring'''


										_UpperCAmelCase    :						List[Any]							=      self.num_labels
										_UpperCAmelCase    :						Union[str, Any]							=      LevitForImageClassification(lowerCamelCase__		)
										model.to(lowerCamelCase__		)
										model.eval()
										_UpperCAmelCase    :						List[Any]							=      model(lowerCamelCase__  ,   labels=lowerCamelCase__		)
										self.parent.assertEqual(result.logits.shape  ,   (self.batch_size, self.num_labels)		)





						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						Tuple							=      self.prepare_config_and_inputs()
										_UpperCAmelCase    :						List[Any]							=      config_and_inputs
										_UpperCAmelCase    :						int							=      {'''pixel_values''': pixel_values}
										return config, inputs_dict







@require_torch
class 			a					(		__lowerCAmelCase							,	__lowerCAmelCase							,	unittest.TestCase			):
						_lowercase        =     (
						    (LevitModel, LevitForImageClassification, LevitForImageClassificationWithTeacher)
						    if is_torch_available()
						    else ()
						)
						_lowercase        =     (
						    {
						        "feature-extraction": LevitModel,
						        "image-classification": (LevitForImageClassification, LevitForImageClassificationWithTeacher),
						    }
						    if is_torch_available()
						    else {}
						)

						_lowercase        =     False
						_lowercase        =     False
						_lowercase        =     False
						_lowercase        =     False
						_lowercase        =     False


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						Dict							=      LevitModelTester(self		)
										_UpperCAmelCase    :						Union[str, Any]							=      ConfigTester(self  ,   config_class=lowerCamelCase__  ,   has_text_modality=lowerCamelCase__  ,   hidden_size=37		)


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										self.create_and_test_config_common_properties()
										self.config_tester.create_and_test_config_to_json_string()
										self.config_tester.create_and_test_config_to_json_file()
										self.config_tester.create_and_test_config_from_and_save_pretrained()
										self.config_tester.create_and_test_config_with_num_labels()
										self.config_tester.check_config_can_be_init_without_params()
										self.config_tester.check_config_arguments_init()


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										return


						@unittest.skip(reason="Levit does not use inputs_embeds"		)
						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										pass


						@unittest.skip(reason="Levit does not support input and output embeddings"		)
						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										pass


						@unittest.skip(reason="Levit does not output attentions"		)
						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										pass


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						str							=      self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
														_UpperCAmelCase    :						Optional[int]							=      model_class(lowerCamelCase__		)
														_UpperCAmelCase    :						Any							=      inspect.signature(model.forward		)
														# signature.parameters is an OrderedDict => so arg_names order is deterministic
														_UpperCAmelCase    :						Optional[int]							=      [*signature.parameters.keys()]

														_UpperCAmelCase    :						Dict							=      ['''pixel_values''']
														self.assertListEqual(arg_names[:1]  ,   lowerCamelCase__		)


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										def check_hidden_states_output(A_  ,   A_  ,   A_		):
														_UpperCAmelCase    :						Dict							=      model_class(lowerCamelCase__		)
														model.to(lowerCamelCase__		)
														model.eval()

														with torch.no_grad():
																		_UpperCAmelCase    :						Union[str, Any]							=      model(**self._prepare_for_class(lowerCamelCase__  ,   lowerCamelCase__		)		)

														_UpperCAmelCase    :						List[str]							=      outputs.hidden_states

														_UpperCAmelCase    :						Optional[int]							=      len(self.model_tester.depths		) + 1
														self.assertEqual(len(lowerCamelCase__		)  ,   lowerCamelCase__		)

														_UpperCAmelCase    :						Any							=      (self.model_tester.image_size, self.model_tester.image_size)
														_UpperCAmelCase    :						Optional[int]							=      image_size[0], image_size[1]
														for _ in range(4		):
																		_UpperCAmelCase    :						str							=      floor(
																		    (
																		        (height + 2 * self.model_tester.padding - self.model_tester.kernel_size)
																		        / self.model_tester.stride
																		    )
																		    + 1		)
																		_UpperCAmelCase    :						str							=      floor(
																		    (
																		        (width + 2 * self.model_tester.padding - self.model_tester.kernel_size)
																		        / self.model_tester.stride
																		    )
																		    + 1		)
														# verify the first hidden states (first block)
														self.assertListEqual(
														    list(hidden_states[0].shape[-2:]		)  ,   [
														        height * width,
														        self.model_tester.hidden_sizes[0],
														    ]  ,   )

										_UpperCAmelCase    :						Any							=      self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
														_UpperCAmelCase    :						str							=      True
														check_hidden_states_output(lowerCamelCase__  ,   lowerCamelCase__  ,   lowerCamelCase__		)

														# check that output_hidden_states also work using config
														del inputs_dict["output_hidden_states"]
														_UpperCAmelCase    :						int							=      True

														check_hidden_states_output(lowerCamelCase__  ,   lowerCamelCase__  ,   lowerCamelCase__		)


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







										'''simple docstring'''


										pass


						def     _UpperCAmelCase			(    self  ,   A_  ,   A_  ,   A_=False		):







										'''simple docstring'''


										_UpperCAmelCase    :						Tuple							=      super()._prepare_for_class(lowerCamelCase__  ,   lowerCamelCase__  ,   return_labels=lowerCamelCase__		)

										if return_labels:
														if model_class.__name__ == "LevitForImageClassificationWithTeacher":
																		del inputs_dict["labels"]

										return inputs_dict


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						Dict							=      self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*lowerCamelCase__		)


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						List[Any]							=      self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*lowerCamelCase__		)


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										if not self.model_tester.is_training:
														return

										_UpperCAmelCase    :						str							=      self.model_tester.prepare_config_and_inputs_for_common()
										_UpperCAmelCase    :						int							=      True

										for model_class in self.all_model_classes:
														# LevitForImageClassificationWithTeacher supports inference-only
														if (
														    model_class in get_values(lowerCamelCase__		)
														    or model_class.__name__ == "LevitForImageClassificationWithTeacher"
														):
																		continue
														_UpperCAmelCase    :						Optional[int]							=      model_class(lowerCamelCase__		)
														model.to(lowerCamelCase__		)
														model.train()
														_UpperCAmelCase    :						Dict							=      self._prepare_for_class(lowerCamelCase__  ,   lowerCamelCase__  ,   return_labels=lowerCamelCase__		)
														_UpperCAmelCase    :						Dict							=      model(**lowerCamelCase__		).loss
														loss.backward()


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						List[Any]							=      self.model_tester.prepare_config_and_inputs_for_common()
										if not self.model_tester.is_training:
														return

										_UpperCAmelCase    :						Tuple							=      False
										_UpperCAmelCase    :						List[Any]							=      True

										for model_class in self.all_model_classes:
														if model_class in get_values(lowerCamelCase__		) or not model_class.supports_gradient_checkpointing:
																		continue
														# LevitForImageClassificationWithTeacher supports inference-only
														if model_class.__name__ == "LevitForImageClassificationWithTeacher":
																		continue
														_UpperCAmelCase    :						Any							=      model_class(lowerCamelCase__		)
														model.gradient_checkpointing_enable()
														model.to(lowerCamelCase__		)
														model.train()
														_UpperCAmelCase    :						Optional[int]							=      self._prepare_for_class(lowerCamelCase__  ,   lowerCamelCase__  ,   return_labels=lowerCamelCase__		)
														_UpperCAmelCase    :						List[Any]							=      model(**lowerCamelCase__		).loss
														loss.backward()


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						Dict							=      self.model_tester.prepare_config_and_inputs_for_common()

										_UpperCAmelCase    :						str							=      [
										    {'''title''': '''multi_label_classification''', '''num_labels''': 2, '''dtype''': torch.float},
										    {'''title''': '''single_label_classification''', '''num_labels''': 1, '''dtype''': torch.long},
										    {'''title''': '''regression''', '''num_labels''': 1, '''dtype''': torch.float},
										]

										for model_class in self.all_model_classes:
														if (
														    model_class
														    not in [
														        *get_values(lowerCamelCase__		),
														    ]
														    or model_class.__name__ == "LevitForImageClassificationWithTeacher"
														):
																		continue

														for problem_type in problem_types:
																		with self.subTest(msg=f'Testing {model_class} with {problem_type["title"]}'		):
																						_UpperCAmelCase    :						Optional[Any]							=      problem_type['''title''']
																						_UpperCAmelCase    :						List[Any]							=      problem_type['''num_labels''']

																						_UpperCAmelCase    :						List[Any]							=      model_class(lowerCamelCase__		)
																						model.to(lowerCamelCase__		)
																						model.train()

																						_UpperCAmelCase    :						Any							=      self._prepare_for_class(lowerCamelCase__  ,   lowerCamelCase__  ,   return_labels=lowerCamelCase__		)

																						if problem_type["num_labels"] > 1:
																										_UpperCAmelCase    :						Dict							=      inputs['''labels'''].unsqueeze(1		).repeat(1  ,   problem_type["num_labels"]		)

																						_UpperCAmelCase    :						Dict							=      inputs['''labels'''].to(problem_type["dtype"]		)

																						# This tests that we do not trigger the warning form PyTorch "Using a target size that is different
																						# to the input size. This will likely lead to incorrect results due to broadcasting. Please ensure
																						# they have the same size." which is a symptom something in wrong for the regression problem.
																						# See https://github.com/huggingface/transformers/issues/11780
																						with warnings.catch_warnings(record=lowerCamelCase__		) as warning_list:
																										_UpperCAmelCase    :						Union[str, Any]							=      model(**lowerCamelCase__		).loss
																						for w in warning_list:
																										if "Using a target size that is different to the input size" in str(w.message		):
																														raise ValueError(
																														    f'Something is going wrong in the regression problem: intercepted {w.message}'		)

																						loss.backward()





						@slow
						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										for model_name in LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														_UpperCAmelCase    :						str							=      LevitModel.from_pretrained(lowerCamelCase__		)
														self.assertIsNotNone(lowerCamelCase__		)







def 							__SCREAMING_SNAKE_CASE							(			)				->    int:
				_UpperCAmelCase    :						Optional[Any]							=      Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"	)
				return image





@require_torch
@require_vision
class 			a					(		unittest.TestCase			):


						@cached_property
						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										return LevitImageProcessor.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]		)





						@slow
						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						Dict							=      LevitForImageClassificationWithTeacher.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]		).to(
										    lowerCamelCase__		)

										_UpperCAmelCase    :						Union[str, Any]							=      self.default_image_processor
										_UpperCAmelCase    :						Tuple							=      prepare_img()
										_UpperCAmelCase    :						List[str]							=      image_processor(images=lowerCamelCase__  ,   return_tensors="pt"		).to(lowerCamelCase__		)

										# forward pass
										with torch.no_grad():
														_UpperCAmelCase    :						List[Any]							=      model(**lowerCamelCase__		)

										# verify the logits
										_UpperCAmelCase    :						List[str]							=      torch.Size((1, 1000)		)
										self.assertEqual(outputs.logits.shape  ,   lowerCamelCase__		)

										_UpperCAmelCase    :						List[str]							=      torch.tensor([1.04_48, -0.37_45, -1.83_17]		).to(lowerCamelCase__		)

										self.assertTrue(torch.allclose(outputs.logits[0, :3]  ,   lowerCamelCase__  ,   atol=1e-4		)		)





from typing import List, Optional, Union

import torch

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


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

SCREAMING_SNAKE_CASE_            = '\n    Examples:\n        ```py\n        >>> import torch\n        >>> import numpy as np\n\n        >>> from diffusers import KandinskyV22PriorPipeline, KandinskyV22ControlnetPipeline\n        >>> from transformers import pipeline\n        >>> from diffusers.utils import load_image\n\n\n        >>> def make_hint(image, depth_estimator):\n        ...     image = depth_estimator(image)["depth"]\n        ...     image = np.array(image)\n        ...     image = image[:, :, None]\n        ...     image = np.concatenate([image, image, image], axis=2)\n        ...     detected_map = torch.from_numpy(image).float() / 255.0\n        ...     hint = detected_map.permute(2, 0, 1)\n        ...     return hint\n\n\n        >>> depth_estimator = pipeline("depth-estimation")\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(\n        ...     "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16\n        ... )\n        >>> pipe_prior = pipe_prior.to("cuda")\n\n        >>> pipe = KandinskyV22ControlnetPipeline.from_pretrained(\n        ...     "kandinsky-community/kandinsky-2-2-controlnet-depth", torch_dtype=torch.float16\n        ... )\n        >>> pipe = pipe.to("cuda")\n\n\n        >>> img = load_image(\n        ...     "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"\n        ...     "/kandinsky/cat.png"\n        ... ).resize((768, 768))\n\n        >>> hint = make_hint(img, depth_estimator).unsqueeze(0).half().to("cuda")\n\n        >>> prompt = "A robot, 4k photo"\n        >>> 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"\n\n        >>> generator = torch.Generator(device="cuda").manual_seed(43)\n\n        >>> image_emb, zero_image_emb = pipe_prior(\n        ...     prompt=prompt, negative_prompt=negative_prior_prompt, generator=generator\n        ... ).to_tuple()\n\n        >>> images = pipe(\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     hint=hint,\n        ...     num_inference_steps=50,\n        ...     generator=generator,\n        ...     height=768,\n        ...     width=768,\n        ... ).images\n\n        >>> images[0].save("robot_cat.png")\n        ```\n'

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





class 			a					(		UpperCAmelCase			):


						def __init__(    self  ,   A_  ,   A_  ,   A_  ,   ):







										'''simple docstring'''


										super().__init__()

										self.register_modules(
										    unet=A_  ,   scheduler=A_  ,   movq=A_  ,   )
										_UpperCAmelCase    :						Optional[int]							=      2 ** (len(self.movq.config.block_out_channels		) - 1)


						def     _UpperCAmelCase			(    self  ,   A_  ,   A_  ,   A_  ,   A_  ,   A_  ,   A_		):







										'''simple docstring'''


										if latents is None:
														_UpperCAmelCase    :						Any							=      randn_tensor(A_  ,   generator=A_  ,   device=A_  ,   dtype=A_		)
										else:
														if latents.shape != shape:
																		raise ValueError(f'Unexpected latents shape, got {latents.shape}, expected {shape}'		)
														_UpperCAmelCase    :						Optional[int]							=      latents.to(A_		)

										_UpperCAmelCase    :						Tuple							=      latents * scheduler.init_noise_sigma
										return latents


						def     _UpperCAmelCase			(    self  ,   A_=0		):







										'''simple docstring'''


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

										_UpperCAmelCase    :						Union[str, Any]							=      torch.device(f'cuda:{gpu_id}'		)

										_UpperCAmelCase    :						Any							=      [
										    self.unet,
										    self.movq,
										]
										for cpu_offloaded_model in models:
														if cpu_offloaded_model is not None:
																		cpu_offload(A_  ,   A_		)


						def     _UpperCAmelCase			(    self  ,   A_=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."		)

										_UpperCAmelCase    :						str							=      torch.device(f'cuda:{gpu_id}'		)

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

										_UpperCAmelCase    :						Optional[int]							=      None
										for cpu_offloaded_model in [self.unet, self.movq]:
														_UpperCAmelCase						,							_UpperCAmelCase    :						Dict							=      cpu_offload_with_hook(A_  ,   A_  ,   prev_module_hook=A_		)

										# We'll offload the last model manually.
										_UpperCAmelCase    :						Optional[Any]							=      hook


						@property
						# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										if not hasattr(self.unet  ,   "_hf_hook"		):
														return self.device
										for module in self.unet.modules():
														if (
														    hasattr(A_  ,   "_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(A_		)
						def __call__(    self  ,   A_  ,   A_  ,   A_  ,   A_ = 512  ,   A_ = 512  ,   A_ = 100  ,   A_ = 4.0  ,   A_ = 1  ,   A_ = None  ,   A_ = None  ,   A_ = "pil"  ,   A_ = True  ,   ):







										'''simple docstring'''


										_UpperCAmelCase    :						str							=      self._execution_device

										_UpperCAmelCase    :						Tuple							=      guidance_scale > 1.0

										if isinstance(A_  ,   A_		):
														_UpperCAmelCase    :						Union[str, Any]							=      torch.cat(A_  ,   dim=0		)
										if isinstance(A_  ,   A_		):
														_UpperCAmelCase    :						Dict							=      torch.cat(A_  ,   dim=0		)
										if isinstance(A_  ,   A_		):
														_UpperCAmelCase    :						Any							=      torch.cat(A_  ,   dim=0		)

										_UpperCAmelCase    :						Optional[Any]							=      image_embeds.shape[0] * num_images_per_prompt

										if do_classifier_free_guidance:
														_UpperCAmelCase    :						Optional[int]							=      image_embeds.repeat_interleave(A_  ,   dim=0		)
														_UpperCAmelCase    :						Union[str, Any]							=      negative_image_embeds.repeat_interleave(A_  ,   dim=0		)
														_UpperCAmelCase    :						Tuple							=      hint.repeat_interleave(A_  ,   dim=0		)

										_UpperCAmelCase    :						List[Any]							=      torch.cat([negative_image_embeds, image_embeds]  ,   dim=0		).to(dtype=self.unet.dtype  ,   device=A_		)
										_UpperCAmelCase    :						List[Any]							=      torch.cat([hint, hint]  ,   dim=0		).to(dtype=self.unet.dtype  ,   device=A_		)
										self.scheduler.set_timesteps(A_  ,   device=A_		)
										_UpperCAmelCase    :						Dict							=      self.scheduler.timesteps

										_UpperCAmelCase    :						Union[str, Any]							=      self.movq.config.latent_channels

										_UpperCAmelCase						,							_UpperCAmelCase    :						Optional[int]							=      downscale_height_and_width(A_  ,   A_  ,   self.movq_scale_factor		)

										# create initial latent
										_UpperCAmelCase    :						Tuple							=      self.prepare_latents(
										    (batch_size, num_channels_latents, height, width)  ,   image_embeds.dtype  ,   A_  ,   A_  ,   A_  ,   self.scheduler  ,   )

										for i, t in enumerate(self.progress_bar(A_		)		):
														# expand the latents if we are doing classifier free guidance
														_UpperCAmelCase    :						List[str]							=      torch.cat([latents] * 2		) if do_classifier_free_guidance else latents

														_UpperCAmelCase    :						Optional[Any]							=      {"image_embeds": image_embeds, "hint": hint}
														_UpperCAmelCase    :						Optional[int]							=      self.unet(
														    sample=A_  ,   timestep=A_  ,   encoder_hidden_states=A_  ,   added_cond_kwargs=A_  ,   return_dict=A_  ,   )[0]

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

														# compute the previous noisy sample x_t -> x_t-1
														_UpperCAmelCase    :						int							=      self.scheduler.step(
														    A_  ,   A_  ,   A_  ,   generator=A_  ,   )[0]
										# post-processing
										_UpperCAmelCase    :						Optional[Any]							=      self.movq.decode(A_  ,   force_not_quantize=A_		)["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"]:
														_UpperCAmelCase    :						Union[str, Any]							=      image * 0.5 + 0.5
														_UpperCAmelCase    :						Dict							=      image.clamp(0  ,   1		)
														_UpperCAmelCase    :						List[str]							=      image.cpu().permute(0  ,   2  ,   3  ,   1		).float().numpy()

										if output_type == "pil":
														_UpperCAmelCase    :						str							=      self.numpy_to_pil(A_		)

										if not return_dict:
														return (image,)

										return ImagePipelineOutput(images=A_		)


