





"""simple docstring"""





import copy
import inspect
import unittest

from transformers import AutoBackbone
from transformers.configuration_utils import PretrainedConfig
from transformers.testing_utils import require_timm, require_torch, torch_device
from transformers.utils.import_utils import is_torch_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor


if is_torch_available():
       import torch

       from transformers import TimmBackbone, TimmBackboneConfig

from ...test_pipeline_mixin import PipelineTesterMixin
class 		SCREAMING_SNAKE_CASE_							:
    """simple docstring"""




    def __init__(     self  :Any						,    __lowercase  :List[Any]						,    __lowercase  :Any=None						,    __lowercase  :Any=None						,    __lowercase  :Union[str, Any]=None						,    __lowercase  :List[Any]="resnet50"						,    __lowercase  :List[str]=3						,    __lowercase  :Optional[int]=32						,    __lowercase  :Optional[int]=3						,    __lowercase  :Any=True						,    __lowercase  :str=True						,    ):
           __lowerCamelCase      :      Any         =parent
           __lowerCamelCase      :      Optional[int]         =out_indices if out_indices is not None else [4]
           __lowerCamelCase      :      Union[str, Any]         =stage_names
           __lowerCamelCase      :      List[str]         =out_features
           __lowerCamelCase      :      Optional[int]         =backbone
           __lowerCamelCase      :      Dict         =batch_size
           __lowerCamelCase      :      Union[str, Any]         =image_size
           __lowerCamelCase      :      List[Any]         =num_channels
           __lowerCamelCase      :      Dict         =use_pretrained_backbone
           __lowerCamelCase      :      List[str]         =is_training




    def 					__lowercase  (     self  :Dict			):
           __lowerCamelCase      :      Tuple         =floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)
           __lowerCamelCase      :      Any         =self.get_config()

           return config, pixel_values




    def 					__lowercase  (     self  :Dict			):
           return TimmBackboneConfig(
               image_size=self.image_size						,    num_channels=self.num_channels						,    out_features=self.out_features						,    out_indices=self.out_indices						,    stage_names=self.stage_names						,    use_pretrained_backbone=self.use_pretrained_backbone						,    backbone=self.backbone						,    )




    def 					__lowercase  (     self  :Any						,    __lowercase  :Optional[Any]						,    __lowercase  :List[Any]			):
           __lowerCamelCase      :      str         =TimmBackbone(config=_A			)
           model.to(_A			)
           model.eval()
           with torch.no_grad():
                  __lowerCamelCase      :      List[str]         =model(_A			)
           self.parent.assertEqual(
               result.feature_map[-1].shape						,    (self.batch_size, model.channels[-1], 14, 14)						,    )




    def 					__lowercase  (     self  :List[Any]			):
           __lowerCamelCase      :      Optional[Any]         =self.prepare_config_and_inputs()
           __lowerCamelCase    ,						__lowerCamelCase      :      List[str]         =config_and_inputs
           __lowerCamelCase      :      Union[str, Any]         ={'''pixel_values''': pixel_values}
           return config, inputs_dict

@require_torch
@require_timm
class 		SCREAMING_SNAKE_CASE_							(    _a       ,  _a       ,  _a       ,  unittest.TestCase   ):
    """simple docstring"""
    __snake_case   :       Optional[Any]	 =		(TimmBackbone,) if is_torch_available() else ()
    __snake_case   :       Optional[int]	 =		{"""feature-extraction""": TimmBackbone} if is_torch_available() else {}
    __snake_case   :       Optional[Any]	 =		False
    __snake_case   :       Union[str, Any]	 =		False
    __snake_case   :       Tuple	 =		False
    __snake_case   :       Optional[Any]	 =		False




    def 					__lowercase  (     self  :Union[str, Any]			):
           __lowerCamelCase      :      int         =TimmBackboneModelTester(self			)
           __lowerCamelCase      :      int         =ConfigTester(self						,    config_class=_A						,    has_text_modality=_A			)




    def 					__lowercase  (     self  :Dict			):
           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 					__lowercase  (     self  :Tuple			):
           __lowerCamelCase      :      Dict         ='''resnet18'''
           __lowerCamelCase      :      Union[str, Any]         ='''microsoft/resnet-18'''

           __lowerCamelCase      :      Union[str, Any]         =AutoBackbone.from_pretrained(_A						,    use_timm_backbone=_A			)
           __lowerCamelCase      :      Any         =AutoBackbone.from_pretrained(_A			)

           self.assertEqual(len(timm_model.out_features			)						,    len(transformers_model.out_features			)			)
           self.assertEqual(len(timm_model.stage_names			)						,    len(transformers_model.stage_names			)			)
           self.assertEqual(timm_model.channels						,    transformers_model.channels			)
           # Out indices are set to the last layer by default. For timm models, we don't know
           # the number of layers in advance, so we set it to (-1,), whereas for transformers
           # models, we set it to [len(stage_names) - 1] (kept for backward compatibility).
           self.assertEqual(timm_model.out_indices						,    (-1,)			)
           self.assertEqual(transformers_model.out_indices						,    [len(timm_model.stage_names			) - 1]			)

           __lowerCamelCase      :      Optional[int]         =AutoBackbone.from_pretrained(_A						,    use_timm_backbone=_A						,    out_indices=[1, 2, 3]			)
           __lowerCamelCase      :      List[Any]         =AutoBackbone.from_pretrained(_A						,    out_indices=[1, 2, 3]			)

           self.assertEqual(timm_model.out_indices						,    transformers_model.out_indices			)
           self.assertEqual(len(timm_model.out_features			)						,    len(transformers_model.out_features			)			)
           self.assertEqual(timm_model.channels						,    transformers_model.channels			)




    @unittest.skip('''TimmBackbone doesn\'t support feed forward chunking'''			)
    def 					__lowercase  (     self  :Optional[int]			):
           pass




    @unittest.skip('''TimmBackbone doesn\'t have num_hidden_layers attribute'''			)
    def 					__lowercase  (     self  :str			):
           pass




    @unittest.skip('''TimmBackbone initialization is managed on the timm side'''			)
    def 					__lowercase  (     self  :int			):
           pass




    @unittest.skip('''TimmBackbone models doesn\'t have inputs_embeds'''			)
    def 					__lowercase  (     self  :Optional[int]			):
           pass




    @unittest.skip('''TimmBackbone models doesn\'t have inputs_embeds'''			)
    def 					__lowercase  (     self  :List[str]			):
           pass




    @unittest.skip('''TimmBackbone model cannot be created without specifying a backbone checkpoint'''			)
    def 					__lowercase  (     self  :str			):
           pass




    @unittest.skip('''Only checkpoints on timm can be loaded into TimmBackbone'''			)
    def 					__lowercase  (     self  :List[str]			):
           pass




    @unittest.skip('''model weights aren\'t tied in TimmBackbone.'''			)
    def 					__lowercase  (     self  :Optional[Any]			):
           pass




    @unittest.skip('''model weights aren\'t tied in TimmBackbone.'''			)
    def 					__lowercase  (     self  :str			):
           pass




    @unittest.skip('''Only checkpoints on timm can be loaded into TimmBackbone'''			)
    def 					__lowercase  (     self  :str			):
           pass




    @unittest.skip('''Only checkpoints on timm can be loaded into TimmBackbone'''			)
    def 					__lowercase  (     self  :Any			):
           pass




    @unittest.skip('''TimmBackbone doesn\'t have hidden size info in its configuration.'''			)
    def 					__lowercase  (     self  :Union[str, Any]			):
           pass




    @unittest.skip('''TimmBackbone doesn\'t support output_attentions.'''			)
    def 					__lowercase  (     self  :Union[str, Any]			):
           pass




    @unittest.skip('''Safetensors is not supported by timm.'''			)
    def 					__lowercase  (     self  :List[Any]			):
           pass




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




    def 					__lowercase  (     self  :int			):
           __lowerCamelCase    ,						__lowerCamelCase      :      Dict         =self.model_tester.prepare_config_and_inputs_for_common()

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

                  __lowerCamelCase      :      Union[str, Any]         =['''pixel_values''']
                  self.assertListEqual(arg_names[:1]						,    _A			)




    def 					__lowercase  (     self  :int			):
           __lowerCamelCase    ,						__lowerCamelCase      :      Any         =self.model_tester.prepare_config_and_inputs_for_common()
           __lowerCamelCase      :      int         =True
           __lowerCamelCase      :      Union[str, Any]         =self.has_attentions

           # no need to test all models as different heads yield the same functionality
           __lowerCamelCase      :      List[str]         =self.all_model_classes[0]
           __lowerCamelCase      :      int         =model_class(_A			)
           model.to(_A			)

           __lowerCamelCase      :      List[Any]         =self._prepare_for_class(_A						,    _A			)
           __lowerCamelCase      :      Union[str, Any]         =model(**_A			)
           __lowerCamelCase      :      Dict         =outputs[0][-1]

           # Encoder-/Decoder-only models
           __lowerCamelCase      :      Dict         =outputs.hidden_states[0]
           hidden_states.retain_grad()

           if self.has_attentions:
                  __lowerCamelCase      :      Union[str, Any]         =outputs.attentions[0]
                  attentions.retain_grad()

           output.flatten()[0].backward(retain_graph=_A			)

           self.assertIsNotNone(hidden_states.grad			)

           if self.has_attentions:
                  self.assertIsNotNone(attentions.grad			)




    def 					__lowercase  (     self  :Union[str, Any]			):
           __lowerCamelCase    ,						__lowerCamelCase      :      Tuple         =self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  __lowerCamelCase      :      Optional[Any]         =model_class(_A			)
                  model.to(_A			)
                  model.eval()
                  __lowerCamelCase      :      Optional[Any]         =model(**_A			)

                  self.assertEqual(len(result.feature_maps			)						,    len(config.out_indices			)			)
                  self.assertEqual(len(model.channels			)						,    len(config.out_indices			)			)

                  # Check output of last stage is taken if out_features=None, out_indices=None
                  __lowerCamelCase      :      Union[str, Any]         =copy.deepcopy(_A			)
                  __lowerCamelCase      :      Optional[int]         =None
                  __lowerCamelCase      :      int         =model_class(_A			)
                  model.to(_A			)
                  model.eval()
                  __lowerCamelCase      :      List[str]         =model(**_A			)

                  self.assertEqual(len(result.feature_maps			)						,    1			)
                  self.assertEqual(len(model.channels			)						,    1			)

                  # Check backbone can be initialized with fresh weights
                  __lowerCamelCase      :      Any         =copy.deepcopy(_A			)
                  __lowerCamelCase      :      Optional[int]         =False
                  __lowerCamelCase      :      int         =model_class(_A			)
                  model.to(_A			)
                  model.eval()
                  __lowerCamelCase      :      Tuple         =model(**_A			)











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


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

_UpperCamelCase							=      '\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 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							Any    ,						SCREAMING_SNAKE_CASE       :							Tuple    ,						SCREAMING_SNAKE_CASE       :							Tuple=8   ):


							'''simple docstring'''







							__lowerCamelCase      :      Tuple         =height // scale_factor**2
							if height % scale_factor**2 != 0:
														new_height += 1
							__lowerCamelCase      :      Tuple         =width // scale_factor**2
							if width % scale_factor**2 != 0:
														new_width += 1
							return new_height * scale_factor, new_width * scale_factor
class 		SCREAMING_SNAKE_CASE_							(    snake_case__   ):
				"""simple docstring"""




				def __init__(     self  :Union[str, Any]						,    __lowercase  :UNetaDConditionModel						,    __lowercase  :DDPMScheduler						,    __lowercase  :VQModel						,    ):
											super().__init__()

											self.register_modules(
											    unet=__lowercase						,    scheduler=__lowercase						,    movq=__lowercase						,    )
											__lowerCamelCase      :      Any         =2 ** (len(self.movq.config.block_out_channels			) - 1)




				def 					__lowercase  (     self  :int						,    __lowercase  :Tuple						,    __lowercase  :Dict						,    __lowercase  :str						,    __lowercase  :Dict						,    __lowercase  :Union[str, Any]						,    __lowercase  :Dict			):
											if latents is None:
																		__lowerCamelCase      :      str         =randn_tensor(__lowercase						,    generator=__lowercase						,    device=__lowercase						,    dtype=__lowercase			)
											else:
																		if latents.shape != shape:
																									raise ValueError(f'Unexpected latents shape, got {latents.shape}, expected {shape}'			)
																		__lowerCamelCase      :      str         =latents.to(__lowercase			)

											__lowerCamelCase      :      Any         =latents * scheduler.init_noise_sigma
											return latents




				def 					__lowercase  (     self  :Dict						,    __lowercase  :Any=0			):
											if is_accelerate_available():
																		from accelerate import cpu_offload
											else:
																		raise ImportError('''Please install accelerate via `pip install accelerate`'''			)

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

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




				def 					__lowercase  (     self  :Optional[int]						,    __lowercase  :List[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.'''			)

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

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

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

											# We'll offload the last model manually.
											__lowerCamelCase      :      Dict         =hook




				@property
				# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
				def 					__lowercase  (     self  :Tuple			):
											if not hasattr(self.unet						,    '''_hf_hook'''			):
																		return self.device
											for module in self.unet.modules():
																		if (
																		    hasattr(__lowercase						,    '''_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(__lowercase			)
				def __call__(     self  :Union[str, Any]						,    __lowercase  :Union[torch.FloatTensor, List[torch.FloatTensor]]						,    __lowercase  :Union[torch.FloatTensor, List[torch.FloatTensor]]						,    __lowercase  :torch.FloatTensor						,    __lowercase  :int = 512						,    __lowercase  :int = 512						,    __lowercase  :int = 100						,    __lowercase  :float = 4.0						,    __lowercase  :int = 1						,    __lowercase  :Optional[Union[torch.Generator, List[torch.Generator]]] = None						,    __lowercase  :Optional[torch.FloatTensor] = None						,    __lowercase  :Optional[str] = "pil"						,    __lowercase  :bool = True						,    ):
											__lowerCamelCase      :      Any         =self._execution_device

											__lowerCamelCase      :      Dict         =guidance_scale > 1.0

											if isinstance(__lowercase						,    __lowercase			):
																		__lowerCamelCase      :      Any         =torch.cat(__lowercase						,    dim=0			)
											if isinstance(__lowercase						,    __lowercase			):
																		__lowerCamelCase      :      Dict         =torch.cat(__lowercase						,    dim=0			)
											if isinstance(__lowercase						,    __lowercase			):
																		__lowerCamelCase      :      Dict         =torch.cat(__lowercase						,    dim=0			)

											__lowerCamelCase      :      List[Any]         =image_embeds.shape[0] * num_images_per_prompt

											if do_classifier_free_guidance:
																		__lowerCamelCase      :      Optional[int]         =image_embeds.repeat_interleave(__lowercase						,    dim=0			)
																		__lowerCamelCase      :      Any         =negative_image_embeds.repeat_interleave(__lowercase						,    dim=0			)
																		__lowerCamelCase      :      List[Any]         =hint.repeat_interleave(__lowercase						,    dim=0			)

											__lowerCamelCase      :      List[str]         =torch.cat([negative_image_embeds, image_embeds]						,    dim=0			).to(dtype=self.unet.dtype						,    device=__lowercase			)
											__lowerCamelCase      :      List[str]         =torch.cat([hint, hint]						,    dim=0			).to(dtype=self.unet.dtype						,    device=__lowercase			)
											self.scheduler.set_timesteps(__lowercase						,    device=__lowercase			)
											__lowerCamelCase      :      str         =self.scheduler.timesteps

											__lowerCamelCase      :      Tuple         =self.movq.config.latent_channels

											__lowerCamelCase    ,						__lowerCamelCase      :      List[str]         =downscale_height_and_width(__lowercase						,    __lowercase						,    self.movq_scale_factor			)

											# create initial latent
											__lowerCamelCase      :      Optional[int]         =self.prepare_latents(
											    (batch_size, num_channels_latents, height, width)						,    image_embeds.dtype						,    __lowercase						,    __lowercase						,    __lowercase						,    self.scheduler						,    )

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

																		__lowerCamelCase      :      List[Any]         ={'''image_embeds''': image_embeds, '''hint''': hint}
																		__lowerCamelCase      :      Union[str, Any]         =self.unet(
																		    sample=__lowercase						,    timestep=__lowercase						,    encoder_hidden_states=__lowercase						,    added_cond_kwargs=__lowercase						,    return_dict=__lowercase						,    )[0]

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

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

											if output_type == "pil":
																		__lowerCamelCase      :      Optional[Any]         =self.numpy_to_pil(__lowercase			)

											if not return_dict:
																		return (image,)

											return ImagePipelineOutput(images=__lowercase			)




