






import gc
import random
import unittest

import numpy as np
import torch

from diffusers import (
    DDIMScheduler,
    KandinskyVaaControlnetPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
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

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()






class 					SCREAMING_SNAKE_CASE_  (       _a  ,						unittest.TestCase ):
       """simple docstring"""


       __lowerCAmelCase    :		Optional[Any]					=KandinskyVaaControlnetPipeline
       __lowerCAmelCase    :		List[str]					=['''image_embeds''', '''negative_image_embeds''', '''hint''']
       __lowerCAmelCase    :		Optional[int]					=['''image_embeds''', '''negative_image_embeds''', '''hint''']
       __lowerCAmelCase    :		Tuple					=[
           '''generator''',
           '''height''',
           '''width''',
           '''latents''',
           '''guidance_scale''',
           '''num_inference_steps''',
           '''return_dict''',
           '''guidance_scale''',
           '''num_images_per_prompt''',
           '''output_type''',
           '''return_dict''',
       ]
       __lowerCAmelCase    :		List[str]					=False





       @property
       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

          return 32





       @property
       def 				UpperCamelCase__						(				self	:Union[str, Any]):



          """simple docstring"""

          return 32





       @property
       def 				UpperCamelCase__						(				self	:str):



          """simple docstring"""

          return self.time_input_dim





       @property
       def 				UpperCamelCase__						(				self	:List[str]):



          """simple docstring"""

          return self.time_input_dim * 4





       @property
       def 				UpperCamelCase__						(				self	:Any):



          """simple docstring"""

          return 100





       @property
       def 				UpperCamelCase__						(				self	:List[str]):



          """simple docstring"""

          torch.manual_seed(0)

          _lowercase              ={
              'in_channels': 8,
              # Out channels is double in channels because predicts mean and variance
              'out_channels': 8,
              'addition_embed_type': 'image_hint',
              'down_block_types': ('ResnetDownsampleBlock2D', 'SimpleCrossAttnDownBlock2D'),
              'up_block_types': ('SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D'),
              'mid_block_type': 'UNetMidBlock2DSimpleCrossAttn',
              'block_out_channels': (self.block_out_channels_a, self.block_out_channels_a * 2),
              'layers_per_block': 1,
              'encoder_hid_dim': self.text_embedder_hidden_size,
              'encoder_hid_dim_type': 'image_proj',
              'cross_attention_dim': self.cross_attention_dim,
              'attention_head_dim': 4,
              'resnet_time_scale_shift': 'scale_shift',
              'class_embed_type': None,
          }

          _lowercase              =UNetaDConditionModel(**snake_case)
          return model





       @property
       def 				UpperCamelCase__						(				self	:Optional[int]):



          """simple docstring"""

          return {
              "block_out_channels": [32, 32, 64, 64],
              "down_block_types": [
                  "DownEncoderBlock2D",
                  "DownEncoderBlock2D",
                  "DownEncoderBlock2D",
                  "AttnDownEncoderBlock2D",
              ],
              "in_channels": 3,
              "latent_channels": 4,
              "layers_per_block": 1,
              "norm_num_groups": 8,
              "norm_type": "spatial",
              "num_vq_embeddings": 12,
              "out_channels": 3,
              "up_block_types": ["AttnUpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"],
              "vq_embed_dim": 4,
          }





       @property
       def 				UpperCamelCase__						(				self	:Optional[int]):



          """simple docstring"""

          torch.manual_seed(0)
          _lowercase              =VQModel(**self.dummy_movq_kwargs)
          return model





       def 				UpperCamelCase__						(				self	:List[Any]):



          """simple docstring"""

          _lowercase              =self.dummy_unet
          _lowercase              =self.dummy_movq

          _lowercase              =DDIMScheduler(
              num_train_timesteps=1000,      beta_schedule='linear',      beta_start=0.0_0_0_8_5,      beta_end=0.0_1_2,      clip_sample=snake_case,      set_alpha_to_one=snake_case,      steps_offset=1,      prediction_type='epsilon',      thresholding=snake_case,      )

          _lowercase              ={
              'unet': unet,
              'scheduler': scheduler,
              'movq': movq,
          }
          return components





       def 				UpperCamelCase__						(				self	:Optional[Any],      snake_case	:List[str],      snake_case	:List[str]=0):



          """simple docstring"""

          _lowercase              =floats_tensor((1, self.text_embedder_hidden_size),      rng=random.Random(snake_case)).to(snake_case)
          _lowercase              =floats_tensor((1, self.text_embedder_hidden_size),      rng=random.Random(seed + 1)).to(
              snake_case)

          # create hint
          _lowercase              =floats_tensor((1, 3, 64, 64),      rng=random.Random(snake_case)).to(snake_case)

          if str(snake_case).startswith('mps'):
             _lowercase              =torch.manual_seed(snake_case)
          else:
             _lowercase              =torch.Generator(device=snake_case).manual_seed(snake_case)
          _lowercase              ={
              'image_embeds': image_embeds,
              'negative_image_embeds': negative_image_embeds,
              'hint': hint,
              'generator': generator,
              'height': 64,
              'width': 64,
              'guidance_scale': 4.0,
              'num_inference_steps': 2,
              'output_type': 'np',
          }
          return inputs





       def 				UpperCamelCase__						(				self	:Tuple):



          """simple docstring"""

          _lowercase              ='cpu'

          _lowercase              =self.get_dummy_components()

          _lowercase              =self.pipeline_class(**snake_case)
          _lowercase              =pipe.to(snake_case)

          pipe.set_progress_bar_config(disable=snake_case)

          _lowercase              =pipe(**self.get_dummy_inputs(snake_case))
          _lowercase              =output.images

          _lowercase              =pipe(
              **self.get_dummy_inputs(snake_case),      return_dict=snake_case,      )[0]

          _lowercase              =image[0, -3:, -3:, -1]
          _lowercase              =image_from_tuple[0, -3:, -3:, -1]

          assert image.shape == (1, 64, 64, 3)

          _lowercase              =np.array(
              [0.6_9_5_9_8_2_6, 0.8_6_8_2_7_9, 0.7_5_5_8_0_9_2, 0.6_8_7_6_9_4_6_7, 0.8_5_8_0_5_8_0_4, 0.6_5_9_7_7_4_9_6, 0.4_4_8_8_5_3_0_2, 0.5_9_5_9_1_1_1, 0.4_2_5_1_5_9_5])

          assert (
              np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
          ), f''' expected_slice {expected_slice}, but got {image_slice.flatten()}'''

          assert (
              np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2
          ), f''' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'''






@slow
@require_torch_gpu
class 					SCREAMING_SNAKE_CASE_  (       unittest.TestCase ):
       """simple docstring"""


       def 				UpperCamelCase__						(				self	:Optional[int]):



          """simple docstring"""

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





       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

          _lowercase              =load_numpy(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
              '/kandinskyv22/kandinskyv22_controlnet_robotcat_fp16.npy')

          _lowercase              =load_image(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
              '/kandinskyv22/hint_image_cat.png')
          _lowercase              =torch.from_numpy(np.array(snake_case)).float() / 2_5_5.0
          _lowercase              =hint.permute(2,      0,      1).unsqueeze(0)

          _lowercase              =KandinskyVaaPriorPipeline.from_pretrained(
              'kandinsky-community/kandinsky-2-2-prior',      torch_dtype=torch.floataa)
          pipe_prior.to(snake_case)

          _lowercase              =KandinskyVaaControlnetPipeline.from_pretrained(
              'kandinsky-community/kandinsky-2-2-controlnet-depth',      torch_dtype=torch.floataa)
          _lowercase              =pipeline.to(snake_case)
          pipeline.set_progress_bar_config(disable=snake_case)

          _lowercase              ='A robot, 4k photo'

          _lowercase              =torch.Generator(device='cuda').manual_seed(0)
          _lowercase       ,       _lowercase              =pipe_prior(
              snake_case,      generator=snake_case,      num_inference_steps=5,      negative_prompt='',      ).to_tuple()

          _lowercase              =torch.Generator(device='cuda').manual_seed(0)
          _lowercase              =pipeline(
              image_embeds=snake_case,      negative_image_embeds=snake_case,      hint=snake_case,      generator=snake_case,      num_inference_steps=100,      output_type='np',      )

          _lowercase              =output.images[0]

          assert image.shape == (512, 512, 3)

          assert_mean_pixel_difference(snake_case,      snake_case)










import json
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
from transformers import (
    CONFIG_MAPPING,
    IMAGE_PROCESSOR_MAPPING,
    AutoConfig,
    AutoImageProcessor,
    CLIPConfig,
    CLIPImageProcessor,
)
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER


sys.path.append(str(Path(__file__).parent.parent.parent.parent / "utils"))

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_image_processing import CustomImageProcessor  # noqa E402






class 					SCREAMING_SNAKE_CASE_  (       unittest.TestCase ):
       """simple docstring"""


       def 				UpperCamelCase__						(				self	:Any):



          """simple docstring"""

          _lowercase              =0





       def 				UpperCamelCase__						(				self	:Optional[int]):



          """simple docstring"""

          _lowercase              =AutoImageProcessor.from_pretrained('openai/clip-vit-base-patch32')
          self.assertIsInstance(snake_case,      snake_case)





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



          """simple docstring"""

          with tempfile.TemporaryDirectory() as tmpdirname:
             _lowercase              =Path(snake_case) / 'preprocessor_config.json'
             _lowercase              =Path(snake_case) / 'config.json'
             json.dump(
                 {'image_processor_type': 'CLIPImageProcessor', 'processor_class': 'CLIPProcessor'},      open(snake_case,      'w'),      )
             json.dump({'model_type': 'clip'},      open(snake_case,      'w'))

             _lowercase              =AutoImageProcessor.from_pretrained(snake_case)
             self.assertIsInstance(snake_case,      snake_case)





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



          """simple docstring"""

          with tempfile.TemporaryDirectory() as tmpdirname:
             _lowercase              =Path(snake_case) / 'preprocessor_config.json'
             _lowercase              =Path(snake_case) / 'config.json'
             json.dump(
                 {'feature_extractor_type': 'CLIPFeatureExtractor', 'processor_class': 'CLIPProcessor'},      open(snake_case,      'w'),      )
             json.dump({'model_type': 'clip'},      open(snake_case,      'w'))

             _lowercase              =AutoImageProcessor.from_pretrained(snake_case)
             self.assertIsInstance(snake_case,      snake_case)





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



          """simple docstring"""

          with tempfile.TemporaryDirectory() as tmpdirname:
             _lowercase              =CLIPConfig()

             # Create a dummy config file with image_proceesor_type
             _lowercase              =Path(snake_case) / 'preprocessor_config.json'
             _lowercase              =Path(snake_case) / 'config.json'
             json.dump(
                 {'image_processor_type': 'CLIPImageProcessor', 'processor_class': 'CLIPProcessor'},      open(snake_case,      'w'),      )
             json.dump({'model_type': 'clip'},      open(snake_case,      'w'))

             # remove image_processor_type to make sure config.json alone is enough to load image processor locally
             _lowercase              =AutoImageProcessor.from_pretrained(snake_case).to_dict()

             config_dict.pop('image_processor_type')
             _lowercase              =CLIPImageProcessor(**snake_case)

             # save in new folder
             model_config.save_pretrained(snake_case)
             config.save_pretrained(snake_case)

             _lowercase              =AutoImageProcessor.from_pretrained(snake_case)

             # make sure private variable is not incorrectly saved
             _lowercase              =json.loads(config.to_json_string())
             self.assertTrue('_processor_class' not in dict_as_saved)

          self.assertIsInstance(snake_case,      snake_case)





       def 				UpperCamelCase__						(				self	:Optional[int]):



          """simple docstring"""

          with tempfile.TemporaryDirectory() as tmpdirname:
             _lowercase              =Path(snake_case) / 'preprocessor_config.json'
             json.dump(
                 {'image_processor_type': 'CLIPImageProcessor', 'processor_class': 'CLIPProcessor'},      open(snake_case,      'w'),      )

             _lowercase              =AutoImageProcessor.from_pretrained(snake_case)
             self.assertIsInstance(snake_case,      snake_case)





       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

          with self.assertRaisesRegex(
              snake_case,      'clip-base is not a local folder and is not a valid model identifier'):
             _lowercase              =AutoImageProcessor.from_pretrained('clip-base')





       def 				UpperCamelCase__						(				self	:Dict):



          """simple docstring"""

          with self.assertRaisesRegex(
              snake_case,      r'aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'):
             _lowercase              =AutoImageProcessor.from_pretrained(snake_case,      revision='aaaaaa')





       def 				UpperCamelCase__						(				self	:List[Any]):



          """simple docstring"""

          with self.assertRaisesRegex(
              snake_case,      'hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json.',      ):
             _lowercase              =AutoImageProcessor.from_pretrained('hf-internal-testing/config-no-model')





       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

          with self.assertRaises(snake_case):
             _lowercase              =AutoImageProcessor.from_pretrained('hf-internal-testing/test_dynamic_image_processor')
          # If remote code is disabled, we can't load this config.
          with self.assertRaises(snake_case):
             _lowercase              =AutoImageProcessor.from_pretrained(
                 'hf-internal-testing/test_dynamic_image_processor',      trust_remote_code=snake_case)

          _lowercase              =AutoImageProcessor.from_pretrained(
              'hf-internal-testing/test_dynamic_image_processor',      trust_remote_code=snake_case)
          self.assertEqual(image_processor.__class__.__name__,      'NewImageProcessor')

          # Test image processor can be reloaded.
          with tempfile.TemporaryDirectory() as tmp_dir:
             image_processor.save_pretrained(snake_case)
             _lowercase              =AutoImageProcessor.from_pretrained(snake_case,      trust_remote_code=snake_case)
          self.assertEqual(reloaded_image_processor.__class__.__name__,      'NewImageProcessor')





       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

          try:
             AutoConfig.register('custom',      snake_case)
             AutoImageProcessor.register(snake_case,      snake_case)
             # Trying to register something existing in the Transformers library will raise an error
             with self.assertRaises(snake_case):
                AutoImageProcessor.register(snake_case,      snake_case)

             with tempfile.TemporaryDirectory() as tmpdirname:
                _lowercase              =Path(snake_case) / 'preprocessor_config.json'
                _lowercase              =Path(snake_case) / 'config.json'
                json.dump(
                    {'feature_extractor_type': 'CLIPFeatureExtractor', 'processor_class': 'CLIPProcessor'},      open(snake_case,      'w'),      )
                json.dump({'model_type': 'clip'},      open(snake_case,      'w'))

                _lowercase              =CustomImageProcessor.from_pretrained(snake_case)

             # Now that the config is registered, it can be used as any other config with the auto-API
             with tempfile.TemporaryDirectory() as tmp_dir:
                image_processor.save_pretrained(snake_case)
                _lowercase              =AutoImageProcessor.from_pretrained(snake_case)
                self.assertIsInstance(snake_case,      snake_case)

          finally:
             if "custom" in CONFIG_MAPPING._extra_content:
                del CONFIG_MAPPING._extra_content["custom"]
             if CustomConfig in IMAGE_PROCESSOR_MAPPING._extra_content:
                del IMAGE_PROCESSOR_MAPPING._extra_content[CustomConfig]





       def 				UpperCamelCase__						(				self	:str):



          """simple docstring"""

          class 					SCREAMING_SNAKE_CASE_  (       _a ):
                 """simple docstring"""


                 __lowerCAmelCase    :		Union[str, Any]					=True

          try:
             AutoConfig.register('custom',      snake_case)
             AutoImageProcessor.register(snake_case,      snake_case)
             # If remote code is not set, the default is to use local
             _lowercase              =AutoImageProcessor.from_pretrained('hf-internal-testing/test_dynamic_image_processor')
             self.assertEqual(image_processor.__class__.__name__,      'NewImageProcessor')
             self.assertTrue(image_processor.is_local)

             # If remote code is disabled, we load the local one.
             _lowercase              =AutoImageProcessor.from_pretrained(
                 'hf-internal-testing/test_dynamic_image_processor',      trust_remote_code=snake_case)
             self.assertEqual(image_processor.__class__.__name__,      'NewImageProcessor')
             self.assertTrue(image_processor.is_local)

             # If remote is enabled, we load from the Hub
             _lowercase              =AutoImageProcessor.from_pretrained(
                 'hf-internal-testing/test_dynamic_image_processor',      trust_remote_code=snake_case)
             self.assertEqual(image_processor.__class__.__name__,      'NewImageProcessor')
             self.assertTrue(not hasattr(snake_case,      'is_local'))

          finally:
             if "custom" in CONFIG_MAPPING._extra_content:
                del CONFIG_MAPPING._extra_content["custom"]
             if CustomConfig in IMAGE_PROCESSOR_MAPPING._extra_content:
                del IMAGE_PROCESSOR_MAPPING._extra_content[CustomConfig]


