



"""simple docstring"""






import argparse
import logging
import os
import re

import tensorflow as tf

from transformers import (
    AutoConfig,
    AutoTokenizer,
    DataCollatorForLanguageModeling,
    PushToHubCallback,
    TFAutoModelForMaskedLM,
    create_optimizer,
)


_lowercase				:				Any   		=		logging.getLogger(__name__)

_lowercase				:				int   		=		tf.data.AUTOTUNE






def   snake_case__     (	):
			"""simple docstring"""
			lowerCamelCase__	:	int				=argparse.ArgumentParser(description='''Train a masked language model on TPU.'''			)
			parser.add_argument(
			    '''--pretrained_model_config'''							,	type=__lowerCamelCase							,	default='''roberta-base'''							,	help='''The model config to use. Note that we don\'t copy the model\'s weights, only the config!'''							,	)
			parser.add_argument(
			    '''--tokenizer'''							,	type=__lowerCamelCase							,	default='''unigram-tokenizer-wikitext'''							,	help='''The name of the tokenizer to load. We use the pretrained tokenizer to initialize the model\'s vocab size.'''							,	)

			parser.add_argument(
			    '''--per_replica_batch_size'''							,	type=__lowerCamelCase							,	default=8							,	help='''Batch size per TPU core.'''							,	)

			parser.add_argument(
			    '''--no_tpu'''							,	action='''store_true'''							,	help='''If set, run on CPU and don\'t try to initialize a TPU. Useful for debugging on non-TPU instances.'''							,	)

			parser.add_argument(
			    '''--tpu_name'''							,	type=__lowerCamelCase							,	help='''Name of TPU resource to initialize. Should be blank on Colab, and \'local\' on TPU VMs.'''							,	default='''local'''							,	)

			parser.add_argument(
			    '''--tpu_zone'''							,	type=__lowerCamelCase							,	help='''Google cloud zone that TPU resource is located in. Only used for non-Colab TPU nodes.'''							,	)

			parser.add_argument(
			    '''--gcp_project'''							,	type=__lowerCamelCase							,	help='''Google cloud project name. Only used for non-Colab TPU nodes.'''			)

			parser.add_argument(
			    '''--bfloat16'''							,	action='''store_true'''							,	help='''Use mixed-precision bfloat16 for training. This is the recommended lower-precision format for TPU.'''							,	)

			parser.add_argument(
			    '''--train_dataset'''							,	type=__lowerCamelCase							,	help='''Path to training dataset to load. If the path begins with `gs://`'''
			    ''' then the dataset will be loaded from a Google Cloud Storage bucket.'''							,	)

			parser.add_argument(
			    '''--shuffle_buffer_size'''							,	type=__lowerCamelCase							,	default=2**18							,	help='''Size of the shuffle buffer (in samples)'''							,	)

			parser.add_argument(
			    '''--eval_dataset'''							,	type=__lowerCamelCase							,	help='''Path to evaluation dataset to load. If the path begins with `gs://`'''
			    ''' then the dataset will be loaded from a Google Cloud Storage bucket.'''							,	)

			parser.add_argument(
			    '''--num_epochs'''							,	type=__lowerCamelCase							,	default=1							,	help='''Number of epochs to train for.'''							,	)

			parser.add_argument(
			    '''--learning_rate'''							,	type=__lowerCamelCase							,	default=1e-4							,	help='''Learning rate to use for training.'''							,	)

			parser.add_argument(
			    '''--weight_decay_rate'''							,	type=__lowerCamelCase							,	default=1e-3							,	help='''Weight decay rate to use for training.'''							,	)

			parser.add_argument(
			    '''--max_length'''							,	type=__lowerCamelCase							,	default=512							,	help='''Maximum length of tokenized sequences. Should match the setting used in prepare_tfrecord_shards.py'''							,	)

			parser.add_argument(
			    '''--mlm_probability'''							,	type=__lowerCamelCase							,	default=0.15							,	help='''Fraction of tokens to mask during training.'''							,	)

			parser.add_argument('''--output_dir'''							,	type=__lowerCamelCase							,	required=__lowerCamelCase							,	help='''Path to save model checkpoints to.'''			)
			parser.add_argument('''--hub_model_id'''							,	type=__lowerCamelCase							,	help='''Model ID to upload to on the Hugging Face Hub.'''			)

			lowerCamelCase__	:	Tuple				=parser.parse_args()
			return args






def   snake_case__     (	__lowerCamelCase       :					Union[str, Any]			):
			"""simple docstring"""
			try:
						if args.tpu_name:
									lowerCamelCase__	:	int				=tf.distribute.cluster_resolver.TPUClusterResolver(
									    args.tpu_name							,	zone=args.tpu_zone							,	project=args.gcp_project			)
						else:
									lowerCamelCase__	:	Any				=tf.distribute.cluster_resolver.TPUClusterResolver()
			except ValueError:
						raise RuntimeError(
						    '''Couldn\'t connect to TPU! Most likely you need to specify --tpu_name, --tpu_zone, or '''
						    '''--gcp_project. When running on a TPU VM, use --tpu_name local.'''			)

			tf.config.experimental_connect_to_cluster(__lowerCamelCase			)
			tf.tpu.experimental.initialize_tpu_system(__lowerCamelCase			)

			return tpu






def   snake_case__     (	__lowerCamelCase       :					int			):
			"""simple docstring"""
			lowerCamelCase__	:	Optional[Any]				=0
			for file in file_list:
						lowerCamelCase__	:	Tuple				=file.split('''/'''			)[-1]
						lowerCamelCase__	:	Dict				=re.search(R'''-\d+-(\d+)\.tfrecord'''							,	__lowerCamelCase			).group(1			)
						lowerCamelCase__	:	Union[str, Any]				=int(__lowerCamelCase			)
						num_samples += sample_count

			return num_samples






def   snake_case__     (	__lowerCamelCase       :					int							,	__lowerCamelCase       :					List[str]							,	__lowerCamelCase       :					Optional[int]							,	__lowerCamelCase       :					List[Any]							,	__lowerCamelCase       :					Union[str, Any]							,	__lowerCamelCase       :					Optional[Any]=None			):
			"""simple docstring"""
			lowerCamelCase__	:	Optional[int]				=count_samples(__lowerCamelCase			)
			lowerCamelCase__	:	Dict				=tf.data.Dataset.from_tensor_slices(__lowerCamelCase			)
			if shuffle:
						lowerCamelCase__	:	List[str]				=dataset.shuffle(len(__lowerCamelCase			)			)
			lowerCamelCase__	:	Dict				=tf.data.TFRecordDataset(__lowerCamelCase							,	num_parallel_reads=__lowerCamelCase			)
			# TF can't infer the total sample count because it doesn't read all the records yet, so we assert it here
			lowerCamelCase__	:	Optional[int]				=dataset.apply(tf.data.experimental.assert_cardinality(__lowerCamelCase			)			)
			lowerCamelCase__	:	str				=dataset.map(__lowerCamelCase							,	num_parallel_calls=__lowerCamelCase			)
			if shuffle:
						assert shuffle_buffer_size is not None
						lowerCamelCase__	:	str				=dataset.shuffle(args.shuffle_buffer_size			)
			lowerCamelCase__	:	Any				=dataset.batch(__lowerCamelCase							,	drop_remainder=__lowerCamelCase			)
			lowerCamelCase__	:	str				=dataset.map(__lowerCamelCase							,	num_parallel_calls=__lowerCamelCase			)
			lowerCamelCase__	:	int				=dataset.prefetch(__lowerCamelCase			)
			return dataset






def   snake_case__     (	__lowerCamelCase       :					Any			):
			"""simple docstring"""
			if not args.no_tpu:
						lowerCamelCase__	:	str				=initialize_tpu(__lowerCamelCase			)
						lowerCamelCase__	:	str				=tf.distribute.TPUStrategy(__lowerCamelCase			)
			else:
						lowerCamelCase__	:	List[str]				=tf.distribute.OneDeviceStrategy(device='''/gpu:0'''			)

			if args.bfloataa:
						tf.keras.mixed_precision.set_global_policy('''mixed_bfloat16'''			)

			lowerCamelCase__	:	Tuple				=AutoTokenizer.from_pretrained(args.tokenizer			)
			lowerCamelCase__	:	Any				=AutoConfig.from_pretrained(args.pretrained_model_config			)
			lowerCamelCase__	:	Tuple				=tokenizer.vocab_size

			lowerCamelCase__	:	Any				=tf.io.gfile.glob(os.path.join(args.train_dataset							,	'''*.tfrecord'''			)			)
			if not training_records:
						raise ValueError(f'''No .tfrecord files found in {args.train_dataset}.'''			)
			lowerCamelCase__	:	int				=tf.io.gfile.glob(os.path.join(args.eval_dataset							,	'''*.tfrecord'''			)			)
			if not eval_records:
						raise ValueError(f'''No .tfrecord files found in {args.eval_dataset}.'''			)

			lowerCamelCase__	:	Optional[int]				=count_samples(__lowerCamelCase			)

			lowerCamelCase__	:	Optional[int]				=num_train_samples // (args.per_replica_batch_size * strategy.num_replicas_in_sync)
			lowerCamelCase__	:	Any				=steps_per_epoch * args.num_epochs

			with strategy.scope():
						lowerCamelCase__	:	Union[str, Any]				=TFAutoModelForMaskedLM.from_config(__lowerCamelCase			)
						model(model.dummy_inputs			)  # Pass some dummy inputs through the model to ensure all the weights are built
						lowerCamelCase__			,     lowerCamelCase__	:	Dict				=create_optimizer(
						    num_train_steps=__lowerCamelCase							,	num_warmup_steps=total_train_steps // 20							,	init_lr=args.learning_rate							,	weight_decay_rate=args.weight_decay_rate							,	)

						# Transformers models compute the right loss for their task by default when labels are passed, and will
						# use this for training unless you specify your own loss function in compile().
						model.compile(optimizer=__lowerCamelCase							,	metrics=['''accuracy''']			)

			def decode_fn(__lowerCamelCase       :					Any			):
						lowerCamelCase__	:	Optional[Any]				={
						    '''input_ids''': tf.io.FixedLenFeature(dtype=tf.intaa							,	shape=(args.max_length,)			),
						    '''attention_mask''': tf.io.FixedLenFeature(dtype=tf.intaa							,	shape=(args.max_length,)			),
						}
						return tf.io.parse_single_example(__lowerCamelCase							,	__lowerCamelCase			)

			# Many of the data collators in Transformers are TF-compilable when return_tensors == "tf", so we can
			# use their methods in our data pipeline.
			lowerCamelCase__	:	Optional[int]				=DataCollatorForLanguageModeling(
			    tokenizer=__lowerCamelCase							,	mlm_probability=args.mlm_probability							,	mlm=__lowerCamelCase							,	return_tensors='''tf'''			)

			def mask_with_collator(__lowerCamelCase       :					Any			):
						# TF really needs an isin() function
						lowerCamelCase__	:	Union[str, Any]				=(
						    ~tf.cast(batch['''attention_mask''']							,	tf.bool			)
						    | (batch['''input_ids'''] == tokenizer.cls_token_id)
						    | (batch['''input_ids'''] == tokenizer.sep_token_id)
						)
						lowerCamelCase__			,     lowerCamelCase__	:	List[Any]				=data_collator.tf_mask_tokens(
						    batch['''input_ids''']							,	vocab_size=len(__lowerCamelCase			)							,	mask_token_id=tokenizer.mask_token_id							,	special_tokens_mask=__lowerCamelCase							,	)
						return batch

			lowerCamelCase__	:	Any				=args.per_replica_batch_size * strategy.num_replicas_in_sync

			lowerCamelCase__	:	Optional[Any]				=prepare_dataset(
			    __lowerCamelCase							,	decode_fn=__lowerCamelCase							,	mask_fn=__lowerCamelCase							,	batch_size=__lowerCamelCase							,	shuffle=__lowerCamelCase							,	shuffle_buffer_size=args.shuffle_buffer_size							,	)

			lowerCamelCase__	:	Any				=prepare_dataset(
			    __lowerCamelCase							,	decode_fn=__lowerCamelCase							,	mask_fn=__lowerCamelCase							,	batch_size=__lowerCamelCase							,	shuffle=__lowerCamelCase							,	)

			lowerCamelCase__	:	Any				=[]
			if args.hub_model_id:
						callbacks.append(
						    PushToHubCallback(output_dir=args.output_dir							,	hub_model_id=args.hub_model_id							,	tokenizer=__lowerCamelCase			)			)

			model.fit(
			    __lowerCamelCase							,	validation_data=__lowerCamelCase							,	epochs=args.num_epochs							,	callbacks=__lowerCamelCase							,	)

			model.save_pretrained(args.output_dir			)


if __name__ == "__main__":
		_lowercase				:				List[Any]   		=		parse_args()
		main(args)









"""simple docstring"""






import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModel,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableUnCLIPImgaImgPipeline, UNetaDConditionModel
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    floats_tensor,
    load_image,
    load_numpy,
    require_torch_gpu,
    skip_mps,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()

class  __SCREAMING_SNAKE_CASE			(   lowerCAmelCase_	,     lowerCAmelCase_	,     lowerCAmelCase_	,     unittest.TestCase     ):


					'''simple docstring'''





					_a         =							StableUnCLIPImgaImgPipeline
					_a         =							TEXT_GUIDED_IMAGE_VARIATION_PARAMS
					_a         =							TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
					_a         =							frozenset(
					    []     )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
					_a         =							frozenset([]     )



					def 	snake_case ( self							:						List[str]  )->      str:
								lowerCamelCase__	:	Dict				=32
								lowerCamelCase__	:	Optional[Any]				=embedder_hidden_size

								# image encoding components

								lowerCamelCase__	:	Dict				=CLIPImageProcessor(crop_size=32,						size=32  )

								torch.manual_seed(0  )
								lowerCamelCase__	:	List[Any]				=CLIPVisionModelWithProjection(
								    CLIPVisionConfig(
								        hidden_size=lowerCamelCase,						projection_dim=lowerCamelCase,						num_hidden_layers=5,						num_attention_heads=4,						image_size=32,						intermediate_size=37,						patch_size=1,						)  )

								# regular denoising components

								torch.manual_seed(0  )
								lowerCamelCase__	:	Optional[int]				=StableUnCLIPImageNormalizer(embedding_dim=lowerCamelCase  )
								lowerCamelCase__	:	Dict				=DDPMScheduler(beta_schedule='''squaredcos_cap_v2'''  )

								torch.manual_seed(0  )
								lowerCamelCase__	:	Optional[int]				=CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  )

								torch.manual_seed(0  )
								lowerCamelCase__	:	Tuple				=CLIPTextModel(
								    CLIPTextConfig(
								        bos_token_id=0,						eos_token_id=2,						hidden_size=lowerCamelCase,						projection_dim=32,						intermediate_size=37,						layer_norm_eps=1E-05,						num_attention_heads=4,						num_hidden_layers=5,						pad_token_id=1,						vocab_size=1000,						)  )

								torch.manual_seed(0  )
								lowerCamelCase__	:	Dict				=UNetaDConditionModel(
								    sample_size=32,						in_channels=4,						out_channels=4,						down_block_types=('''CrossAttnDownBlock2D''', '''DownBlock2D'''),						up_block_types=('''UpBlock2D''', '''CrossAttnUpBlock2D'''),						block_out_channels=(32, 64),						attention_head_dim=(2, 4),						class_embed_type='''projection''',						projection_class_embeddings_input_dim=embedder_projection_dim * 2,						cross_attention_dim=lowerCamelCase,						layers_per_block=1,						upcast_attention=lowerCamelCase,						use_linear_projection=lowerCamelCase,						)

								torch.manual_seed(0  )
								lowerCamelCase__	:	Union[str, Any]				=DDIMScheduler(
								    beta_schedule='''scaled_linear''',						beta_start=0.00_085,						beta_end=0.012,						prediction_type='''v_prediction''',						set_alpha_to_one=lowerCamelCase,						steps_offset=1,						)

								torch.manual_seed(0  )
								lowerCamelCase__	:	Optional[int]				=AutoencoderKL()

								lowerCamelCase__	:	int				={
								    # image encoding components
								    '''feature_extractor''': feature_extractor,
								    '''image_encoder''': image_encoder.eval(),
								    # image noising components
								    '''image_normalizer''': image_normalizer.eval(),
								    '''image_noising_scheduler''': image_noising_scheduler,
								    # regular denoising components
								    '''tokenizer''': tokenizer,
								    '''text_encoder''': text_encoder.eval(),
								    '''unet''': unet.eval(),
								    '''scheduler''': scheduler,
								    '''vae''': vae.eval(),
								}

								return components



					def 	snake_case ( self							:						str,						lowerCamelCase							:						Dict,						lowerCamelCase							:						Any=0,						lowerCamelCase							:						str=True  )->      List[str]:
								if str(lowerCamelCase  ).startswith('''mps'''  ):
											lowerCamelCase__	:	List[Any]				=torch.manual_seed(lowerCamelCase  )
								else:
											lowerCamelCase__	:	Any				=torch.Generator(device=lowerCamelCase  ).manual_seed(lowerCamelCase  )

								lowerCamelCase__	:	Dict				=floats_tensor((1, 3, 32, 32),						rng=random.Random(lowerCamelCase  )  ).to(lowerCamelCase  )

								if pil_image:
											lowerCamelCase__	:	int				=input_image * 0.5 + 0.5
											lowerCamelCase__	:	Dict				=input_image.clamp(0,						1  )
											lowerCamelCase__	:	List[str]				=input_image.cpu().permute(0,						2,						3,						1  ).float().numpy()
											lowerCamelCase__	:	Dict				=DiffusionPipeline.numpy_to_pil(lowerCamelCase  )[0]

								return {
								    "prompt": "An anime racoon running a marathon",
								    "image": input_image,
								    "generator": generator,
								    "num_inference_steps": 2,
								    "output_type": "np",
								}



					@skip_mps
					def 	snake_case ( self							:						List[str]  )->      Optional[Any]:
								lowerCamelCase__	:	Dict				='''cpu'''  # ensure determinism for the device-dependent torch.Generator
								lowerCamelCase__	:	str				=self.get_dummy_components()
								lowerCamelCase__	:	int				=StableUnCLIPImgaImgPipeline(**lowerCamelCase  )
								lowerCamelCase__	:	Any				=sd_pipe.to(lowerCamelCase  )
								sd_pipe.set_progress_bar_config(disable=lowerCamelCase  )

								lowerCamelCase__	:	Dict				=self.get_dummy_inputs(lowerCamelCase  )
								inputs.update({'''image_embeds''': None}  )
								lowerCamelCase__	:	Any				=sd_pipe(**lowerCamelCase  ).images
								lowerCamelCase__	:	List[Any]				=image[0, -3:, -3:, -1]

								assert image.shape == (1, 32, 32, 3)
								lowerCamelCase__	:	Union[str, Any]				=np.array([0.3_872, 0.7_224, 0.5_601, 0.4_741, 0.6_872, 0.5_814, 0.4_636, 0.3_867, 0.5_078]  )

								assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1E-3



					def 	snake_case ( self							:						int  )->      Tuple:
								lowerCamelCase__	:	Tuple				=torch_device in ['''cpu''', '''mps''']

								self._test_attention_slicing_forward_pass(test_max_difference=lowerCamelCase  )



					def 	snake_case ( self							:						int  )->      Optional[Any]:
								lowerCamelCase__	:	List[Any]				=torch_device in ['''cpu''', '''mps''']

								self._test_inference_batch_single_identical(test_max_difference=lowerCamelCase  )




					@unittest.skipIf(
					    torch_device != '''cuda''' or not is_xformers_available(),						reason='''XFormers attention is only available with CUDA and `xformers` installed''',						)
					def 	snake_case ( self							:						List[str]  )->      List[str]:
								self._test_xformers_attention_forwardGenerator_pass(test_max_difference=lowerCamelCase  )





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


					'''simple docstring'''





					def 	snake_case ( self							:						List[Any]  )->      Dict:
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()



					def 	snake_case ( self							:						Optional[int]  )->      int:
								lowerCamelCase__	:	Tuple				=load_image(
								    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'''  )

								lowerCamelCase__	:	Optional[int]				=load_numpy(
								    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_img2img_anime_turtle_fp16.npy'''  )

								lowerCamelCase__	:	Optional[Any]				=StableUnCLIPImgaImgPipeline.from_pretrained(
								    '''fusing/stable-unclip-2-1-l-img2img''',						torch_dtype=torch.floataa  )
								pipe.to(lowerCamelCase  )
								pipe.set_progress_bar_config(disable=lowerCamelCase  )
								# stable unclip will oom when integration tests are run on a V100,
								# so turn on memory savings
								pipe.enable_attention_slicing()
								pipe.enable_sequential_cpu_offload()

								lowerCamelCase__	:	int				=torch.Generator(device='''cpu'''  ).manual_seed(0  )
								lowerCamelCase__	:	Any				=pipe(lowerCamelCase,						'''anime turle''',						generator=lowerCamelCase,						output_type='''np'''  )

								lowerCamelCase__	:	List[Any]				=output.images[0]

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

								assert_mean_pixel_difference(lowerCamelCase,						lowerCamelCase  )



					def 	snake_case ( self							:						Optional[int]  )->      Tuple:
								lowerCamelCase__	:	Any				=load_image(
								    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'''  )

								lowerCamelCase__	:	str				=load_numpy(
								    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_h_img2img_anime_turtle_fp16.npy'''  )

								lowerCamelCase__	:	Optional[int]				=StableUnCLIPImgaImgPipeline.from_pretrained(
								    '''fusing/stable-unclip-2-1-h-img2img''',						torch_dtype=torch.floataa  )
								pipe.to(lowerCamelCase  )
								pipe.set_progress_bar_config(disable=lowerCamelCase  )
								# stable unclip will oom when integration tests are run on a V100,
								# so turn on memory savings
								pipe.enable_attention_slicing()
								pipe.enable_sequential_cpu_offload()

								lowerCamelCase__	:	str				=torch.Generator(device='''cpu'''  ).manual_seed(0  )
								lowerCamelCase__	:	Tuple				=pipe(lowerCamelCase,						'''anime turle''',						generator=lowerCamelCase,						output_type='''np'''  )

								lowerCamelCase__	:	Tuple				=output.images[0]

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

								assert_mean_pixel_difference(lowerCamelCase,						lowerCamelCase  )




					def 	snake_case ( self							:						Optional[int]  )->      List[str]:
								lowerCamelCase__	:	int				=load_image(
								    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'''  )

								torch.cuda.empty_cache()
								torch.cuda.reset_max_memory_allocated()
								torch.cuda.reset_peak_memory_stats()

								lowerCamelCase__	:	Any				=StableUnCLIPImgaImgPipeline.from_pretrained(
								    '''fusing/stable-unclip-2-1-h-img2img''',						torch_dtype=torch.floataa  )
								lowerCamelCase__	:	Optional[Any]				=pipe.to(lowerCamelCase  )
								pipe.set_progress_bar_config(disable=lowerCamelCase  )
								pipe.enable_attention_slicing()
								pipe.enable_sequential_cpu_offload()

								lowerCamelCase__	:	List[Any]				=pipe(
								    lowerCamelCase,						'''anime turtle''',						num_inference_steps=2,						output_type='''np''',						)

								lowerCamelCase__	:	Optional[int]				=torch.cuda.max_memory_allocated()
								# make sure that less than 7 GB is allocated
								assert mem_bytes < 7 * 10**9




