
"""simple docstring"""


import os
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import torch
from torch import nn

from ...models.controlnet import ControlNetModel, ControlNetOutput
from ...models.modeling_utils import ModelMixin
from ...utils import logging


_lowercase							      = logging.get_logger(__name__)







class  lowerCAmelCase_ (					UpperCamelCase_			):



						'''simple docstring'''


						def __init__(       self   :					List[Any]     ,A_   :					Union[List[ControlNetModel], Tuple[ControlNetModel]]					)   ->      str:
							super().__init__()
							A        =      nn.ModuleList(UpperCamelCase__					)


						def   _SCREAMING_SNAKE_CASE					(       self   :					Optional[Any]     ,A_   :					torch.FloatTensor     ,A_   :					Union[torch.Tensor, float, int]     ,A_   :					torch.Tensor     ,A_   :					List[torch.tensor]     ,A_   :					List[float]     ,A_   :					Optional[torch.Tensor] = None     ,A_   :					Optional[torch.Tensor] = None     ,A_   :					Optional[torch.Tensor] = None     ,A_   :					Optional[Dict[str, Any]] = None     ,A_   :					bool = False     ,A_   :					bool = True     ,)   ->      Dict:





							for i, (image, scale, controlnet) in enumerate(zip(UpperCamelCase__     ,UpperCamelCase__     ,self.nets					)					):
								A        =      controlnet(
								    UpperCamelCase__     ,UpperCamelCase__     ,UpperCamelCase__     ,UpperCamelCase__     ,UpperCamelCase__     ,UpperCamelCase__     ,UpperCamelCase__     ,UpperCamelCase__     ,UpperCamelCase__     ,UpperCamelCase__     ,UpperCamelCase__     ,)

								# merge samples
								if i == 0:
									A        =      down_samples, mid_sample
								else:
									A        =      [
									    samples_prev + samples_curr
									    for samples_prev, samples_curr in zip(UpperCamelCase__     ,UpperCamelCase__					)
									]
									mid_block_res_sample += mid_sample

							return down_block_res_samples, mid_block_res_sample


						def   _SCREAMING_SNAKE_CASE					(       self   :					Any     ,A_   :					Union[str, os.PathLike]     ,A_   :					bool = True     ,A_   :					Callable = None     ,A_   :					bool = False     ,A_   :					Optional[str] = None     ,)   ->      Union[str, Any]:
							A        =      0
							A        =      save_directory
							for controlnet in self.nets:
								controlnet.save_pretrained(
								    UpperCamelCase__     ,is_main_process=UpperCamelCase__     ,save_function=UpperCamelCase__     ,safe_serialization=UpperCamelCase__     ,variant=UpperCamelCase__     ,)

								idx += 1
								A        =      model_path_to_save + F'_{idx}'




						@classmethod
						def   _SCREAMING_SNAKE_CASE					(       cls   :					Dict     ,A_   :					Optional[Union[str, os.PathLike]]     ,**A_   :					Optional[int]					)   ->      str:
							A        =      0
							A        =      []

							# load controlnet and append to list until no controlnet directory exists anymore
							# first controlnet has to be saved under `./mydirectory/controlnet` to be compliant with `DiffusionPipeline.from_prertained`
							# second, third, ... controlnets have to be saved under `./mydirectory/controlnet_1`, `./mydirectory/controlnet_2`, ...
							A        =      pretrained_model_path
							while os.path.isdir(UpperCamelCase__					):
								A        =      ControlNetModel.from_pretrained(UpperCamelCase__     ,**UpperCamelCase__					)
								controlnets.append(UpperCamelCase__					)

								idx += 1
								A        =      pretrained_model_path + F'_{idx}'

							logger.info(F'{len(UpperCamelCase__					)} controlnets loaded from {pretrained_model_path}.'					)

							if len(UpperCamelCase__					) == 0:
								raise ValueError(
								    F'No ControlNets found under {os.path.dirname(UpperCamelCase__					)}. Expected at least {pretrained_model_path + "_0"}.'					)

							return cls(UpperCamelCase__					)
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 					lowercase__						(   unittest.TestCase):


						def        __A       (       self	:				List[Any]						):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:    Dict							=							0


						def        __A       (       self	:				Dict						):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:    str							=							AutoImageProcessor.from_pretrained('''openai/clip-vit-base-patch32'''						)
												self.assertIsInstance(UpperCamelCase__    ,		UpperCamelCase__						)


						def        __A       (       self	:				Dict						):





												'''simple docstring'''





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

																		SCREAMING_SNAKE_CASE							:    Any							=							AutoImageProcessor.from_pretrained(UpperCamelCase__						)
																		self.assertIsInstance(UpperCamelCase__    ,		UpperCamelCase__						)


						def        __A       (       self	:				int						):





												'''simple docstring'''





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

																		SCREAMING_SNAKE_CASE							:    Optional[Any]							=							AutoImageProcessor.from_pretrained(UpperCamelCase__						)
																		self.assertIsInstance(UpperCamelCase__    ,		UpperCamelCase__						)


						def        __A       (       self	:				List[str]						):





												'''simple docstring'''





												with tempfile.TemporaryDirectory() as tmpdirname:
																		SCREAMING_SNAKE_CASE							:    Dict							=							CLIPConfig()

																		# Create a dummy config file with image_proceesor_type
																		SCREAMING_SNAKE_CASE							:    Dict							=							Path(UpperCamelCase__						) / '''preprocessor_config.json'''
																		SCREAMING_SNAKE_CASE							:    List[str]							=							Path(UpperCamelCase__						) / '''config.json'''
																		json.dump(
																		    {'''image_processor_type''': '''CLIPImageProcessor''', '''processor_class''': '''CLIPProcessor'''}    ,		open(UpperCamelCase__    ,		'''w'''						)    ,		)
																		json.dump({'''model_type''': '''clip'''}    ,		open(UpperCamelCase__    ,		'''w'''						)						)

																		# remove image_processor_type to make sure config.json alone is enough to load image processor locally
																		SCREAMING_SNAKE_CASE							:    List[Any]							=							AutoImageProcessor.from_pretrained(UpperCamelCase__						).to_dict()

																		config_dict.pop('''image_processor_type'''						)
																		SCREAMING_SNAKE_CASE							:    str							=							CLIPImageProcessor(**UpperCamelCase__						)

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

																		SCREAMING_SNAKE_CASE							:    Union[str, Any]							=							AutoImageProcessor.from_pretrained(UpperCamelCase__						)

																		# make sure private variable is not incorrectly saved
																		SCREAMING_SNAKE_CASE							:    List[Any]							=							json.loads(config.to_json_string()						)
																		self.assertTrue('''_processor_class''' not in dict_as_saved						)

												self.assertIsInstance(UpperCamelCase__    ,		UpperCamelCase__						)


						def        __A       (       self	:				Tuple						):





												'''simple docstring'''





												with tempfile.TemporaryDirectory() as tmpdirname:
																		SCREAMING_SNAKE_CASE							:    Optional[int]							=							Path(UpperCamelCase__						) / '''preprocessor_config.json'''
																		json.dump(
																		    {'''image_processor_type''': '''CLIPImageProcessor''', '''processor_class''': '''CLIPProcessor'''}    ,		open(UpperCamelCase__    ,		'''w'''						)    ,		)

																		SCREAMING_SNAKE_CASE							:    Dict							=							AutoImageProcessor.from_pretrained(UpperCamelCase__						)
																		self.assertIsInstance(UpperCamelCase__    ,		UpperCamelCase__						)


						def        __A       (       self	:				int						):





												'''simple docstring'''





												with self.assertRaisesRegex(
												    UpperCamelCase__    ,		'''clip-base is not a local folder and is not a valid model identifier'''						):
																		SCREAMING_SNAKE_CASE							:    List[str]							=							AutoImageProcessor.from_pretrained('''clip-base'''						)


						def        __A       (       self	:				List[str]						):





												'''simple docstring'''





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


						def        __A       (       self	:				Dict						):





												'''simple docstring'''





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


						def        __A       (       self	:				List[Any]						):





												'''simple docstring'''





												with self.assertRaises(UpperCamelCase__						):
																		SCREAMING_SNAKE_CASE							:    str							=							AutoImageProcessor.from_pretrained('''hf-internal-testing/test_dynamic_image_processor'''						)
												# If remote code is disabled, we can't load this config.
												with self.assertRaises(UpperCamelCase__						):
																		SCREAMING_SNAKE_CASE							:    List[Any]							=							AutoImageProcessor.from_pretrained(
																		    '''hf-internal-testing/test_dynamic_image_processor'''    ,		trust_remote_code=UpperCamelCase__						)

												SCREAMING_SNAKE_CASE							:    List[Any]							=							AutoImageProcessor.from_pretrained(
												    '''hf-internal-testing/test_dynamic_image_processor'''    ,		trust_remote_code=UpperCamelCase__						)
												self.assertEqual(image_processor.__class__.__name__    ,		'''NewImageProcessor'''						)

												# Test image processor can be reloaded.
												with tempfile.TemporaryDirectory() as tmp_dir:
																		image_processor.save_pretrained(UpperCamelCase__						)
																		SCREAMING_SNAKE_CASE							:    Optional[int]							=							AutoImageProcessor.from_pretrained(UpperCamelCase__    ,		trust_remote_code=UpperCamelCase__						)
												self.assertEqual(reloaded_image_processor.__class__.__name__    ,		'''NewImageProcessor'''						)


						def        __A       (       self	:				Optional[Any]						):





												'''simple docstring'''





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

																		with tempfile.TemporaryDirectory() as tmpdirname:
																								SCREAMING_SNAKE_CASE							:    Union[str, Any]							=							Path(UpperCamelCase__						) / '''preprocessor_config.json'''
																								SCREAMING_SNAKE_CASE							:    Any							=							Path(UpperCamelCase__						) / '''config.json'''
																								json.dump(
																								    {'''feature_extractor_type''': '''CLIPFeatureExtractor''', '''processor_class''': '''CLIPProcessor'''}    ,		open(UpperCamelCase__    ,		'''w'''						)    ,		)
																								json.dump({'''model_type''': '''clip'''}    ,		open(UpperCamelCase__    ,		'''w'''						)						)

																								SCREAMING_SNAKE_CASE							:    Tuple							=							CustomImageProcessor.from_pretrained(UpperCamelCase__						)

																		# 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(UpperCamelCase__						)
																								SCREAMING_SNAKE_CASE							:    Any							=							AutoImageProcessor.from_pretrained(UpperCamelCase__						)
																								self.assertIsInstance(UpperCamelCase__    ,		UpperCamelCase__						)

												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        __A       (       self	:				Any						):





												'''simple docstring'''





												class 					lowercase__						(   UpperCamelCase_):
																		UpperCamelCase_     =		True

												try:
																		AutoConfig.register('''custom'''    ,		UpperCamelCase__						)
																		AutoImageProcessor.register(UpperCamelCase__    ,		UpperCamelCase__						)
																		# If remote code is not set, the default is to use local
																		SCREAMING_SNAKE_CASE							:    str							=							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.
																		SCREAMING_SNAKE_CASE							:    Union[str, Any]							=							AutoImageProcessor.from_pretrained(
																		    '''hf-internal-testing/test_dynamic_image_processor'''    ,		trust_remote_code=UpperCamelCase__						)
																		self.assertEqual(image_processor.__class__.__name__    ,		'''NewImageProcessor'''						)
																		self.assertTrue(image_processor.is_local						)

																		# If remote is enabled, we load from the Hub
																		SCREAMING_SNAKE_CASE							:    Optional[Any]							=							AutoImageProcessor.from_pretrained(
																		    '''hf-internal-testing/test_dynamic_image_processor'''    ,		trust_remote_code=UpperCamelCase__						)
																		self.assertEqual(image_processor.__class__.__name__    ,		'''NewImageProcessor'''						)
																		self.assertTrue(not hasattr(UpperCamelCase__    ,		'''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]





