



from typing import List, Optional, Union

import torch
from transformers import (
    XLMRobertaTokenizer,
)

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


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

SCREAMING_SNAKE_CASE__   			=   '''
    Examples:
        ```py
        >>> from diffusers import KandinskyPipeline, KandinskyPriorPipeline
        >>> import torch

        >>> pipe_prior = KandinskyPriorPipeline.from_pretrained("kandinsky-community/Kandinsky-2-1-prior")
        >>> pipe_prior.to("cuda")

        >>> prompt = "red cat, 4k photo"
        >>> out = pipe_prior(prompt)
        >>> image_emb = out.image_embeds
        >>> negative_image_emb = out.negative_image_embeds

        >>> pipe = KandinskyPipeline.from_pretrained("kandinsky-community/kandinsky-2-1")
        >>> pipe.to("cuda")

        >>> image = pipe(
        ...     prompt,
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=negative_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=100,
        ... ).images

        >>> image[0].save("cat.png")
        ```
'''


def A    (							__UpperCamelCase			,						__UpperCamelCase			,						__UpperCamelCase=8       )    ->   Tuple:
			A__	  =     h // scale_factor**2
			if h % scale_factor**2 != 0:
						new_h += 1
			A__	  =     w // scale_factor**2
			if w % scale_factor**2 != 0:
						new_w += 1
			return new_h * scale_factor, new_w * scale_factor






class   __lowerCAmelCase  (		UpperCAmelCase_   ):
					"""simple docstring"""






					def __init__(  self       :							Tuple					,       _snake_case       :							MultilingualCLIP					,       _snake_case       :							XLMRobertaTokenizer					,       _snake_case       :							UNetaDConditionModel					,       _snake_case       :							Union[DDIMScheduler, DDPMScheduler]					,       _snake_case       :							VQModel					,       ):






								"""simple docstring"""


								super().__init__()

								self.register_modules(
								    text_encoder=_snake_case					,       tokenizer=_snake_case					,       unet=_snake_case					,       scheduler=_snake_case					,       movq=_snake_case					,       )
								A__	  =     2 ** (len(self.movq.config.block_out_channels       ) - 1)



					def      _a				(  self       :							List[Any]					,       _snake_case       :							List[str]					,       _snake_case       :							int					,       _snake_case       :							Tuple					,       _snake_case       :							Any					,       _snake_case       :							Tuple					,       _snake_case       :							Optional[Any]       ):






								"""simple docstring"""


								if latents is None:
											A__	  =     randn_tensor(_snake_case					,       generator=_snake_case					,       device=_snake_case					,       dtype=_snake_case       )
								else:
											if latents.shape != shape:
														raise ValueError(F'''Unexpected latents shape, got {latents.shape}, expected {shape}'''       )
											A__	  =     latents.to(_snake_case       )

								A__	  =     latents * scheduler.init_noise_sigma
								return latents



					def      _a				(  self       :							Dict					,       _snake_case       :							List[str]					,       _snake_case       :							Any					,       _snake_case       :							Union[str, Any]					,       _snake_case       :							Any					,       _snake_case       :							Tuple=None					,       ):






								"""simple docstring"""


								A__	  =     len(_snake_case       ) if isinstance(_snake_case					,       _snake_case       ) else 1
								# get prompt text embeddings
								A__	  =     self.tokenizer(
								    _snake_case					,       padding='max_length'					,       truncation=_snake_case					,       max_length=77					,       return_attention_mask=_snake_case					,       add_special_tokens=_snake_case					,       return_tensors='pt'					,       )

								A__	  =     text_inputs.input_ids
								A__	  =     self.tokenizer(_snake_case					,       padding='longest'					,       return_tensors='pt'       ).input_ids

								if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(_snake_case					,       _snake_case       ):
											A__	  =     self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]       )
											logger.warning(
											    'The following part of your input was truncated because CLIP can only handle sequences up to'
											    F''' {self.tokenizer.model_max_length} tokens: {removed_text}'''       )

								A__	  =     text_input_ids.to(_snake_case       )
								A__	  =     text_inputs.attention_mask.to(_snake_case       )

								A__  ,			A__	  =     self.text_encoder(
								    input_ids=_snake_case					,       attention_mask=_snake_case       )

								A__	  =     prompt_embeds.repeat_interleave(_snake_case					,       dim=0       )
								A__	  =     text_encoder_hidden_states.repeat_interleave(_snake_case					,       dim=0       )
								A__	  =     text_mask.repeat_interleave(_snake_case					,       dim=0       )

								if do_classifier_free_guidance:
											A__	  =     42
											if negative_prompt is None:
														A__	  =     [''] * batch_size
											elif type(_snake_case       ) is not type(_snake_case       ):
														raise TypeError(
														    F'''`negative_prompt` should be the same type to `prompt`, but got {type(_snake_case       )} !='''
														    F''' {type(_snake_case       )}.'''       )
											elif isinstance(_snake_case					,       _snake_case       ):
														A__	  =     [negative_prompt]
											elif batch_size != len(_snake_case       ):
														raise ValueError(
														    F'''`negative_prompt`: {negative_prompt} has batch size {len(_snake_case       )}, but `prompt`:'''
														    F''' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches'''
														    ' the batch size of `prompt`.'       )
											else:
														A__	  =     negative_prompt

											A__	  =     self.tokenizer(
											    _snake_case					,       padding='max_length'					,       max_length=77					,       truncation=_snake_case					,       return_attention_mask=_snake_case					,       add_special_tokens=_snake_case					,       return_tensors='pt'					,       )
											A__	  =     uncond_input.input_ids.to(_snake_case       )
											A__	  =     uncond_input.attention_mask.to(_snake_case       )

											A__  ,			A__	  =     self.text_encoder(
											    input_ids=_snake_case					,       attention_mask=_snake_case       )

											# duplicate unconditional embeddings for each generation per prompt, using mps friendly method

											A__	  =     negative_prompt_embeds.shape[1]
											A__	  =     negative_prompt_embeds.repeat(1					,       _snake_case       )
											A__	  =     negative_prompt_embeds.view(batch_size * num_images_per_prompt					,       _snake_case       )

											A__	  =     uncond_text_encoder_hidden_states.shape[1]
											A__	  =     uncond_text_encoder_hidden_states.repeat(1					,       _snake_case					,       1       )
											A__	  =     uncond_text_encoder_hidden_states.view(
											    batch_size * num_images_per_prompt					,       _snake_case					,       -1       )
											A__	  =     uncond_text_mask.repeat_interleave(_snake_case					,       dim=0       )

											# done duplicates

											# For classifier free guidance, we need to do two forward passes.
											# Here we concatenate the unconditional and text embeddings into a single batch
											# to avoid doing two forward passes
											A__	  =     torch.cat([negative_prompt_embeds, prompt_embeds]       )
											A__	  =     torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]       )

											A__	  =     torch.cat([uncond_text_mask, text_mask]       )

								return prompt_embeds, text_encoder_hidden_states, text_mask



					def      _a				(  self       :							Union[str, Any]					,       _snake_case       :							int=0       ):






								"""simple docstring"""


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

								A__	  =     torch.device(F'''cuda:{gpu_id}'''       )

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



					def      _a				(  self       :							Any					,       _snake_case       :							Optional[Any]=0       ):






								"""simple docstring"""


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

								A__	  =     torch.device(F'''cuda:{gpu_id}'''       )

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

								A__	  =     None
								for cpu_offloaded_model in [self.text_encoder, self.unet, self.movq]:
											A__  ,			A__	  =     cpu_offload_with_hook(_snake_case					,       _snake_case					,       prev_module_hook=_snake_case       )

								if self.safety_checker is not None:
											A__  ,			A__	  =     cpu_offload_with_hook(self.safety_checker					,       _snake_case					,       prev_module_hook=_snake_case       )

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



					@property
					# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
					def      _a				(  self       :							Union[str, Any]       ):






								"""simple docstring"""


								if not hasattr(self.unet					,       '_hf_hook'       ):
											return self.device
								for module in self.unet.modules():
											if (
											    hasattr(_snake_case					,       '_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(_snake_case       )
					def __call__(  self       :							List[str]					,       _snake_case       :							Union[str, List[str]]					,       _snake_case       :							Union[torch.FloatTensor, List[torch.FloatTensor]]					,       _snake_case       :							Union[torch.FloatTensor, List[torch.FloatTensor]]					,       _snake_case       :							Optional[Union[str, List[str]]] = None					,       _snake_case       :							int = 5_12					,       _snake_case       :							int = 5_12					,       _snake_case       :							int = 1_00					,       _snake_case       :							float = 4.0					,       _snake_case       :							int = 1					,       _snake_case       :							Optional[Union[torch.Generator, List[torch.Generator]]] = None					,       _snake_case       :							Optional[torch.FloatTensor] = None					,       _snake_case       :							Optional[str] = "pil"					,       _snake_case       :							bool = True					,       ):






								"""simple docstring"""


								if isinstance(_snake_case					,       _snake_case       ):
											A__	  =     1
								elif isinstance(_snake_case					,       _snake_case       ):
											A__	  =     len(_snake_case       )
								else:
											raise ValueError(F'''`prompt` has to be of type `str` or `list` but is {type(_snake_case       )}'''       )

								A__	  =     self._execution_device

								A__	  =     batch_size * num_images_per_prompt
								A__	  =     guidance_scale > 1.0

								A__  ,			A__  ,			A__	  =     self._encode_prompt(
								    _snake_case					,       _snake_case					,       _snake_case					,       _snake_case					,       _snake_case       )

								if isinstance(_snake_case					,       _snake_case       ):
											A__	  =     torch.cat(_snake_case					,       dim=0       )
								if isinstance(_snake_case					,       _snake_case       ):
											A__	  =     torch.cat(_snake_case					,       dim=0       )

								if do_classifier_free_guidance:
											A__	  =     image_embeds.repeat_interleave(_snake_case					,       dim=0       )
											A__	  =     negative_image_embeds.repeat_interleave(_snake_case					,       dim=0       )

								A__	  =     torch.cat([negative_image_embeds, image_embeds]					,       dim=0       ).to(
								    dtype=prompt_embeds.dtype					,       device=_snake_case       )

								self.scheduler.set_timesteps(_snake_case					,       device=_snake_case       )
								A__	  =     self.scheduler.timesteps

								A__	  =     self.unet.config.in_channels

								A__  ,			A__	  =     get_new_h_w(_snake_case					,       _snake_case					,       self.movq_scale_factor       )

								# create initial latent
								A__	  =     self.prepare_latents(
								    (batch_size, num_channels_latents, height, width)					,       text_encoder_hidden_states.dtype					,       _snake_case					,       _snake_case					,       _snake_case					,       self.scheduler					,       )

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

											A__	  =     {'text_embeds': prompt_embeds, 'image_embeds': image_embeds}
											A__	  =     self.unet(
											    sample=_snake_case					,       timestep=_snake_case					,       encoder_hidden_states=_snake_case					,       added_cond_kwargs=_snake_case					,       return_dict=_snake_case					,       )[0]

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

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

											# compute the previous noisy sample x_t -> x_t-1
											A__	  =     self.scheduler.step(
											    _snake_case					,       _snake_case					,       _snake_case					,       generator=_snake_case					,       ).prev_sample
								# post-processing
								A__	  =     self.movq.decode(_snake_case					,       force_not_quantize=_snake_case       )['sample']

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

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

								if output_type == "pil":
											A__	  =     self.numpy_to_pil(_snake_case       )

								if not return_dict:
											return (image,)

								return ImagePipelineOutput(images=_snake_case       )









import inspect
import unittest

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

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


if is_torch_available():
							import torch

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


if is_vision_available():
							from PIL import Image

							from transformers import SegformerImageProcessor


class   __lowerCAmelCase  (		UpperCAmelCase_   ):
					"""simple docstring"""






					def      _a				(  self       :							List[str]       ):






								"""simple docstring"""


								A__	  =     self.config_class(**self.inputs_dict       )
								self.parent.assertTrue(hasattr(_snake_case					,       'hidden_sizes'       )       )
								self.parent.assertTrue(hasattr(_snake_case					,       'num_attention_heads'       )       )
								self.parent.assertTrue(hasattr(_snake_case					,       'num_encoder_blocks'       )       )







class   __lowerCAmelCase  :
					"""simple docstring"""






					def __init__(  self       :							Any					,       _snake_case       :							str					,       _snake_case       :							Union[str, Any]=13					,       _snake_case       :							Any=64					,       _snake_case       :							Optional[Any]=3					,       _snake_case       :							Dict=4					,       _snake_case       :							Tuple=[2, 2, 2, 2]					,       _snake_case       :							str=[8, 4, 2, 1]					,       _snake_case       :							Union[str, Any]=[16, 32, 64, 1_28]					,       _snake_case       :							int=[1, 4, 8, 16]					,       _snake_case       :							List[str]=[1, 2, 4, 8]					,       _snake_case       :							int=True					,       _snake_case       :							int=True					,       _snake_case       :							Union[str, Any]="gelu"					,       _snake_case       :							Optional[int]=0.1					,       _snake_case       :							Tuple=0.1					,       _snake_case       :							Dict=0.02					,       _snake_case       :							Tuple=3					,       _snake_case       :							int=None					,       ):






								"""simple docstring"""


								A__	  =     parent
								A__	  =     batch_size
								A__	  =     image_size
								A__	  =     num_channels
								A__	  =     num_encoder_blocks
								A__	  =     sr_ratios
								A__	  =     depths
								A__	  =     hidden_sizes
								A__	  =     downsampling_rates
								A__	  =     num_attention_heads
								A__	  =     is_training
								A__	  =     use_labels
								A__	  =     hidden_act
								A__	  =     hidden_dropout_prob
								A__	  =     attention_probs_dropout_prob
								A__	  =     initializer_range
								A__	  =     num_labels
								A__	  =     scope



					def      _a				(  self       :							int       ):






								"""simple docstring"""


								A__	  =     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

								A__	  =     None
								if self.use_labels:
											A__	  =     ids_tensor([self.batch_size, self.image_size, self.image_size]					,       self.num_labels       )

								A__	  =     self.get_config()
								return config, pixel_values, labels



					def      _a				(  self       :							int       ):






								"""simple docstring"""


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



					def      _a				(  self       :							int					,       _snake_case       :							Optional[Any]					,       _snake_case       :							int					,       _snake_case       :							Any       ):






								"""simple docstring"""


								A__	  =     SegformerModel(config=_snake_case       )
								model.to(_snake_case       )
								model.eval()
								A__	  =     model(_snake_case       )
								A__	  =     A__	  =     self.image_size // (self.downsampling_rates[-1] * 2)
								self.parent.assertEqual(
								    result.last_hidden_state.shape					,       (self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)       )



					def      _a				(  self       :							Union[str, Any]					,       _snake_case       :							Union[str, Any]					,       _snake_case       :							Tuple					,       _snake_case       :							Dict       ):






								"""simple docstring"""


								A__	  =     self.num_labels
								A__	  =     SegformerForSemanticSegmentation(_snake_case       )
								model.to(_snake_case       )
								model.eval()
								A__	  =     model(_snake_case       )
								self.parent.assertEqual(
								    result.logits.shape					,       (self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)       )
								A__	  =     model(_snake_case					,       labels=_snake_case       )
								self.parent.assertEqual(
								    result.logits.shape					,       (self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)       )
								self.parent.assertGreater(result.loss					,       0.0       )



					def      _a				(  self       :							List[str]					,       _snake_case       :							Optional[Any]					,       _snake_case       :							Union[str, Any]					,       _snake_case       :							List[str]       ):






								"""simple docstring"""


								A__	  =     1
								A__	  =     SegformerForSemanticSegmentation(config=_snake_case       )
								model.to(_snake_case       )
								model.eval()
								A__	  =     torch.randint(0					,       1					,       (self.batch_size, self.image_size, self.image_size)       ).to(_snake_case       )
								A__	  =     model(_snake_case					,       labels=_snake_case       )
								self.parent.assertGreater(result.loss					,       0.0       )



					def      _a				(  self       :							List[Any]       ):






								"""simple docstring"""


								A__	  =     self.prepare_config_and_inputs()
								A__  ,			A__  ,			A__	  =     config_and_inputs
								A__	  =     {'pixel_values': pixel_values}
								return config, inputs_dict







@require_torch
class   __lowerCAmelCase  (		UpperCAmelCase_				,  UpperCAmelCase_				,  unittest.TestCase   ):
					"""simple docstring"""






					A__					:		Optional[int]      							=     (
					    (
					        SegformerModel,
					        SegformerForSemanticSegmentation,
					        SegformerForImageClassification,
					    )
					    if is_torch_available()
					    else ()
					)
					A__					:		Union[str, Any]      							=     (
					    {
					        "feature-extraction": SegformerModel,
					        "image-classification": SegformerForImageClassification,
					        "image-segmentation": SegformerForSemanticSegmentation,
					    }
					    if is_torch_available()
					    else {}
					)

					A__					:		Optional[Any]      							=     True
					A__					:		str      							=     False
					A__					:		Tuple      							=     False
					A__					:		Dict      							=     False



					def      _a				(  self       :							Union[str, Any]       ):






								"""simple docstring"""


								A__	  =     SegformerModelTester(self       )
								A__	  =     SegformerConfigTester(self					,       config_class=_snake_case       )



					def      _a				(  self       :							Optional[int]       ):






								"""simple docstring"""


								self.config_tester.run_common_tests()



					def      _a				(  self       :							Optional[Any]       ):






								"""simple docstring"""


								A__	  =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_snake_case       )



					def      _a				(  self       :							List[Any]       ):






								"""simple docstring"""


								A__	  =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_binary_image_segmentation(*_snake_case       )



					def      _a				(  self       :							Tuple       ):






								"""simple docstring"""


								A__	  =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_segmentation(*_snake_case       )



					@unittest.skip('SegFormer does not use inputs_embeds'       )
					def      _a				(  self       :							List[Any]       ):






								"""simple docstring"""


								pass



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






								"""simple docstring"""


								pass



					def      _a				(  self       :							Dict       ):






								"""simple docstring"""


								A__  ,			A__	  =     self.model_tester.prepare_config_and_inputs_for_common()

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

											A__	  =     ['pixel_values']
											self.assertListEqual(arg_names[:1]					,       _snake_case       )



					def      _a				(  self       :							Dict       ):






								"""simple docstring"""


								A__  ,			A__	  =     self.model_tester.prepare_config_and_inputs_for_common()
								A__	  =     True

								for model_class in self.all_model_classes:
											A__	  =     True
											A__	  =     False
											A__	  =     True
											A__	  =     model_class(_snake_case       )
											model.to(_snake_case       )
											model.eval()
											with torch.no_grad():
														A__	  =     model(**self._prepare_for_class(_snake_case					,       _snake_case       )       )
											A__	  =     outputs.attentions

											A__	  =     sum(self.model_tester.depths       )
											self.assertEqual(len(_snake_case       )					,       _snake_case       )

											# check that output_attentions also work using config
											del inputs_dict["output_attentions"]
											A__	  =     True
											A__	  =     model_class(_snake_case       )
											model.to(_snake_case       )
											model.eval()
											with torch.no_grad():
														A__	  =     model(**self._prepare_for_class(_snake_case					,       _snake_case       )       )
											A__	  =     outputs.attentions

											self.assertEqual(len(_snake_case       )					,       _snake_case       )

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

											# verify the last attentions (last block, last layer)
											A__	  =     (self.model_tester.image_size // 32) ** 2
											A__	  =     (self.model_tester.image_size // (32 * self.model_tester.sr_ratios[-1])) ** 2
											self.assertListEqual(
											    list(attentions[-1].shape[-3:]       )					,       [self.model_tester.num_attention_heads[-1], expected_seq_len, expected_reduced_seq_len]					,       )
											A__	  =     len(_snake_case       )

											# Check attention is always last and order is fine
											A__	  =     True
											A__	  =     True
											A__	  =     model_class(_snake_case       )
											model.to(_snake_case       )
											model.eval()
											with torch.no_grad():
														A__	  =     model(**self._prepare_for_class(_snake_case					,       _snake_case       )       )

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

											A__	  =     outputs.attentions

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



					def      _a				(  self       :							Union[str, Any]       ):






								"""simple docstring"""


								def check_hidden_states_output(_snake_case       :							Dict					,       _snake_case       :							int					,       _snake_case       :							List[Any]       ):
											A__	  =     model_class(_snake_case       )
											model.to(_snake_case       )
											model.eval()

											with torch.no_grad():
														A__	  =     model(**self._prepare_for_class(_snake_case					,       _snake_case       )       )

											A__	  =     outputs.hidden_states

											A__	  =     self.model_tester.num_encoder_blocks
											self.assertEqual(len(_snake_case       )					,       _snake_case       )

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

								A__  ,			A__	  =     self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
											A__	  =     True
											check_hidden_states_output(_snake_case					,       _snake_case					,       _snake_case       )

											# check that output_hidden_states also work using config
											del inputs_dict["output_hidden_states"]
											A__	  =     True

											check_hidden_states_output(_snake_case					,       _snake_case					,       _snake_case       )



					def      _a				(  self       :							Tuple       ):






								"""simple docstring"""


								if not self.model_tester.is_training:
											return

								A__  ,			A__	  =     self.model_tester.prepare_config_and_inputs_for_common()
								A__	  =     True

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

											A__	  =     model_class(_snake_case       )
											model.to(_snake_case       )
											model.train()
											A__	  =     self._prepare_for_class(_snake_case					,       _snake_case					,       return_labels=_snake_case       )
											A__	  =     model(**_snake_case       ).loss
											loss.backward()



					@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'       )
					def      _a				(  self       :							Optional[Any]       ):






								"""simple docstring"""


								pass



					@slow
					def      _a				(  self       :							Tuple       ):






								"""simple docstring"""


								for model_name in SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											A__	  =     SegformerModel.from_pretrained(_snake_case       )
											self.assertIsNotNone(_snake_case       )







def A    (							)    ->   str:
			A__	  =     Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'       )
			return image






@require_torch
class   __lowerCAmelCase  (		unittest.TestCase   ):
					"""simple docstring"""






					@slow
					def      _a				(  self       :							Dict       ):






								"""simple docstring"""


								A__	  =     SegformerImageProcessor(
								    image_scale=(5_12, 5_12)					,       keep_ratio=_snake_case					,       align=_snake_case					,       do_random_crop=_snake_case       )
								A__	  =     SegformerForSemanticSegmentation.from_pretrained('nvidia/segformer-b0-finetuned-ade-512-512'       ).to(
								    _snake_case       )

								A__	  =     prepare_img()
								A__	  =     image_processor(images=_snake_case					,       return_tensors='pt'       )
								A__	  =     encoded_inputs.pixel_values.to(_snake_case       )

								with torch.no_grad():
											A__	  =     model(_snake_case       )

								A__	  =     torch.Size((1, model.config.num_labels, 1_28, 1_28)       )
								self.assertEqual(outputs.logits.shape					,       _snake_case       )

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



					@slow
					def      _a				(  self       :							Optional[Any]       ):






								"""simple docstring"""


								A__	  =     SegformerImageProcessor(
								    image_scale=(5_12, 5_12)					,       keep_ratio=_snake_case					,       align=_snake_case					,       do_random_crop=_snake_case       )
								A__	  =     SegformerForSemanticSegmentation.from_pretrained(
								    'nvidia/segformer-b1-finetuned-cityscapes-1024-1024'       ).to(_snake_case       )

								A__	  =     prepare_img()
								A__	  =     image_processor(images=_snake_case					,       return_tensors='pt'       )
								A__	  =     encoded_inputs.pixel_values.to(_snake_case       )

								with torch.no_grad():
											A__	  =     model(_snake_case       )

								A__	  =     torch.Size((1, model.config.num_labels, 1_28, 1_28)       )
								self.assertEqual(outputs.logits.shape					,       _snake_case       )

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



					@slow
					def      _a				(  self       :							Any       ):






								"""simple docstring"""


								A__	  =     SegformerImageProcessor(
								    image_scale=(5_12, 5_12)					,       keep_ratio=_snake_case					,       align=_snake_case					,       do_random_crop=_snake_case       )
								A__	  =     SegformerForSemanticSegmentation.from_pretrained('nvidia/segformer-b0-finetuned-ade-512-512'       ).to(
								    _snake_case       )

								A__	  =     prepare_img()
								A__	  =     image_processor(images=_snake_case					,       return_tensors='pt'       )
								A__	  =     encoded_inputs.pixel_values.to(_snake_case       )

								with torch.no_grad():
											A__	  =     model(_snake_case       )

								A__	  =     outputs.logits.detach().cpu()

								A__	  =     image_processor.post_process_semantic_segmentation(outputs=_snake_case					,       target_sizes=[(5_00, 3_00)]       )
								A__	  =     torch.Size((5_00, 3_00)       )
								self.assertEqual(segmentation[0].shape					,       _snake_case       )

								A__	  =     image_processor.post_process_semantic_segmentation(outputs=_snake_case       )
								A__	  =     torch.Size((1_28, 1_28)       )
								self.assertEqual(segmentation[0].shape					,       _snake_case       )




