






"""simple docstring"""




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 	__UpperCAmelCase		(				_A,  unittest.TestCase	):


	'''simple docstring'''





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


	@property
	def 				UpperCamelCase_  (     self					):



		'''simple docstring'''

		return 3_2


	@property
	def 				UpperCamelCase_  (     self					):



		'''simple docstring'''

		return 3_2


	@property
	def 				UpperCamelCase_  (     self					):



		'''simple docstring'''

		return self.time_input_dim


	@property
	def 				UpperCamelCase_  (     self					):



		'''simple docstring'''

		return self.time_input_dim * 4


	@property
	def 				UpperCamelCase_  (     self					):



		'''simple docstring'''

		return 1_0_0


	@property
	def 				UpperCamelCase_  (     self					):



		'''simple docstring'''

		torch.manual_seed(0					)

		_SCREAMING_SNAKE_CASE						={
		    '''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,
		}

		_SCREAMING_SNAKE_CASE						=UNetaDConditionModel(**__lowerCAmelCase					)
		return model


	@property
	def 				UpperCamelCase_  (     self					):



		'''simple docstring'''

		return {
		    "block_out_channels": [3_2, 3_2, 6_4, 6_4],
		    "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": 1_2,
		    "out_channels": 3,
		    "up_block_types": ["AttnUpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"],
		    "vq_embed_dim": 4,
		}


	@property
	def 				UpperCamelCase_  (     self					):



		'''simple docstring'''

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


	def 				UpperCamelCase_  (     self					):



		'''simple docstring'''

		_SCREAMING_SNAKE_CASE						=self.dummy_unet
		_SCREAMING_SNAKE_CASE						=self.dummy_movq

		_SCREAMING_SNAKE_CASE						=DDIMScheduler(
		    num_train_timesteps=1_0_0_0      ,      beta_schedule='''linear'''      ,      beta_start=0.00085      ,      beta_end=0.012      ,      clip_sample=__lowerCAmelCase      ,      set_alpha_to_one=__lowerCAmelCase      ,      steps_offset=1      ,      prediction_type='''epsilon'''      ,      thresholding=__lowerCAmelCase      ,      )

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


	def 				UpperCamelCase_  (     self      ,      _A      ,      _A=0					):



		'''simple docstring'''

		_SCREAMING_SNAKE_CASE						=floats_tensor((1, self.text_embedder_hidden_size)      ,      rng=random.Random(__lowerCAmelCase					)					).to(__lowerCAmelCase					)
		_SCREAMING_SNAKE_CASE						=floats_tensor((1, self.text_embedder_hidden_size)      ,      rng=random.Random(seed + 1					)					).to(
		    __lowerCAmelCase					)

		# create hint
		_SCREAMING_SNAKE_CASE						=floats_tensor((1, 3, 6_4, 6_4)      ,      rng=random.Random(__lowerCAmelCase					)					).to(__lowerCAmelCase					)

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







	def 				UpperCamelCase_  (     self					):



		'''simple docstring'''

		_SCREAMING_SNAKE_CASE						='''cpu'''

		_SCREAMING_SNAKE_CASE						=self.get_dummy_components()

		_SCREAMING_SNAKE_CASE						=self.pipeline_class(**__lowerCAmelCase					)
		_SCREAMING_SNAKE_CASE						=pipe.to(__lowerCAmelCase					)

		pipe.set_progress_bar_config(disable=__lowerCAmelCase					)

		_SCREAMING_SNAKE_CASE						=pipe(**self.get_dummy_inputs(__lowerCAmelCase					)					)
		_SCREAMING_SNAKE_CASE						=output.images

		_SCREAMING_SNAKE_CASE						=pipe(
		    **self.get_dummy_inputs(__lowerCAmelCase					)      ,      return_dict=__lowerCAmelCase      ,      )[0]

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

		assert image.shape == (1, 6_4, 6_4, 3)

		_SCREAMING_SNAKE_CASE						=np.array(
		    [0.6959826, 0.868279, 0.7558092, 0.68769467, 0.85805804, 0.65977496, 0.44885302, 0.5959111, 0.4251595]					)

		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 	__UpperCAmelCase		(				unittest.TestCase	):


	'''simple docstring'''





	def 				UpperCamelCase_  (     self					):



		'''simple docstring'''

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







	def 				UpperCamelCase_  (     self					):



		'''simple docstring'''

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

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

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

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

		_SCREAMING_SNAKE_CASE						='''A robot, 4k photo'''

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

		_SCREAMING_SNAKE_CASE						=torch.Generator(device='''cuda'''					).manual_seed(0					)
		_SCREAMING_SNAKE_CASE						=pipeline(
		    image_embeds=__lowerCAmelCase      ,      negative_image_embeds=__lowerCAmelCase      ,      hint=__lowerCAmelCase      ,      generator=__lowerCAmelCase      ,      num_inference_steps=1_0_0      ,      output_type='''np'''      ,      )

		_SCREAMING_SNAKE_CASE						=output.images[0]

		assert image.shape == (5_1_2, 5_1_2, 3)

		assert_mean_pixel_difference(__lowerCAmelCase      ,      __lowerCAmelCase					)









import collections
import os
import re
from pathlib import Path


UpperCAmelCase_		 =					"""src/transformers"""


# Matches is_xxx_available()
UpperCAmelCase_		 =					re.compile(r"""is\_([a-z_]*)_available()""")
# Catches a one-line _import_struct = {xxx}
UpperCAmelCase_		 =					re.compile(r"""^_import_structure\s+=\s+\{([^\}]+)\}""")
# Catches a line with a key-values pattern: "bla": ["foo", "bar"]
UpperCAmelCase_		 =					re.compile(r"""\s+\"\S*\":\s+\[([^\]]*)\]""")
# Catches a line if not is_foo_available
UpperCAmelCase_		 =					re.compile(r"""^\s*if\s+not\s+is\_[a-z_]*\_available\(\)""")
# Catches a line _import_struct["bla"].append("foo")
UpperCAmelCase_		 =					re.compile(r"""^\s*_import_structure\[\"\S*\"\]\.append\(\"(\S*)\"\)""")
# Catches a line _import_struct["bla"].extend(["foo", "bar"]) or _import_struct["bla"] = ["foo", "bar"]
UpperCAmelCase_		 =					re.compile(r"""^\s*_import_structure\[\S*\](?:\.extend\(|\s*=\s+)\[([^\]]*)\]""")
# Catches a line with an object between quotes and a comma:     "MyModel",
UpperCAmelCase_		 =					re.compile(r"""^\s+\"([^\"]+)\",""")
# Catches a line with objects between brackets only:    ["foo", "bar"],
UpperCAmelCase_		 =					re.compile(r"""^\s+\[([^\]]+)\]""")
# Catches a line with from foo import bar, bla, boo
UpperCAmelCase_		 =					re.compile(r"""\s+from\s+\S*\s+import\s+([^\(\s].*)\n""")
# Catches a line with try:
UpperCAmelCase_		 =					re.compile(r"""^\s*try:""")
# Catches a line with else:
UpperCAmelCase_		 =					re.compile(r"""^\s*else:""")

def 	SCREAMING_SNAKE_CASE_							(	_snake_case	:Optional[int]				)  ->     Any:
   if _re_test_backend.search(_snake_case				) is None:
      return None
   _A      =       [b[0] for b in _re_backend.findall(_snake_case				)]
   backends.sort()
   return "_and_".join(_snake_case				)

def 	SCREAMING_SNAKE_CASE_							(	_snake_case	:Any				)  ->     Any:
   with open(_snake_case	,   '''r'''	,   encoding='''utf-8'''	,   newline='''\n'''				) as f:
      _A      =       f.readlines()

   _A      =       0
   while line_index < len(_snake_case				) and not lines[line_index].startswith('''_import_structure = {'''				):
      line_index += 1

   # If this is a traditional init, just return.
   if line_index >= len(_snake_case				):
      return None

   # First grab the objects without a specific backend in _import_structure
   _A      =       []
   while not lines[line_index].startswith('''if TYPE_CHECKING'''				) and find_backend(lines[line_index]				) is None:
      _A      =       lines[line_index]
      # If we have everything on a single line, let's deal with it.
      if _re_one_line_import_struct.search(_snake_case				):
         _A      =       _re_one_line_import_struct.search(_snake_case				).groups()[0]
         _A      =       re.findall(r'''\[([^\]]+)\]'''	,   _snake_case				)
         for imp in imports:
            objects.extend([obj[1:-1] for obj in imp.split(''', '''				)]				)
         line_index += 1
         continue
      _A      =       _re_import_struct_key_value.search(_snake_case				)
      if single_line_import_search is not None:
         _A      =       [obj[1:-1] for obj in single_line_import_search.groups()[0].split(''', '''				) if len(_snake_case				) > 0]
         objects.extend(_snake_case				)
      elif line.startswith(''' ''' * 8 + '''"'''				):
         objects.append(line[9:-3]				)
      line_index += 1

   _A      =       {'''none''': objects}
   # Let's continue with backend-specific objects in _import_structure
   while not lines[line_index].startswith('''if TYPE_CHECKING'''				):
      # If the line is an if not is_backend_available, we grab all objects associated.
      _A      =       find_backend(lines[line_index]				)
      # Check if the backend declaration is inside a try block:
      if _re_try.search(lines[line_index - 1]				) is None:
         _A      =       None

      if backend is not None:
         line_index += 1

         # Scroll until we hit the else block of try-except-else
         while _re_else.search(lines[line_index]				) is None:
            line_index += 1

         line_index += 1

         _A      =       []
         # Until we unindent, add backend objects to the list
         while len(lines[line_index]				) <= 1 or lines[line_index].startswith(''' ''' * 4				):
            _A      =       lines[line_index]
            if _re_import_struct_add_one.search(_snake_case				) is not None:
               objects.append(_re_import_struct_add_one.search(_snake_case				).groups()[0]				)
            elif _re_import_struct_add_many.search(_snake_case				) is not None:
               _A      =       _re_import_struct_add_many.search(_snake_case				).groups()[0].split(''', '''				)
               _A      =       [obj[1:-1] for obj in imports if len(_snake_case				) > 0]
               objects.extend(_snake_case				)
            elif _re_between_brackets.search(_snake_case				) is not None:
               _A      =       _re_between_brackets.search(_snake_case				).groups()[0].split(''', '''				)
               _A      =       [obj[1:-1] for obj in imports if len(_snake_case				) > 0]
               objects.extend(_snake_case				)
            elif _re_quote_object.search(_snake_case				) is not None:
               objects.append(_re_quote_object.search(_snake_case				).groups()[0]				)
            elif line.startswith(''' ''' * 8 + '''"'''				):
               objects.append(line[9:-3]				)
            elif line.startswith(''' ''' * 12 + '''"'''				):
               objects.append(line[13:-3]				)
            line_index += 1

         _A      =       objects
      else:
         line_index += 1

    # At this stage we are in the TYPE_CHECKING part, first grab the objects without a specific backend
   _A      =       []
   while (
       line_index < len(_snake_case				)
       and find_backend(lines[line_index]				) is None
       and not lines[line_index].startswith('''else'''				)
   ):
      _A      =       lines[line_index]
      _A      =       _re_import.search(_snake_case				)
      if single_line_import_search is not None:
         objects.extend(single_line_import_search.groups()[0].split(''', '''				)				)
      elif line.startswith(''' ''' * 8				):
         objects.append(line[8:-2]				)
      line_index += 1

   _A      =       {'''none''': objects}
   # Let's continue with backend-specific objects
   while line_index < len(_snake_case				):
      # If the line is an if is_backend_available, we grab all objects associated.
      _A      =       find_backend(lines[line_index]				)
      # Check if the backend declaration is inside a try block:
      if _re_try.search(lines[line_index - 1]				) is None:
         _A      =       None

      if backend is not None:
         line_index += 1

         # Scroll until we hit the else block of try-except-else
         while _re_else.search(lines[line_index]				) is None:
            line_index += 1

         line_index += 1

         _A      =       []
         # Until we unindent, add backend objects to the list
         while len(lines[line_index]				) <= 1 or lines[line_index].startswith(''' ''' * 8				):
            _A      =       lines[line_index]
            _A      =       _re_import.search(_snake_case				)
            if single_line_import_search is not None:
               objects.extend(single_line_import_search.groups()[0].split(''', '''				)				)
            elif line.startswith(''' ''' * 12				):
               objects.append(line[12:-2]				)
            line_index += 1

         _A      =       objects
      else:
         line_index += 1

   return import_dict_objects, type_hint_objects

def 	SCREAMING_SNAKE_CASE_							(	_snake_case	:List[str]	,   _snake_case	:Dict				)  ->     Any:

   def find_duplicates(_snake_case	:Any				):
      return [k for k, v in collections.Counter(_snake_case				).items() if v > 1]

   if list(import_dict_objects.keys()				) != list(type_hint_objects.keys()				):
      return ["Both sides of the init do not have the same backends!"]

   _A      =       []
   for key in import_dict_objects.keys():
      _A      =       find_duplicates(import_dict_objects[key]				)
      if duplicate_imports:
         errors.append(F'''Duplicate _import_structure definitions for: {duplicate_imports}'''				)
      _A      =       find_duplicates(type_hint_objects[key]				)
      if duplicate_type_hints:
         errors.append(F'''Duplicate TYPE_CHECKING objects for: {duplicate_type_hints}'''				)

      if sorted(set(import_dict_objects[key]				)				) != sorted(set(type_hint_objects[key]				)				):
         _A      =       '''base imports''' if key == '''none''' else F'''{key} backend'''
         errors.append(F'''Differences for {name}:'''				)
         for a in type_hint_objects[key]:
            if a not in import_dict_objects[key]:
               errors.append(F'''  {a} in TYPE_HINT but not in _import_structure.'''				)
         for a in import_dict_objects[key]:
            if a not in type_hint_objects[key]:
               errors.append(F'''  {a} in _import_structure but not in TYPE_HINT.'''				)
   return errors

def 	SCREAMING_SNAKE_CASE_							(	)  ->     int:
   _A      =       []
   for root, _, files in os.walk(_snake_case				):
      if "__init__.py" in files:
         _A      =       os.path.join(_snake_case	,   '''__init__.py'''				)
         _A      =       parse_init(_snake_case				)
         if objects is not None:
            _A      =       analyze_results(*_snake_case				)
            if len(_snake_case				) > 0:
               _A      =       F'''Problem in {fname}, both halves do not define the same objects.\n{errors[0]}'''
               failures.append('''\n'''.join(_snake_case				)				)
   if len(_snake_case				) > 0:
      raise ValueError('''\n\n'''.join(_snake_case				)				)

def 	SCREAMING_SNAKE_CASE_							(	)  ->     Optional[int]:
   _A      =       []
   for path, directories, files in os.walk(_snake_case				):
      for folder in directories:
         # Ignore private modules
         if folder.startswith('''_'''				):
            directories.remove(_snake_case				)
            continue
         # Ignore leftovers from branches (empty folders apart from pycache)
         if len(list((Path(_snake_case				) / folder).glob('''*.py'''				)				)				) == 0:
            continue
         _A      =       str((Path(_snake_case				) / folder).relative_to(_snake_case				)				)
         _A      =       short_path.replace(os.path.sep	,   '''.'''				)
         submodules.append(_snake_case				)
      for fname in files:
         if fname == "__init__.py":
            continue
         _A      =       str((Path(_snake_case				) / fname).relative_to(_snake_case				)				)
         _A      =       short_path.replace('''.py'''	,   ''''''				).replace(os.path.sep	,   '''.'''				)
         if len(submodule.split('''.'''				)				) == 1:
            submodules.append(_snake_case				)
   return submodules


UpperCAmelCase_		 =					[
    """convert_pytorch_checkpoint_to_tf2""",
    """modeling_flax_pytorch_utils""",
    """models.esm.openfold_utils""",
]

def 	SCREAMING_SNAKE_CASE_							(	)  ->     List[str]:
   # This is to make sure the transformers module imported is the one in the repo.
   from transformers.utils import direct_transformers_import

   _A      =       direct_transformers_import(_snake_case				)

   _A      =       set(transformers._import_structure.keys()				)
   # This contains all the base keys of the _import_structure object defined in the init, but if the user is missing
   # some optional dependencies, they may not have all of them. Thus we read the init to read all additions and
   # (potentiall re-) add them.
   with open(os.path.join(_snake_case	,   '''__init__.py'''				)	,   '''r'''				) as f:
      _A      =       f.read()
   import_structure_keys.update(set(re.findall(r'''import_structure\[\"([^\"]*)\"\]'''	,   _snake_case				)				)				)

   _A      =       [
       module
       for module in get_transformers_submodules()
       if module not in IGNORE_SUBMODULES and module not in import_structure_keys
   ]

   if len(_snake_case				) > 0:
      _A      =       '''\n'''.join(F'''- {module}''' for module in module_not_registered				)
      raise ValueError(
          '''The following submodules are not properly registed in the main init of Transformers:\n'''
          F'''{list_of_modules}\n'''
          '''Make sure they appear somewhere in the keys of `_import_structure` with an empty list as value.'''				)


if __name__ == "__main__":
   check_all_inits()
   check_submodules()



