





import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class 	UpperCamelCase_							(  UpperCamelCase__					,	UpperCamelCase__					,	unittest.TestCase     ):
							lowerCamelCase_							     =       StableDiffusionSAGPipeline
							lowerCamelCase_							     =       TEXT_TO_IMAGE_PARAMS
							lowerCamelCase_							     =       TEXT_TO_IMAGE_BATCH_PARAMS
							lowerCamelCase_							     =       TEXT_TO_IMAGE_IMAGE_PARAMS
							lowerCamelCase_							     =       TEXT_TO_IMAGE_IMAGE_PARAMS
							lowerCamelCase_							     =       False


							def   _snake_case	(				self     :List[Any]		) -> str:



									"""simple docstring"""



									torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   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				,					)
									SCREAMING_SNAKE_CASE__								=   DDIMScheduler(
									    beta_start=0.0_0_0_8_5				,					beta_end=0.0_1_2				,					beta_schedule="""scaled_linear"""				,					clip_sample=__A				,					set_alpha_to_one=__A				,					)
									torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   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		)
									SCREAMING_SNAKE_CASE__								=   CLIPTextConfig(
									    bos_token_id=0				,					eos_token_id=2				,					hidden_size=32				,					intermediate_size=37				,					layer_norm_eps=1E-05				,					num_attention_heads=4				,					num_hidden_layers=5				,					pad_token_id=1				,					vocab_size=1000				,					)
									SCREAMING_SNAKE_CASE__								=   CLIPTextModel(__A		)
									SCREAMING_SNAKE_CASE__								=   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""		)

									SCREAMING_SNAKE_CASE__								=   {
									    """unet""": unet,
									    """scheduler""": scheduler,
									    """vae""": vae,
									    """text_encoder""": text_encoder,
									    """tokenizer""": tokenizer,
									    """safety_checker""": None,
									    """feature_extractor""": None,
									}
									return components


							def   _snake_case	(				self     :Optional[Any]				,					__A     :List[str]				,					__A     :Optional[int]=0		) -> Any:



									"""simple docstring"""



									if str(__A		).startswith("""mps"""		):
											SCREAMING_SNAKE_CASE__								=   torch.manual_seed(__A		)
									else:
											SCREAMING_SNAKE_CASE__								=   torch.Generator(device=__A		).manual_seed(__A		)
									SCREAMING_SNAKE_CASE__								=   {
									    """prompt""": """.""",
									    """generator""": generator,
									    """num_inference_steps""": 2,
									    """guidance_scale""": 1.0,
									    """sag_scale""": 1.0,
									    """output_type""": """numpy""",
									}
									return inputs


							def   _snake_case	(				self     :int		) -> int:



									"""simple docstring"""



									super().test_inference_batch_single_identical(expected_max_diff=3E-3		)




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


							def   _snake_case	(				self     :Union[str, Any]		) -> Union[str, Any]:



									"""simple docstring"""



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


							def   _snake_case	(				self     :Optional[Any]		) -> Tuple:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   StableDiffusionSAGPipeline.from_pretrained("""CompVis/stable-diffusion-v1-4"""		)
									SCREAMING_SNAKE_CASE__								=   sag_pipe.to(__A		)
									sag_pipe.set_progress_bar_config(disable=__A		)

									SCREAMING_SNAKE_CASE__								=   """."""
									SCREAMING_SNAKE_CASE__								=   torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   sag_pipe(
									    [prompt]				,					generator=__A				,					guidance_scale=7.5				,					sag_scale=1.0				,					num_inference_steps=20				,					output_type="""np"""		)

									SCREAMING_SNAKE_CASE__								=   output.images

									SCREAMING_SNAKE_CASE__								=   image[0, -3:, -3:, -1]

									assert image.shape == (1, 512, 512, 3)
									SCREAMING_SNAKE_CASE__								=   np.array([0.1_5_6_8, 0.1_7_3_8, 0.1_6_9_5, 0.1_6_9_3, 0.1_5_0_7, 0.1_7_0_5, 0.1_5_4_7, 0.1_7_5_1, 0.1_9_4_9]		)

									assert np.abs(image_slice.flatten() - expected_slice		).max() < 5E-2


							def   _snake_case	(				self     :Any		) -> Dict:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   StableDiffusionSAGPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base"""		)
									SCREAMING_SNAKE_CASE__								=   sag_pipe.to(__A		)
									sag_pipe.set_progress_bar_config(disable=__A		)

									SCREAMING_SNAKE_CASE__								=   """."""
									SCREAMING_SNAKE_CASE__								=   torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   sag_pipe(
									    [prompt]				,					generator=__A				,					guidance_scale=7.5				,					sag_scale=1.0				,					num_inference_steps=20				,					output_type="""np"""		)

									SCREAMING_SNAKE_CASE__								=   output.images

									SCREAMING_SNAKE_CASE__								=   image[0, -3:, -3:, -1]

									assert image.shape == (1, 512, 512, 3)
									SCREAMING_SNAKE_CASE__								=   np.array([0.3_4_5_9, 0.2_8_7_6, 0.2_5_3_7, 0.3_0_0_2, 0.2_6_7_1, 0.2_1_6_0, 0.3_0_2_6, 0.2_2_6_2, 0.2_3_7_1]		)

									assert np.abs(image_slice.flatten() - expected_slice		).max() < 5E-2


							def   _snake_case	(				self     :List[str]		) -> Tuple:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   StableDiffusionSAGPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base"""		)
									SCREAMING_SNAKE_CASE__								=   sag_pipe.to(__A		)
									sag_pipe.set_progress_bar_config(disable=__A		)

									SCREAMING_SNAKE_CASE__								=   """."""
									SCREAMING_SNAKE_CASE__								=   torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   sag_pipe(
									    [prompt]				,					width=768				,					height=512				,					generator=__A				,					guidance_scale=7.5				,					sag_scale=1.0				,					num_inference_steps=20				,					output_type="""np"""				,					)

									SCREAMING_SNAKE_CASE__								=   output.images

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






import argparse
import logging
import pickle
from collections import Counter


logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s', datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO
)
_lowerCamelCase 					=      logging.getLogger(__name__)

if __name__ == "__main__":
	_lowerCamelCase 					=      argparse.ArgumentParser(
	    description='Token Counts for smoothing the masking probabilities in MLM (cf XLM/word2vec)'
	)
	parser.add_argument(
	    '--data_file', type=str, default='data/dump.bert-base-uncased.pickle', help='The binarized dataset.'
	)
	parser.add_argument(
	    '--token_counts_dump', type=str, default='data/token_counts.bert-base-uncased.pickle', help='The dump file.'
	)
	parser.add_argument('--vocab_size', default=30522, type=int)
	_lowerCamelCase 					=      parser.parse_args()

	logger.info(F'''Loading data from {args.data_file}''')
	with open(args.data_file, 'rb') as fp:
		_lowerCamelCase 					=      pickle.load(fp)

	logger.info('Counting occurrences for MLM.')
	_lowerCamelCase 					=      Counter()
	for tk_ids in data:
		counter.update(tk_ids)
	_lowerCamelCase 					=      [0] * args.vocab_size
	for k, v in counter.items():
		_lowerCamelCase 					=      v

	logger.info(F'''Dump to {args.token_counts_dump}''')
	with open(args.token_counts_dump, 'wb') as handle:
		pickle.dump(counts, handle, protocol=pickle.HIGHEST_PROTOCOL)