



import unittest

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

from diffusers import DDIMScheduler, LDMPipeline, UNetaDModel, VQModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()





class 		UpperCAmelCase						(						unittest.TestCase				):

   '''simple docstring'''





   @property
   def 	__lowerCAmelCase							(	self			):
         torch.manual_seed(0			)
         _lowerCAmelCase 			=  UNetaDModel(
             block_out_channels=(32, 64)       ,     layers_per_block=2       ,     sample_size=32       ,     in_channels=3       ,     out_channels=3       ,     down_block_types=('''DownBlock2D''', '''AttnDownBlock2D''')       ,     up_block_types=('''AttnUpBlock2D''', '''UpBlock2D''')       ,     )
         return model





   @property
   def 	__lowerCAmelCase							(	self			):
         torch.manual_seed(0			)
         _lowerCAmelCase 			=  VQModel(
             block_out_channels=[32, 64]       ,     in_channels=3       ,     out_channels=3       ,     down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']       ,     up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']       ,     latent_channels=3       ,     )
         return model





   @property
   def 	__lowerCAmelCase							(	self			):
         torch.manual_seed(0			)
         _lowerCAmelCase 			=  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=1_000       ,     )
         return CLIPTextModel(_lowerCAmelCase			)





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  self.dummy_uncond_unet
         _lowerCAmelCase 			=  DDIMScheduler()
         _lowerCAmelCase 			=  self.dummy_vq_model

         _lowerCAmelCase 			=  LDMPipeline(unet=_lowerCAmelCase       ,     vqvae=_lowerCAmelCase       ,     scheduler=_lowerCAmelCase			)
         ldm.to(_lowerCAmelCase			)
         ldm.set_progress_bar_config(disable=_lowerCAmelCase			)

         _lowerCAmelCase 			=  torch.manual_seed(0			)
         _lowerCAmelCase 			=  ldm(generator=_lowerCAmelCase       ,     num_inference_steps=2       ,     output_type='''numpy'''			).images

         _lowerCAmelCase 			=  torch.manual_seed(0			)
         _lowerCAmelCase 			=  ldm(generator=_lowerCAmelCase       ,     num_inference_steps=2       ,     output_type='''numpy'''       ,     return_dict=_lowerCAmelCase			)[0]

         _lowerCAmelCase 			=  image[0, -3:, -3:, -1]
         _lowerCAmelCase 			=  image_from_tuple[0, -3:, -3:, -1]

         assert image.shape == (1, 64, 64, 3)
         _lowerCAmelCase 			=  np.array([0.8_512, 0.818, 0.6_411, 0.6_808, 0.4_465, 0.5_618, 0.46, 0.6_231, 0.5_172]			)
         _lowerCAmelCase 			=  1E-2 if torch_device != '''mps''' else 3E-2

         assert np.abs(image_slice.flatten() - expected_slice			).max() < tolerance
         assert np.abs(image_from_tuple_slice.flatten() - expected_slice			).max() < tolerance






@slow
@require_torch
class 		UpperCAmelCase						(						unittest.TestCase				):

   '''simple docstring'''





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  LDMPipeline.from_pretrained('''CompVis/ldm-celebahq-256'''			)
         ldm.to(_lowerCAmelCase			)
         ldm.set_progress_bar_config(disable=_lowerCAmelCase			)

         _lowerCAmelCase 			=  torch.manual_seed(0			)
         _lowerCAmelCase 			=  ldm(generator=_lowerCAmelCase       ,     num_inference_steps=5       ,     output_type='''numpy'''			).images

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

         assert image.shape == (1, 256, 256, 3)
         _lowerCAmelCase 			=  np.array([0.4_399, 0.44_975, 0.46_825, 0.474, 0.4_359, 0.4_581, 0.45_095, 0.4_341, 0.4_447]			)
         _lowerCAmelCase 			=  1E-2 if torch_device != '''mps''' else 3E-2

         assert np.abs(image_slice.flatten() - expected_slice			).max() < tolerance




# 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.

import argparse
import os

from accelerate.utils import ComputeEnvironment

from .cluster import get_cluster_input
from .config_args import cache_dir, default_config_file, default_yaml_config_file, load_config_from_file  # noqa: F401
from .config_utils import _ask_field, _ask_options, _convert_compute_environment  # noqa: F401
from .sagemaker import get_sagemaker_input


UpperCAmelCase_									=							"Launches a series of prompts to create and save a `default_config.yaml` configuration file for your training system. Should always be ran first on your machine"
def 							UpperCAmelCase__		(				)->Any:
      _lowerCAmelCase 			=  _ask_options(
          '''In which compute environment are you running?'''     ,				['''This machine''', '''AWS (Amazon SageMaker)''']     ,				_convert_compute_environment     ,				)
      if compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER:
            _lowerCAmelCase 			=  get_sagemaker_input()
      else:
            _lowerCAmelCase 			=  get_cluster_input()
      return config
def 							UpperCAmelCase__		(				_SCREAMING_SNAKE_CASE   :	int=None				)->str:
      if subparsers is not None:
            _lowerCAmelCase 			=  subparsers.add_parser('''config'''     ,				description=_SCREAMING_SNAKE_CASE				)
      else:
            _lowerCAmelCase 			=  argparse.ArgumentParser('''Accelerate config command'''     ,				description=_SCREAMING_SNAKE_CASE				)

      parser.add_argument(
          '''--config_file'''     ,				default=_SCREAMING_SNAKE_CASE     ,				help=(
              '''The path to use to store the config file. Will default to a file named default_config.yaml in the cache '''
              '''location, which is the content of the environment `HF_HOME` suffixed with \'accelerate\', or if you don\'t have '''
              '''such an environment variable, your cache directory (\'~/.cache\' or the content of `XDG_CACHE_HOME`) suffixed '''
              '''with \'huggingface\'.'''
          )     ,				)

      if subparsers is not None:
            parser.set_defaults(func=_SCREAMING_SNAKE_CASE				)
      return parser
def 							UpperCAmelCase__		(				_SCREAMING_SNAKE_CASE   :	Dict				)->str:
      _lowerCAmelCase 			=  get_user_input()
      if args.config_file is not None:
            _lowerCAmelCase 			=  args.config_file
      else:
            if not os.path.isdir(_SCREAMING_SNAKE_CASE				):
                  os.makedirs(_SCREAMING_SNAKE_CASE				)
            _lowerCAmelCase 			=  default_yaml_config_file

      if config_file.endswith('''.json'''				):
            config.to_json_file(_SCREAMING_SNAKE_CASE				)
      else:
            config.to_yaml_file(_SCREAMING_SNAKE_CASE				)
      print(f'''accelerate configuration saved at {config_file}'''				)
def 							UpperCAmelCase__		(				)->List[Any]:
      _lowerCAmelCase 			=  config_command_parser()
      _lowerCAmelCase 			=  parser.parse_args()
      config_command(_SCREAMING_SNAKE_CASE				)


if __name__ == "__main__":
    main()