


"""simple docstring"""






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

import transformers
from transformers import (
    CONFIG_MAPPING,
    FEATURE_EXTRACTOR_MAPPING,
    AutoConfig,
    AutoFeatureExtractor,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
)
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, get_tests_dir


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

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402


_lowerCamelCase     =  get_tests_dir('fixtures')
_lowerCamelCase     =  get_tests_dir('fixtures/dummy_feature_extractor_config.json')
_lowerCamelCase     =  get_tests_dir('fixtures/dummy-config.json')


class 			lowerCamelCase_		(			unittest.TestCase   ):





				"""simple docstring"""







				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				0



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base-960h"     )
								self.assertIsInstance(__lowerCAmelCase       ,		__lowerCAmelCase     )



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				AutoFeatureExtractor.from_pretrained(__lowerCAmelCase     )
								self.assertIsInstance(__lowerCAmelCase       ,		__lowerCAmelCase     )



				def   lowerCAmelCase__		(						self     ):






								with tempfile.TemporaryDirectory() as tmpdirname:
												SCREAMING_SNAKE_CASE__     							=				WavaVecaConfig()

												# remove feature_extractor_type to make sure config.json alone is enough to load feature processor locally
												SCREAMING_SNAKE_CASE__     							=				AutoFeatureExtractor.from_pretrained(__lowerCAmelCase     ).to_dict()

												config_dict.pop("feature_extractor_type"     )
												SCREAMING_SNAKE_CASE__     							=				WavaVecaFeatureExtractor(**__lowerCAmelCase     )

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

												SCREAMING_SNAKE_CASE__     							=				AutoFeatureExtractor.from_pretrained(__lowerCAmelCase     )

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

								self.assertIsInstance(__lowerCAmelCase       ,		__lowerCAmelCase     )



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				AutoFeatureExtractor.from_pretrained(__lowerCAmelCase     )
								self.assertIsInstance(__lowerCAmelCase       ,		__lowerCAmelCase     )



				def   lowerCAmelCase__		(						self     ):






								with self.assertRaisesRegex(
								    __lowerCAmelCase       ,		"bert-base is not a local folder and is not a valid model identifier"     ):
												SCREAMING_SNAKE_CASE__     							=				AutoFeatureExtractor.from_pretrained("bert-base"     )



				def   lowerCAmelCase__		(						self     ):






								with self.assertRaisesRegex(
								    __lowerCAmelCase       ,		R"aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"     ):
												SCREAMING_SNAKE_CASE__     							=				AutoFeatureExtractor.from_pretrained(__lowerCAmelCase       ,		revision="aaaaaa"     )



				def   lowerCAmelCase__		(						self     ):






								with self.assertRaisesRegex(
								    __lowerCAmelCase       ,		"hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json."       ,		):
												SCREAMING_SNAKE_CASE__     							=				AutoFeatureExtractor.from_pretrained("hf-internal-testing/config-no-model"     )



				def   lowerCAmelCase__		(						self     ):






								with self.assertRaises(__lowerCAmelCase     ):
												SCREAMING_SNAKE_CASE__     							=				AutoFeatureExtractor.from_pretrained(
												    "hf-internal-testing/test_dynamic_feature_extractor"     )
								# If remote code is disabled, we can't load this config.
								with self.assertRaises(__lowerCAmelCase     ):
												SCREAMING_SNAKE_CASE__     							=				AutoFeatureExtractor.from_pretrained(
												    "hf-internal-testing/test_dynamic_feature_extractor"       ,		trust_remote_code=__lowerCAmelCase     )

								SCREAMING_SNAKE_CASE__     							=				AutoFeatureExtractor.from_pretrained(
								    "hf-internal-testing/test_dynamic_feature_extractor"       ,		trust_remote_code=__lowerCAmelCase     )
								self.assertEqual(feature_extractor.__class__.__name__       ,		"NewFeatureExtractor"     )

								# Test feature extractor can be reloaded.
								with tempfile.TemporaryDirectory() as tmp_dir:
												feature_extractor.save_pretrained(__lowerCAmelCase     )
												SCREAMING_SNAKE_CASE__     							=				AutoFeatureExtractor.from_pretrained(__lowerCAmelCase       ,		trust_remote_code=__lowerCAmelCase     )
								self.assertEqual(reloaded_feature_extractor.__class__.__name__       ,		"NewFeatureExtractor"     )



				def   lowerCAmelCase__		(						self     ):






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

												# Now that the config is registered, it can be used as any other config with the auto-API
												SCREAMING_SNAKE_CASE__     							=				CustomFeatureExtractor.from_pretrained(__lowerCAmelCase     )
												with tempfile.TemporaryDirectory() as tmp_dir:
																feature_extractor.save_pretrained(__lowerCAmelCase     )
																SCREAMING_SNAKE_CASE__     							=				AutoFeatureExtractor.from_pretrained(__lowerCAmelCase     )
																self.assertIsInstance(__lowerCAmelCase       ,		__lowerCAmelCase     )

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






				def   lowerCAmelCase__		(						self     ):


								class 			lowerCamelCase_		(			UpperCAmelCase__   ):





												"""simple docstring"""







												_lowerCAmelCase      :      List[Any]      							=					True

								try:
												AutoConfig.register("custom"       ,		__lowerCAmelCase     )
												AutoFeatureExtractor.register(__lowerCAmelCase       ,		__lowerCAmelCase     )
												# If remote code is not set, the default is to use local
												SCREAMING_SNAKE_CASE__     							=				AutoFeatureExtractor.from_pretrained(
												    "hf-internal-testing/test_dynamic_feature_extractor"     )
												self.assertEqual(feature_extractor.__class__.__name__       ,		"NewFeatureExtractor"     )
												self.assertTrue(feature_extractor.is_local     )

												# If remote code is disabled, we load the local one.
												SCREAMING_SNAKE_CASE__     							=				AutoFeatureExtractor.from_pretrained(
												    "hf-internal-testing/test_dynamic_feature_extractor"       ,		trust_remote_code=__lowerCAmelCase     )
												self.assertEqual(feature_extractor.__class__.__name__       ,		"NewFeatureExtractor"     )
												self.assertTrue(feature_extractor.is_local     )

												# If remote is enabled, we load from the Hub
												SCREAMING_SNAKE_CASE__     							=				AutoFeatureExtractor.from_pretrained(
												    "hf-internal-testing/test_dynamic_feature_extractor"       ,		trust_remote_code=__lowerCAmelCase     )
												self.assertEqual(feature_extractor.__class__.__name__       ,		"NewFeatureExtractor"     )
												self.assertTrue(not hasattr(__lowerCAmelCase       ,		"is_local"     )     )

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








"""simple docstring"""






import numpy as np
import torch
from imwatermark import WatermarkEncoder


# Copied from https://github.com/Stability-AI/generative-models/blob/613af104c6b85184091d42d374fef420eddb356d/scripts/demo/streamlit_helpers.py#L66
_lowerCamelCase     =  0b1011_0011_1110_1100_1001_0000_0111_1011_1011_0001_1001_1110
# bin(x)[2:] gives bits of x as str, use int to convert them to 0/1
_lowerCamelCase     =  [int(bit) for bit in bin(WATERMARK_MESSAGE)[2:]]


class 			lowerCamelCase_		:





				"""simple docstring"""







				def __init__(						self     ):
								SCREAMING_SNAKE_CASE__     							=				WATERMARK_BITS
								SCREAMING_SNAKE_CASE__     							=				WatermarkEncoder()

								self.encoder.set_watermark("bits"       ,		self.watermark     )






				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__     ):
								# can't encode images that are smaller than 256
								if images.shape[-1] < 256:
												return images

								SCREAMING_SNAKE_CASE__     							=				(255 * (images / 2 + 0.5)).cpu().permute(0       ,		2       ,		3       ,		1     ).float().numpy()

								SCREAMING_SNAKE_CASE__     							=				[self.encoder.encode(UpperCAmelCase__       ,		"dwtDct"     ) for image in images]

								SCREAMING_SNAKE_CASE__     							=				torch.from_numpy(np.array(UpperCAmelCase__     )     ).permute(0       ,		3       ,		1       ,		2     )

								SCREAMING_SNAKE_CASE__     							=				torch.clamp(2 * (images / 255 - 0.5)       ,		min=-1.0       ,		max=1.0     )
								return images
