





import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, CycleDiffusionPipeline, DDIMScheduler, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class 							_lowerCAmelCase     (      UpperCAmelCase_				,     UpperCAmelCase_				,     unittest.TestCase				):


  '''simple docstring'''
  a_				:							str          =CycleDiffusionPipeline
  a_				:							Dict          =TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
      """negative_prompt""",
      """height""",
      """width""",
      """negative_prompt_embeds""",
  }
  a_				:							int          =PipelineTesterMixin.required_optional_params - {"""latents"""}
  a_				:							Optional[Any]          =TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({"""source_prompt"""}				)
  a_				:							int          =IMAGE_TO_IMAGE_IMAGE_PARAMS
  a_				:							Tuple          =IMAGE_TO_IMAGE_IMAGE_PARAMS







  def UpperCamelCase_			(   self   :       Optional[Any]				):


    '''simple docstring'''





    torch.manual_seed(0				)
    _snake_case :					List[str]								=						UNetaDConditionModel(
        block_out_channels=(32, 64)   ,					layers_per_block=2   ,					sample_size=32   ,					in_channels=4   ,					out_channels=4   ,					down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')   ,					up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')   ,					cross_attention_dim=32   ,					)
    _snake_case :					List[Any]								=						DDIMScheduler(
        beta_start=0.0_00_85   ,					beta_end=0.0_12   ,					beta_schedule='scaled_linear'   ,					num_train_timesteps=10_00   ,					clip_sample=_lowercase   ,					set_alpha_to_one=_lowercase   ,					)
    torch.manual_seed(0				)
    _snake_case :					List[Any]								=						AutoencoderKL(
        block_out_channels=[32, 64]   ,					in_channels=3   ,					out_channels=3   ,					down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']   ,					up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']   ,					latent_channels=4   ,					)
    torch.manual_seed(0				)
    _snake_case :					List[str]								=						CLIPTextConfig(
        bos_token_id=0   ,					eos_token_id=2   ,					hidden_size=32   ,					intermediate_size=37   ,					layer_norm_eps=1e-0_5   ,					num_attention_heads=4   ,					num_hidden_layers=5   ,					pad_token_id=1   ,					vocab_size=10_00   ,					)
    _snake_case :					int								=						CLIPTextModel(_lowercase				)
    _snake_case :					Tuple								=						CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'				)

    _snake_case :					List[Any]								=						{
        'unet': unet,
        'scheduler': scheduler,
        'vae': vae,
        'text_encoder': text_encoder,
        'tokenizer': tokenizer,
        'safety_checker': None,
        'feature_extractor': None,
    }
    return components







  def UpperCamelCase_			(   self   :       List[str]   ,					UpperCamelCase   :       Optional[Any]   ,					UpperCamelCase   :       Dict=0				):


    '''simple docstring'''





    _snake_case :					Optional[int]								=						floats_tensor((1, 3, 32, 32)   ,					rng=random.Random(_lowercase				)				).to(_lowercase				)
    _snake_case :					Dict								=						image / 2 + 0.5
    if str(_lowercase				).startswith('mps'				):
      _snake_case :					Any								=						torch.manual_seed(_lowercase				)
    else:
      _snake_case :					int								=						torch.Generator(device=_lowercase				).manual_seed(_lowercase				)
    _snake_case :					int								=						{
        'prompt': 'An astronaut riding an elephant',
        'source_prompt': 'An astronaut riding a horse',
        'image': image,
        'generator': generator,
        'num_inference_steps': 2,
        'eta': 0.1,
        'strength': 0.8,
        'guidance_scale': 3,
        'source_guidance_scale': 1,
        'output_type': 'numpy',
    }
    return inputs







  def UpperCamelCase_			(   self   :       Tuple				):


    '''simple docstring'''





    _snake_case :					Any								=						'cpu'  # ensure determinism for the device-dependent torch.Generator

    _snake_case :					str								=						self.get_dummy_components()
    _snake_case :					Any								=						CycleDiffusionPipeline(**_lowercase				)
    _snake_case :					Any								=						pipe.to(_lowercase				)
    pipe.set_progress_bar_config(disable=_lowercase				)

    _snake_case :					str								=						self.get_dummy_inputs(_lowercase				)
    _snake_case :					Optional[Any]								=						pipe(**_lowercase				)
    _snake_case :					List[str]								=						output.images

    _snake_case :					Optional[Any]								=						images[0, -3:, -3:, -1]

    assert images.shape == (1, 32, 32, 3)
    _snake_case :					Tuple								=						np.array([0.44_59, 0.49_43, 0.45_44, 0.66_43, 0.54_74, 0.43_27, 0.57_01, 0.59_59, 0.51_79]				)

    assert np.abs(image_slice.flatten() - expected_slice				).max() < 1e-2







  @unittest.skipIf(torch_device != 'cuda'   ,					'This test requires a GPU'				)
  def UpperCamelCase_			(   self   :       Dict				):


    '''simple docstring'''





    _snake_case :					Tuple								=						self.get_dummy_components()
    for name, module in components.items():
      if hasattr(_lowercase   ,					'half'				):
        _snake_case :					str								=						module.half()
    _snake_case :					List[str]								=						CycleDiffusionPipeline(**_lowercase				)
    _snake_case :					List[str]								=						pipe.to(_lowercase				)
    pipe.set_progress_bar_config(disable=_lowercase				)

    _snake_case :					Tuple								=						self.get_dummy_inputs(_lowercase				)
    _snake_case :					str								=						pipe(**_lowercase				)
    _snake_case :					Tuple								=						output.images

    _snake_case :					Dict								=						images[0, -3:, -3:, -1]

    assert images.shape == (1, 32, 32, 3)
    _snake_case :					Optional[Any]								=						np.array([0.35_06, 0.45_43, 0.4_46, 0.45_75, 0.51_95, 0.41_55, 0.52_73, 0.5_18, 0.41_16]				)

    assert np.abs(image_slice.flatten() - expected_slice				).max() < 1e-2







  @skip_mps
  def UpperCamelCase_			(   self   :       str				):


    '''simple docstring'''





    return super().test_save_load_local()







  @unittest.skip('non-deterministic pipeline'				)
  def UpperCamelCase_			(   self   :       Optional[Any]				):


    '''simple docstring'''





    return super().test_inference_batch_single_identical()







  @skip_mps
  def UpperCamelCase_			(   self   :       Optional[Any]				):


    '''simple docstring'''





    return super().test_dict_tuple_outputs_equivalent()







  @skip_mps
  def UpperCamelCase_			(   self   :       Optional[Any]				):


    '''simple docstring'''





    return super().test_save_load_optional_components()







  @skip_mps
  def UpperCamelCase_			(   self   :       Union[str, Any]				):


    '''simple docstring'''





    return super().test_attention_slicing_forward_pass()




@slow
@require_torch_gpu
class 							_lowerCAmelCase     (      unittest.TestCase				):


  '''simple docstring'''







  def UpperCamelCase_			(   self   :       int				):


    '''simple docstring'''





    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()







  def UpperCamelCase_			(   self   :       str				):


    '''simple docstring'''





    _snake_case :					Dict								=						load_image(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
        '/cycle-diffusion/black_colored_car.png'				)
    _snake_case :					Optional[Any]								=						load_numpy(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/cycle-diffusion/blue_colored_car_fp16.npy'				)
    _snake_case :					Optional[Any]								=						init_image.resize((5_12, 5_12)				)

    _snake_case :					Dict								=						'CompVis/stable-diffusion-v1-4'
    _snake_case :					Optional[Any]								=						DDIMScheduler.from_pretrained(_lowercase   ,					subfolder='scheduler'				)
    _snake_case :					Optional[Any]								=						CycleDiffusionPipeline.from_pretrained(
        _lowercase   ,					scheduler=_lowercase   ,					safety_checker=_lowercase   ,					torch_dtype=torch.floataa   ,					revision='fp16'				)

    pipe.to(_lowercase				)
    pipe.set_progress_bar_config(disable=_lowercase				)
    pipe.enable_attention_slicing()

    _snake_case :					str								=						'A black colored car'
    _snake_case :					Union[str, Any]								=						'A blue colored car'

    _snake_case :					Optional[Any]								=						torch.manual_seed(0				)
    _snake_case :					Optional[int]								=						pipe(
        prompt=_lowercase   ,					source_prompt=_lowercase   ,					image=_lowercase   ,					num_inference_steps=1_00   ,					eta=0.1   ,					strength=0.85   ,					guidance_scale=3   ,					source_guidance_scale=1   ,					generator=_lowercase   ,					output_type='np'   ,					)
    _snake_case :					Optional[int]								=						output.images

    # the values aren't exactly equal, but the images look the same visually
    assert np.abs(image - expected_image				).max() < 5e-1







  def UpperCamelCase_			(   self   :       Tuple				):


    '''simple docstring'''





    _snake_case :					Dict								=						load_image(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
        '/cycle-diffusion/black_colored_car.png'				)
    _snake_case :					int								=						load_numpy(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/cycle-diffusion/blue_colored_car.npy'				)
    _snake_case :					str								=						init_image.resize((5_12, 5_12)				)

    _snake_case :					Optional[Any]								=						'CompVis/stable-diffusion-v1-4'
    _snake_case :					Optional[Any]								=						DDIMScheduler.from_pretrained(_lowercase   ,					subfolder='scheduler'				)
    _snake_case :					int								=						CycleDiffusionPipeline.from_pretrained(_lowercase   ,					scheduler=_lowercase   ,					safety_checker=_lowercase				)

    pipe.to(_lowercase				)
    pipe.set_progress_bar_config(disable=_lowercase				)
    pipe.enable_attention_slicing()

    _snake_case :					Tuple								=						'A black colored car'
    _snake_case :					Dict								=						'A blue colored car'

    _snake_case :					Dict								=						torch.manual_seed(0				)
    _snake_case :					str								=						pipe(
        prompt=_lowercase   ,					source_prompt=_lowercase   ,					image=_lowercase   ,					num_inference_steps=1_00   ,					eta=0.1   ,					strength=0.85   ,					guidance_scale=3   ,					source_guidance_scale=1   ,					generator=_lowercase   ,					output_type='np'   ,					)
    _snake_case :					Union[str, Any]								=						output.images

    assert np.abs(image - expected_image				).max() < 2e-2







import copy
from typing import Dict, List, Optional

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto import CONFIG_MAPPING


SCREAMING_SNAKE_CASE		:Optional[Any]						=     {
    """facebook/mask2former-swin-small-coco-instance""": (
        """https://huggingface.co/facebook/mask2former-swin-small-coco-instance/blob/main/config.json"""
    )
    # See all Mask2Former models at https://huggingface.co/models?filter=mask2former
}

SCREAMING_SNAKE_CASE		:Union[str, Any]						=     logging.get_logger(__name__)






class 						__magic_name__					(							snake_case				):
     UpperCamelCase_	:Optional[int]										=      """mask2former"""
     UpperCamelCase_	:Dict										=      ["""swin"""]
     UpperCamelCase_	:List[Any]										=      {"""hidden_size""": """hidden_dim"""}






     def __init__(				self		,						_lowercase = None		,						_lowercase = 256		,						_lowercase = 256		,						_lowercase = 256		,						_lowercase = 1_024		,						_lowercase = "relu"		,						_lowercase = 6		,						_lowercase = 10		,						_lowercase = 8		,						_lowercase = 0.0		,						_lowercase = 2_048		,						_lowercase = False		,						_lowercase = False		,						_lowercase = 4		,						_lowercase = 255		,						_lowercase = 100		,						_lowercase = 0.1		,						_lowercase = 2.0		,						_lowercase = 5.0		,						_lowercase = 5.0		,						_lowercase = 12_544		,						_lowercase = 3.0		,						_lowercase = 0.75		,						_lowercase = 0.02		,						_lowercase = 1.0		,						_lowercase = True		,						_lowercase = [4, 8, 16, 32]		,						_lowercase = None		,						**_lowercase		,						)->     Union[str, Any]:
         if backbone_config is None:
             logger.info("`backbone_config` is `None`. Initializing the config with the default `Swin` backbone."	)
             UpperCamelCase_			  =						CONFIG_MAPPING["swin"](
                 image_size=224		,						in_channels=3		,						patch_size=4		,						embed_dim=96		,						depths=[2, 2, 18, 2]		,						num_heads=[3, 6, 12, 24]		,						window_size=7		,						drop_path_rate=0.3		,						use_absolute_embeddings=_lowercase		,						out_features=["stage1", "stage2", "stage3", "stage4"]		,						)

         if isinstance(_lowercase		,						_lowercase	):
             UpperCamelCase_			  =						backbone_config.pop("model_type"	)
             UpperCamelCase_			  =						CONFIG_MAPPING[backbone_model_type]
             UpperCamelCase_			  =						config_class.from_dict(_lowercase	)

         # verify that the backbone is supported
         if backbone_config.model_type not in self.backbones_supported:
             logger.warning_once(
                 F"Backbone {backbone_config.model_type} is not a supported model and may not be compatible with Mask2Former. "
                 F"Supported model types: {','.join(self.backbones_supported	)}"	)

         UpperCamelCase_			  =						backbone_config
         UpperCamelCase_			  =						feature_size
         UpperCamelCase_			  =						mask_feature_size
         UpperCamelCase_			  =						hidden_dim
         UpperCamelCase_			  =						encoder_feedforward_dim
         UpperCamelCase_			  =						activation_function
         UpperCamelCase_			  =						encoder_layers
         UpperCamelCase_			  =						decoder_layers
         UpperCamelCase_			  =						num_attention_heads
         UpperCamelCase_			  =						dropout
         UpperCamelCase_			  =						dim_feedforward
         UpperCamelCase_			  =						pre_norm
         UpperCamelCase_			  =						enforce_input_projection
         UpperCamelCase_			  =						common_stride
         UpperCamelCase_			  =						ignore_value
         UpperCamelCase_			  =						num_queries
         UpperCamelCase_			  =						no_object_weight
         UpperCamelCase_			  =						class_weight
         UpperCamelCase_			  =						mask_weight
         UpperCamelCase_			  =						dice_weight
         UpperCamelCase_			  =						train_num_points
         UpperCamelCase_			  =						oversample_ratio
         UpperCamelCase_			  =						importance_sample_ratio
         UpperCamelCase_			  =						init_std
         UpperCamelCase_			  =						init_xavier_std
         UpperCamelCase_			  =						use_auxiliary_loss
         UpperCamelCase_			  =						feature_strides
         UpperCamelCase_			  =						output_auxiliary_logits
         UpperCamelCase_			  =						decoder_layers

         super().__init__(**_lowercase	)






     @classmethod
     def 	UpperCAmelCase_				(				cls		,						_lowercase		,						**_lowercase	)->     Optional[int]:
         return cls(
             backbone_config=_lowercase		,						**_lowercase		,						)






     def 	UpperCAmelCase_				(				self	)->     Dict[str, any]:
         UpperCamelCase_			  =						copy.deepcopy(self.__dict__	)
         UpperCamelCase_			  =						self.backbone_config.to_dict()
         UpperCamelCase_			  =						self.__class__.model_type
         return output
