




"""simple docstring"""




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


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

lowerCAmelCase							:			str     							=							"""
    Examples:
        ```py
        >>> import torch
        >>> import numpy as np

        >>> from diffusers import KandinskyV22PriorPipeline, KandinskyV22ControlnetPipeline
        >>> 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 = KandinskyV22PriorPipeline.from_pretrained(
        ...     \"kandinsky-community/kandinsky-2-2-prior\", torch_dtype=torch.float16
        ... )
        >>> pipe_prior = pipe_prior.to(\"cuda\")

        >>> pipe = KandinskyV22ControlnetPipeline.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)

        >>> image_emb, zero_image_emb = pipe_prior(
        ...     prompt=prompt, negative_prompt=negative_prior_prompt, generator=generator
        ... ).to_tuple()

        >>> images = pipe(
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=zero_image_emb,
        ...     hint=hint,
        ...     num_inference_steps=50,
        ...     generator=generator,
        ...     height=768,
        ...     width=768,
        ... ).images

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



def     a__ (				snake_case__       , snake_case__       , snake_case__=8				)						->					Union[str, Any]:
			lowerCamelCase			      =					height // scale_factor**2
			if height % scale_factor**2 != 0:
						new_height += 1
			lowerCamelCase			      =					width // scale_factor**2
			if width % scale_factor**2 != 0:
						new_width += 1
			return new_height * scale_factor, new_width * scale_factor






class     __magic_name__							(       UpperCAmelCase__				):






				'''simple docstring'''





				def __init__(       self		, _a		, _a		, _a		, ):





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

							self.register_modules(
							    unet=_a		, scheduler=_a		, movq=_a		, )
							lowerCamelCase			      =					2 ** (len(self.movq.config.block_out_channels ) - 1)





				def    _lowerCAmelCase				(       self		, _a		, _a		, _a		, _a		, _a		, _a ):





							"""simple docstring"""
							if latents is None:
										lowerCamelCase			      =					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}' )
										lowerCamelCase			      =					latents.to(_a )

							lowerCamelCase			      =					latents * scheduler.init_noise_sigma
							return latents





				def    _lowerCAmelCase				(       self		, _a=0 ):





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

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

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





				def    _lowerCAmelCase				(       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.""" )

							lowerCamelCase			      =					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)

							lowerCamelCase			      =					None
							for cpu_offloaded_model in [self.unet, self.movq]:
										lowerCamelCase      ,   lowerCamelCase			      =					cpu_offload_with_hook(_a		, _a		, prev_module_hook=_a )

							# We'll offload the last model manually.
							lowerCamelCase			      =					hook





				@property
				# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
				def    _lowerCAmelCase				(       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"""
							lowerCamelCase			      =					self._execution_device

							lowerCamelCase			      =					guidance_scale > 1.0

							if isinstance(_a		, _a ):
										lowerCamelCase			      =					torch.cat(_a		, dim=0 )
							if isinstance(_a		, _a ):
										lowerCamelCase			      =					torch.cat(_a		, dim=0 )
							if isinstance(_a		, _a ):
										lowerCamelCase			      =					torch.cat(_a		, dim=0 )

							lowerCamelCase			      =					image_embeds.shape[0] * num_images_per_prompt

							if do_classifier_free_guidance:
										lowerCamelCase			      =					image_embeds.repeat_interleave(_a		, dim=0 )
										lowerCamelCase			      =					negative_image_embeds.repeat_interleave(_a		, dim=0 )
										lowerCamelCase			      =					hint.repeat_interleave(_a		, dim=0 )

							lowerCamelCase			      =					torch.cat([negative_image_embeds, image_embeds]		, dim=0 ).to(dtype=self.unet.dtype		, device=_a )
							lowerCamelCase			      =					torch.cat([hint, hint]		, dim=0 ).to(dtype=self.unet.dtype		, device=_a )
							self.scheduler.set_timesteps(_a		, device=_a )
							lowerCamelCase			      =					self.scheduler.timesteps

							lowerCamelCase			      =					self.movq.config.latent_channels

							lowerCamelCase      ,   lowerCamelCase			      =					downscale_height_and_width(_a		, _a		, self.movq_scale_factor )

							# create initial latent
							lowerCamelCase			      =					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
										lowerCamelCase			      =					torch.cat([latents] * 2 ) if do_classifier_free_guidance else latents

										lowerCamelCase			      =					{"""image_embeds""": image_embeds, """hint""": hint}
										lowerCamelCase			      =					self.unet(
										    sample=_a		, timestep=_a		, encoder_hidden_states=_a		, added_cond_kwargs=_a		, return_dict=_a		, )[0]

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

										# compute the previous noisy sample x_t -> x_t-1
										lowerCamelCase			      =					self.scheduler.step(
										    _a		, _a		, _a		, generator=_a		, )[0]
							# post-processing
							lowerCamelCase			      =					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"]:
										lowerCamelCase			      =					image * 0.5 + 0.5
										lowerCamelCase			      =					image.clamp(0		, 1 )
										lowerCamelCase			      =					image.cpu().permute(0		, 2		, 3		, 1 ).float().numpy()

							if output_type == "pil":
										lowerCamelCase			      =					self.numpy_to_pil(_a )

							if not return_dict:
										return (image,)

							return ImagePipelineOutput(images=_a )









"""simple docstring"""




import inspect
import unittest

from transformers import ViTMSNConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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 torch import nn

					from transformers import ViTMSNForImageClassification, ViTMSNModel
					from transformers.models.vit_msn.modeling_vit_msn import VIT_MSN_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
					from PIL import Image

					from transformers import ViTImageProcessor






class     __magic_name__							:






				'''simple docstring'''





				def __init__(       self		, _a		, _a=13		, _a=30		, _a=2		, _a=3		, _a=True		, _a=True		, _a=32		, _a=5		, _a=4		, _a=37		, _a="gelu"		, _a=0.1		, _a=0.1		, _a=10		, _a=0.02		, _a=None		, ):





							"""simple docstring"""
							lowerCamelCase			      =					parent
							lowerCamelCase			      =					batch_size
							lowerCamelCase			      =					image_size
							lowerCamelCase			      =					patch_size
							lowerCamelCase			      =					num_channels
							lowerCamelCase			      =					is_training
							lowerCamelCase			      =					use_labels
							lowerCamelCase			      =					hidden_size
							lowerCamelCase			      =					num_hidden_layers
							lowerCamelCase			      =					num_attention_heads
							lowerCamelCase			      =					intermediate_size
							lowerCamelCase			      =					hidden_act
							lowerCamelCase			      =					hidden_dropout_prob
							lowerCamelCase			      =					attention_probs_dropout_prob
							lowerCamelCase			      =					type_sequence_label_size
							lowerCamelCase			      =					initializer_range
							lowerCamelCase			      =					scope

							# in ViT MSN, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
							lowerCamelCase			      =					(image_size // patch_size) ** 2
							lowerCamelCase			      =					num_patches + 1





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

							lowerCamelCase			      =					None
							if self.use_labels:
										lowerCamelCase			      =					ids_tensor([self.batch_size]		, self.type_sequence_label_size )

							lowerCamelCase			      =					self.get_config()

							return config, pixel_values, labels





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							return ViTMSNConfig(
							    image_size=self.image_size		, patch_size=self.patch_size		, num_channels=self.num_channels		, hidden_size=self.hidden_size		, num_hidden_layers=self.num_hidden_layers		, num_attention_heads=self.num_attention_heads		, intermediate_size=self.intermediate_size		, 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    _lowerCAmelCase				(       self		, _a		, _a		, _a ):





							"""simple docstring"""
							lowerCamelCase			      =					ViTMSNModel(config=_a )
							model.to(_a )
							model.eval()
							lowerCamelCase			      =					model(_a )
							self.parent.assertEqual(result.last_hidden_state.shape		, (self.batch_size, self.seq_length, self.hidden_size) )





				def    _lowerCAmelCase				(       self		, _a		, _a		, _a ):





							"""simple docstring"""
							lowerCamelCase			      =					self.type_sequence_label_size
							lowerCamelCase			      =					ViTMSNForImageClassification(_a )
							model.to(_a )
							model.eval()
							lowerCamelCase			      =					model(_a		, labels=_a )
							print("""Pixel and labels shape: {pixel_values.shape}, {labels.shape}""" )
							print("""Labels: {labels}""" )
							self.parent.assertEqual(result.logits.shape		, (self.batch_size, self.type_sequence_label_size) )

							# test greyscale images
							lowerCamelCase			      =					1
							lowerCamelCase			      =					ViTMSNForImageClassification(_a )
							model.to(_a )
							model.eval()

							lowerCamelCase			      =					floats_tensor([self.batch_size, 1, self.image_size, self.image_size] )
							lowerCamelCase			      =					model(_a )
							self.parent.assertEqual(result.logits.shape		, (self.batch_size, self.type_sequence_label_size) )





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					self.prepare_config_and_inputs()
							lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase			      =					config_and_inputs
							lowerCamelCase			      =					{"""pixel_values""": pixel_values}
							return config, inputs_dict






@require_torch
class     __magic_name__							(       UpperCAmelCase__						,					UpperCAmelCase__						,					unittest.TestCase				):






				'''simple docstring'''
				__UpperCamelCase			  =						(ViTMSNModel, ViTMSNForImageClassification) if is_torch_available() else ()
				__UpperCamelCase			  =						(
				    {"feature-extraction": ViTMSNModel, "image-classification": ViTMSNForImageClassification}
				    if is_torch_available()
				    else {}
				)

				__UpperCamelCase			  =						False
				__UpperCamelCase			  =						False
				__UpperCamelCase			  =						False
				__UpperCamelCase			  =						False





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					ViTMSNModelTester(self )
							lowerCamelCase			      =					ConfigTester(self		, config_class=_a		, has_text_modality=_a		, hidden_size=37 )





				def    _lowerCAmelCase				(       self ):





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





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





							"""simple docstring"""
							pass





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase      ,   lowerCamelCase			      =					self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
										lowerCamelCase			      =					model_class(_a )
										self.assertIsInstance(model.get_input_embeddings()		, (nn.Module) )
										lowerCamelCase			      =					model.get_output_embeddings()
										self.assertTrue(x is None or isinstance(_a		, nn.Linear ) )





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase      ,   lowerCamelCase			      =					self.model_tester.prepare_config_and_inputs_for_common()

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

										lowerCamelCase			      =					["""pixel_values"""]
										self.assertListEqual(arg_names[:1]		, _a )





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*_a )





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*_a )





				@slow
				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							for model_name in VIT_MSN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowerCamelCase			      =					ViTMSNModel.from_pretrained(_a )
										self.assertIsNotNone(_a )




def     a__ (				)						->					Any:
			lowerCamelCase			      =					Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""				)
			return image






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






				'''simple docstring'''





				@cached_property
				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							return ViTImageProcessor.from_pretrained("""facebook/vit-msn-small""" ) if is_vision_available() else None





				@slow
				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							torch.manual_seed(2 )
							lowerCamelCase			      =					ViTMSNForImageClassification.from_pretrained("""facebook/vit-msn-small""" ).to(_a )

							lowerCamelCase			      =					self.default_image_processor
							lowerCamelCase			      =					prepare_img()
							lowerCamelCase			      =					image_processor(images=_a		, return_tensors="""pt""" ).to(_a )

							# forward pass
							with torch.no_grad():
										lowerCamelCase			      =					model(**_a )

							# verify the logits
							lowerCamelCase			      =					torch.Size((1, 1_000) )
							self.assertEqual(outputs.logits.shape		, _a )

							lowerCamelCase			      =					torch.tensor([-0.0_803, -0.4_454, -0.2_375] ).to(_a )

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



