
# Copyright 2022 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
import subprocess

from packaging.version import Version, parse

from accelerate.commands.config.config_args import default_config_file, load_config_from_file


lowerCamelCase__        =		"""Run commands across TPU VMs for initial setup before running `accelerate launch`."""
def    __lowerCAmelCase    (_UpperCamelCase=None		):







	if subparsers is not None:
		__lowerCAmelCase      :						Optional[Any]						      =				subparsers.add_parser('tpu-config'	,  description=_description		)
	else:
		__lowerCAmelCase      :						Union[str, Any]						      =				argparse.ArgumentParser('Accelerate tpu-config command'	,  description=_description		)
	# Core arguments
	__lowerCAmelCase      :						List[Any]						      =				parser.add_argument_group(
	    'Config Arguments'	,  'Arguments that can be configured through `accelerate config`.'		)
	config_args.add_argument(
	    '--config_file'	,  type=_UpperCamelCase	,  default=_UpperCamelCase	,  help='Path to the config file to use for accelerate.'	,  )
	config_args.add_argument(
	    '--tpu_name'	,  default=_UpperCamelCase	,  help='The name of the TPU to use. If not specified, will use the TPU specified in the config file.'	,  )
	config_args.add_argument(
	    '--tpu_zone'	,  default=_UpperCamelCase	,  help='The zone of the TPU to use. If not specified, will use the zone specified in the config file.'	,  )
	__lowerCAmelCase      :						List[str]						      =				parser.add_argument_group('TPU Arguments'	,  'Arguments for options ran inside the TPU.'		)
	pod_args.add_argument(
	    '--use_alpha'	,  action='store_true'	,  help='Whether to use `gcloud alpha` when running the TPU training script instead of `gcloud`.'	,  )
	pod_args.add_argument(
	    '--command_file'	,  default=_UpperCamelCase	,  help='The path to the file containing the commands to run on the pod on startup.'	,  )
	pod_args.add_argument(
	    '--command'	,  action='append'	,  nargs='+'	,  help='A command to run on the pod. Can be passed multiple times.'	,  )
	pod_args.add_argument(
	    '--install_accelerate'	,  action='store_true'	,  help='Whether to install accelerate on the pod. Defaults to False.'	,  )
	pod_args.add_argument(
	    '--accelerate_version'	,  default='latest'	,  help='The version of accelerate to install on the pod. If not specified, will use the latest pypi version. Specify \'dev\' to install from GitHub.'	,  )
	pod_args.add_argument(
	    '--debug'	,  action='store_true'	,  help='If set, will print the command that would be run instead of running it.'		)

	if subparsers is not None:
		parser.set_defaults(func=_UpperCamelCase		)
	return parser
def    __lowerCAmelCase    (_UpperCamelCase		):
	__lowerCAmelCase      :						List[Any]						      =				None

	# Get the default from the config file if it exists.
	if args.config_file is not None or os.path.isfile(_UpperCamelCase		):
		__lowerCAmelCase      :						Optional[int]						      =				load_config_from_file(args.config_file		)
		if not args.command_file and defaults.command_file is not None and not args.command:
			__lowerCAmelCase      :						Tuple						      =				defaults.command_file
		if not args.command and defaults.commands is not None:
			__lowerCAmelCase      :						Optional[int]						      =				defaults.commands
		if not args.tpu_name:
			__lowerCAmelCase      :						Union[str, Any]						      =				defaults.tpu_name
		if not args.tpu_zone:
			__lowerCAmelCase      :						Tuple						      =				defaults.tpu_zone
	if args.accelerate_version == "dev":
		__lowerCAmelCase      :						List[Any]						      =				'''git+https://github.com/huggingface/accelerate.git'''
	elif args.accelerate_version == "latest":
		__lowerCAmelCase      :						List[str]						      =				'''accelerate -U'''
	elif isinstance(parse(args.accelerate_version		)	,  _UpperCamelCase		):
		__lowerCAmelCase      :						List[str]						      =				F"accelerate=={args.accelerate_version}"

	if not args.command_file and not args.command:
		raise ValueError('You must specify either a command file or a command to run on the pod.'		)

	if args.command_file:
		with open(args.command_file	,  'r'		) as f:
			__lowerCAmelCase      :						str						      =				[f.read().splitlines()]

    # To turn list of lists into list of strings
	if isinstance(args.command[0]	,  _UpperCamelCase		):
		__lowerCAmelCase      :						List[Any]						      =				[line for cmd in args.command for line in cmd]
	# Default to the shared folder and install accelerate
	__lowerCAmelCase      :						Union[str, Any]						      =				['''cd /usr/share''']
	if args.install_accelerate:
		new_cmd += [F"pip install {args.accelerate_version}"]
	new_cmd += args.command
	__lowerCAmelCase      :						int						      =				'''; '''.join(_UpperCamelCase		)

	# Then send it to gcloud
	# Eventually try to use google-api-core to do this instead of subprocess
	__lowerCAmelCase      :						Optional[Any]						      =				['''gcloud''']
	if args.use_alpha:
		cmd += ["alpha"]
	cmd += [
	    "compute",
	    "tpus",
	    "tpu-vm",
	    "ssh",
	    args.tpu_name,
	    "--zone",
	    args.tpu_zone,
	    "--command",
	    args.command,
	    "--worker",
	    "all",
	]
	if args.debug:
		print(F"Running {' '.join(_UpperCamelCase		)}"		)
		return
	subprocess.run(_UpperCamelCase		)
	print('Successfully setup pod.'		)
def    __lowerCAmelCase    ():
	__lowerCAmelCase      :						Union[str, Any]						      =				tpu_command_parser()
	__lowerCAmelCase      :						List[str]						      =				parser.parse_args()

	tpu_command_launcher(_UpperCamelCase		)

"""simple docstring"""







import gc
import unittest

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

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference




class 			A__	(     _lowerCamelCase	,					unittest.TestCase):
		A_						:	Optional[int]           =     ShapEPipeline
		A_						:	str           =     ['prompt']
		A_						:	Any           =     ['prompt']
		A_						:	List[Any]           =     [
		    'num_images_per_prompt',
		    'num_inference_steps',
		    'generator',
		    'latents',
		    'guidance_scale',
		    'frame_size',
		    'output_type',
		    'return_dict',
		]
		A_						:	Optional[int]           =     False

		@property
		def 							__lowerCamelCase							(				self			):
			return 32

		@property
		def 							__lowerCamelCase							(				self			):
			return 32

		@property
		def 							__lowerCamelCase							(				self			):
			return self.time_input_dim * 4

		@property
		def 							__lowerCamelCase							(				self			):
			return 8

		@property
		def 							__lowerCamelCase							(				self			):
			__lowerCAmelCase      :						Tuple						      =				CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'			)
			return tokenizer

		@property
		def 							__lowerCamelCase							(				self			):
			torch.manual_seed(0			)
			__lowerCAmelCase      :						str						      =				CLIPTextConfig(
			    bos_token_id=0	,		eos_token_id=2	,		hidden_size=self.text_embedder_hidden_size	,		projection_dim=self.text_embedder_hidden_size	,		intermediate_size=37	,		layer_norm_eps=1E-05	,		num_attention_heads=4	,		num_hidden_layers=5	,		pad_token_id=1	,		vocab_size=10_00	,		)
			return CLIPTextModelWithProjection(_SCREAMING_SNAKE_CASE			)

		@property
		def 							__lowerCamelCase							(				self			):
			torch.manual_seed(0			)

			__lowerCAmelCase      :						Optional[Any]						      =				{
			    'num_attention_heads': 2,
			    'attention_head_dim': 16,
			    'embedding_dim': self.time_input_dim,
			    'num_embeddings': 32,
			    'embedding_proj_dim': self.text_embedder_hidden_size,
			    'time_embed_dim': self.time_embed_dim,
			    'num_layers': 1,
			    'clip_embed_dim': self.time_input_dim * 2,
			    'additional_embeddings': 0,
			    'time_embed_act_fn': 'gelu',
			    'norm_in_type': 'layer',
			    'encoder_hid_proj_type': None,
			    'added_emb_type': None,
			}

			__lowerCAmelCase      :						int						      =				PriorTransformer(**_SCREAMING_SNAKE_CASE			)
			return model

		@property
		def 							__lowerCamelCase							(				self			):
			torch.manual_seed(0			)

			__lowerCAmelCase      :						Dict						      =				{
			    'param_shapes': (
			        (self.renderer_dim, 93),
			        (self.renderer_dim, 8),
			        (self.renderer_dim, 8),
			        (self.renderer_dim, 8),
			    ),
			    'd_latent': self.time_input_dim,
			    'd_hidden': self.renderer_dim,
			    'n_output': 12,
			    'background': (
			        0.1,
			        0.1,
			        0.1,
			    ),
			}
			__lowerCAmelCase      :						Union[str, Any]						      =				ShapERenderer(**_SCREAMING_SNAKE_CASE			)
			return model

		def 							__lowerCamelCase							(				self			):
			__lowerCAmelCase      :						List[Any]						      =				self.dummy_prior
			__lowerCAmelCase      :						str						      =				self.dummy_text_encoder
			__lowerCAmelCase      :						List[Any]						      =				self.dummy_tokenizer
			__lowerCAmelCase      :						str						      =				self.dummy_renderer

			__lowerCAmelCase      :						List[Any]						      =				HeunDiscreteScheduler(
			    beta_schedule='exp'	,		num_train_timesteps=10_24	,		prediction_type='sample'	,		use_karras_sigmas=_SCREAMING_SNAKE_CASE	,		clip_sample=_SCREAMING_SNAKE_CASE	,		clip_sample_range=1.0	,		)
			__lowerCAmelCase      :						int						      =				{
			    'prior': prior,
			    'text_encoder': text_encoder,
			    'tokenizer': tokenizer,
			    'renderer': renderer,
			    'scheduler': scheduler,
			}

			return components

		def 							__lowerCamelCase							(				self	,		_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE=0			):
			if str(_SCREAMING_SNAKE_CASE			).startswith('mps'			):
				__lowerCAmelCase      :						int						      =				torch.manual_seed(_SCREAMING_SNAKE_CASE			)
			else:
				__lowerCAmelCase      :						Dict						      =				torch.Generator(device=_SCREAMING_SNAKE_CASE			).manual_seed(_SCREAMING_SNAKE_CASE			)
			__lowerCAmelCase      :						Optional[Any]						      =				{
			    'prompt': 'horse',
			    'generator': generator,
			    'num_inference_steps': 1,
			    'frame_size': 32,
			    'output_type': 'np',
			}
			return inputs

		def 							__lowerCamelCase							(				self			):
			__lowerCAmelCase      :						Any						      =				'cpu'

			__lowerCAmelCase      :						int						      =				self.get_dummy_components()

			__lowerCAmelCase      :						Union[str, Any]						      =				self.pipeline_class(**_SCREAMING_SNAKE_CASE			)
			__lowerCAmelCase      :						Union[str, Any]						      =				pipe.to(_SCREAMING_SNAKE_CASE			)

			pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE			)

			__lowerCAmelCase      :						Optional[int]						      =				pipe(**self.get_dummy_inputs(_SCREAMING_SNAKE_CASE			)			)
			__lowerCAmelCase      :						Optional[Any]						      =				output.images[0]
			__lowerCAmelCase      :						Optional[int]						      =				image[0, -3:, -3:, -1]

			assert image.shape == (20, 32, 32, 3)

			__lowerCAmelCase      :						str						      =				np.array(
			    [
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			    ]			)

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

		def 							__lowerCamelCase							(				self			):
			# NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
			self._test_inference_batch_consistent(batch_sizes=[1, 2]			)

		def 							__lowerCamelCase							(				self			):
			__lowerCAmelCase      :						str						      =				torch_device == 'cpu'
			__lowerCAmelCase      :						str						      =				True

			self._test_inference_batch_single_identical(
			    batch_size=2	,		test_max_difference=_SCREAMING_SNAKE_CASE	,		relax_max_difference=_SCREAMING_SNAKE_CASE	,		)







		def 							__lowerCamelCase							(				self			):
			__lowerCAmelCase      :						str						      =				self.get_dummy_components()
			__lowerCAmelCase      :						List[Any]						      =				self.pipeline_class(**_SCREAMING_SNAKE_CASE			)
			__lowerCAmelCase      :						Tuple						      =				pipe.to(_SCREAMING_SNAKE_CASE			)
			pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE			)

			__lowerCAmelCase      :						int						      =				1
			__lowerCAmelCase      :						List[Any]						      =				2

			__lowerCAmelCase      :						int						      =				self.get_dummy_inputs(_SCREAMING_SNAKE_CASE			)

			for key in inputs.keys():
				if key in self.batch_params:
					__lowerCAmelCase      :						Dict						      =				batch_size * [inputs[key]]

			__lowerCAmelCase      :						Any						      =				pipe(**_SCREAMING_SNAKE_CASE	,		num_images_per_prompt=_SCREAMING_SNAKE_CASE			)[0]

			assert images.shape[0] == batch_size * num_images_per_prompt




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

		def 							__lowerCamelCase							(				self			):
			# clean up the VRAM after each test
			super().tearDown()
			gc.collect()
			torch.cuda.empty_cache()







		def 							__lowerCamelCase							(				self			):
			__lowerCAmelCase      :						List[Any]						      =				load_numpy(
			    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
			    '/shap_e/test_shap_e_np_out.npy'			)
			__lowerCAmelCase      :						Dict						      =				ShapEPipeline.from_pretrained('openai/shap-e'			)
			__lowerCAmelCase      :						int						      =				pipe.to(_SCREAMING_SNAKE_CASE			)
			pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE			)

			__lowerCAmelCase      :						Any						      =				torch.Generator(device=_SCREAMING_SNAKE_CASE			).manual_seed(0			)

			__lowerCAmelCase      :						List[str]						      =				pipe(
			    'a shark'	,		generator=_SCREAMING_SNAKE_CASE	,		guidance_scale=15.0	,		num_inference_steps=64	,		frame_size=64	,		output_type='np'	,		).images[0]

			assert images.shape == (20, 64, 64, 3)

			assert_mean_pixel_difference(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE			)