
import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()





class   __UpperCamelCase      (				A__      ,		A__      ,		A__      ,		unittest.TestCase    ):






     """simple docstring"""







     lowerCAmelCase_									=       StableUnCLIPPipeline
     lowerCAmelCase_									=       TEXT_TO_IMAGE_PARAMS
     lowerCAmelCase_									=       TEXT_TO_IMAGE_BATCH_PARAMS
     lowerCAmelCase_									=       TEXT_TO_IMAGE_IMAGE_PARAMS
     lowerCAmelCase_									=       TEXT_TO_IMAGE_IMAGE_PARAMS

     # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
     lowerCAmelCase_									=       False



     def       UpperCAmelCase__				(							self				:				Optional[int]      ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE	:     Union[str, Any] =						32
       __SCREAMING_SNAKE_CASE	:     str =						embedder_hidden_size

       # prior components

       torch.manual_seed(0      )
       __SCREAMING_SNAKE_CASE	:     List[Any] =						CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''      )

       torch.manual_seed(0      )
       __SCREAMING_SNAKE_CASE	:     Dict =						CLIPTextModelWithProjection(
           CLIPTextConfig(
               bos_token_id=0			,		eos_token_id=2			,		hidden_size=_A			,		projection_dim=_A			,		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      )
       __SCREAMING_SNAKE_CASE	:     List[Any] =						PriorTransformer(
           num_attention_heads=2			,		attention_head_dim=12			,		embedding_dim=_A			,		num_layers=1			,		)

       torch.manual_seed(0      )
       __SCREAMING_SNAKE_CASE	:     Optional[int] =						DDPMScheduler(
           variance_type='''fixed_small_log'''			,		prediction_type='''sample'''			,		num_train_timesteps=1000			,		clip_sample=_A			,		clip_sample_range=5.0			,		beta_schedule='''squaredcos_cap_v2'''			,		)

       # regular denoising components

       torch.manual_seed(0      )
       __SCREAMING_SNAKE_CASE	:     List[Any] =						StableUnCLIPImageNormalizer(embedding_dim=_A      )
       __SCREAMING_SNAKE_CASE	:     str =						DDPMScheduler(beta_schedule='''squaredcos_cap_v2'''      )

       torch.manual_seed(0      )
       __SCREAMING_SNAKE_CASE	:     Any =						CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''      )

       torch.manual_seed(0      )
       __SCREAMING_SNAKE_CASE	:     Union[str, Any] =						CLIPTextModel(
           CLIPTextConfig(
               bos_token_id=0			,		eos_token_id=2			,		hidden_size=_A			,		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      )
       __SCREAMING_SNAKE_CASE	:     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=_A			,		layers_per_block=1			,		upcast_attention=_A			,		use_linear_projection=_A			,		)

       torch.manual_seed(0      )
       __SCREAMING_SNAKE_CASE	:     Any =						DDIMScheduler(
           beta_schedule='''scaled_linear'''			,		beta_start=0.0_00_85			,		beta_end=0.0_12			,		prediction_type='''v_prediction'''			,		set_alpha_to_one=_A			,		steps_offset=1			,		)

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

       __SCREAMING_SNAKE_CASE	:     List[str] =						{
           # prior components
           "prior_tokenizer": prior_tokenizer,
           "prior_text_encoder": prior_text_encoder,
           "prior": prior,
           "prior_scheduler": prior_scheduler,
           # image noising components
           "image_normalizer": image_normalizer,
           "image_noising_scheduler": image_noising_scheduler,
           # regular denoising components
           "tokenizer": tokenizer,
           "text_encoder": text_encoder,
           "unet": unet,
           "scheduler": scheduler,
           "vae": vae,
       }

       return components



     def       UpperCAmelCase__				(							self				:				Optional[int]			,		_A				:				int			,		_A				:				Optional[Any]=0      ):







       """simple docstring"""




       if str(_A      ).startswith('''mps'''      ):
         __SCREAMING_SNAKE_CASE	:     int =						torch.manual_seed(_A      )
       else:
         __SCREAMING_SNAKE_CASE	:     Any =						torch.Generator(device=_A      ).manual_seed(_A      )
       __SCREAMING_SNAKE_CASE	:     Any =						{
           "prompt": "A painting of a squirrel eating a burger",
           "generator": generator,
           "num_inference_steps": 2,
           "prior_num_inference_steps": 2,
           "output_type": "numpy",
       }
       return inputs



     def       UpperCAmelCase__				(							self				:				Optional[Any]      ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE	:     Tuple =						torch_device == "cpu"

       self._test_attention_slicing_forward_pass(test_max_difference=_A      )




     def       UpperCAmelCase__				(							self				:				str      ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE	:     Union[str, Any] =						torch_device in ["cpu", "mps"]

       self._test_inference_batch_single_identical(test_max_difference=_A      )






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






     """simple docstring"""







     def       UpperCAmelCase__				(							self				:				Optional[Any]      ):







       """simple docstring"""




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



     def       UpperCAmelCase__				(							self				:				Any      ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE	:     List[str] =						load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy'''      )

       __SCREAMING_SNAKE_CASE	:     Any =						StableUnCLIPPipeline.from_pretrained('''fusing/stable-unclip-2-1-l'''			,		torch_dtype=torch.floataa      )
       pipe.to(_A      )
       pipe.set_progress_bar_config(disable=_A      )
       # 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()

       __SCREAMING_SNAKE_CASE	:     int =						torch.Generator(device='''cpu'''      ).manual_seed(0      )
       __SCREAMING_SNAKE_CASE	:     Optional[Any] =						pipe('''anime turle'''			,		generator=_A			,		output_type='''np'''      )

       __SCREAMING_SNAKE_CASE	:     Optional[Any] =						output.images[0]

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

       assert_mean_pixel_difference(_A			,		_A      )




     def       UpperCAmelCase__				(							self				:				List[str]      ):







       """simple docstring"""




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

       __SCREAMING_SNAKE_CASE	:     Union[str, Any] =						StableUnCLIPPipeline.from_pretrained('''fusing/stable-unclip-2-1-l'''			,		torch_dtype=torch.floataa      )
       __SCREAMING_SNAKE_CASE	:     int =						pipe.to(_A      )
       pipe.set_progress_bar_config(disable=_A      )
       pipe.enable_attention_slicing()
       pipe.enable_sequential_cpu_offload()

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

       __SCREAMING_SNAKE_CASE	:     Dict =						torch.cuda.max_memory_allocated()
       # make sure that less than 7 GB is allocated
       assert mem_bytes < 7 * 10**9

# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from argparse import ArgumentParser

from accelerate.commands.config import get_config_parser
from accelerate.commands.env import env_command_parser
from accelerate.commands.launch import launch_command_parser
from accelerate.commands.test import test_command_parser
from accelerate.commands.tpu import tpu_command_parser




def       _snake_case	(   ):
		"""simple docstring"""







		SCREAMING_SNAKE_CASE_  :		int				      =    ArgumentParser("Accelerate CLI tool"   , usage="accelerate <command> [<args>]"   , allow_abbrev=lowerCAmelCase	)
		SCREAMING_SNAKE_CASE_  :		Optional[int]				      =    parser.add_subparsers(help="accelerate command helpers"	)

		# Register commands
		get_config_parser(subparsers=lowerCAmelCase	)
		env_command_parser(subparsers=lowerCAmelCase	)
		launch_command_parser(subparsers=lowerCAmelCase	)
		tpu_command_parser(subparsers=lowerCAmelCase	)
		test_command_parser(subparsers=lowerCAmelCase	)

		# Let's go
		SCREAMING_SNAKE_CASE_  :		Dict				      =    parser.parse_args()

		if not hasattr(lowerCAmelCase   , "func"	):
				parser.print_help()
				exit(1	)

		# Run
		args.func(lowerCAmelCase	)


if __name__ == "__main__":
		main()
