






import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    TextToVideoSDPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import is_xformers_available, load_numpy, skip_mps, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

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


enable_full_determinism()



@skip_mps
class   __lowercase    (			UpperCAmelCase_		,	unittest.TestCase     ):




			"""simple docstring"""







			_UpperCAmelCase       :							List[Any]      							=					TextToVideoSDPipeline
			_UpperCAmelCase       :							List[str]      							=					TEXT_TO_IMAGE_PARAMS
			_UpperCAmelCase       :							Tuple      							=					TEXT_TO_IMAGE_BATCH_PARAMS
			# No `output_type`.
			_UpperCAmelCase       :							Any      							=					frozenset(
			    [
			        '''num_inference_steps''',
			        '''generator''',
			        '''latents''',
			        '''return_dict''',
			        '''callback''',
			        '''callback_steps''',
			    ]     )



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[int]):
									torch.manual_seed(0)
									SCREAMING_SNAKE_CASE_:			List[str]  		=	UNetaDConditionModel(
									    block_out_channels=(32, 64, 64, 64)       ,						layers_per_block=2       ,						sample_size=32       ,						in_channels=4       ,						out_channels=4       ,						down_block_types=("CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "DownBlock3D")       ,						up_block_types=("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D")       ,						cross_attention_dim=32       ,						attention_head_dim=4       ,						)
									SCREAMING_SNAKE_CASE_:			str  		=	DDIMScheduler(
									    beta_start=0.0_0085       ,						beta_end=0.012       ,						beta_schedule="scaled_linear"       ,						clip_sample=lowerCAmelCase__       ,						set_alpha_to_one=lowerCAmelCase__       ,						)
									torch.manual_seed(0)
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	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       ,						sample_size=128       ,						)
									torch.manual_seed(0)
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	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       ,						hidden_act="gelu"       ,						projection_dim=512       ,						)
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	CLIPTextModel(lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			Any  		=	CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

									SCREAMING_SNAKE_CASE_:			Dict  		=	{
									    "unet": unet,
									    "scheduler": scheduler,
									    "vae": vae,
									    "text_encoder": text_encoder,
									    "tokenizer": tokenizer,
									}
									return components



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[Any]       ,						lowerCAmelCase__  :					Dict       ,						lowerCAmelCase__  :					str=0):
									if str(lowerCAmelCase__).startswith("mps"):
															SCREAMING_SNAKE_CASE_:			Dict  		=	torch.manual_seed(lowerCAmelCase__)
									else:
															SCREAMING_SNAKE_CASE_:			str  		=	torch.Generator(device=lowerCAmelCase__).manual_seed(lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	{
									    "prompt": "A painting of a squirrel eating a burger",
									    "generator": generator,
									    "num_inference_steps": 2,
									    "guidance_scale": 6.0,
									    "output_type": "pt",
									}
									return inputs



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Any):
									SCREAMING_SNAKE_CASE_:			Tuple  		=	"cpu"  # ensure determinism for the device-dependent torch.Generator
									SCREAMING_SNAKE_CASE_:			int  		=	self.get_dummy_components()
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	TextToVideoSDPipeline(**lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			List[str]  		=	sd_pipe.to(lowerCAmelCase__)
									sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__)

									SCREAMING_SNAKE_CASE_:			List[str]  		=	self.get_dummy_inputs(lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			str  		=	"np"
									SCREAMING_SNAKE_CASE_:			Dict  		=	sd_pipe(**lowerCAmelCase__).frames
									SCREAMING_SNAKE_CASE_:			Any  		=	frames[0][-3:, -3:, -1]

									assert frames[0].shape == (64, 64, 3)
									SCREAMING_SNAKE_CASE_:			Optional[int]  		=	np.array([158.0, 160.0, 153.0, 125.0, 100.0, 121.0, 111.0, 93.0, 113.0])

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



			def 	_SCREAMING_SNAKE_CASE      (			self  :					List[str]):
									self._test_attention_slicing_forward_pass(test_mean_pixel_difference=lowerCAmelCase__       ,						expected_max_diff=3E-3)



			@unittest.skipIf(
			    torch_device != "cuda" or not is_xformers_available()       ,						reason="XFormers attention is only available with CUDA and `xformers` installed"       ,						)
			def 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[int]):
									self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=lowerCAmelCase__       ,						expected_max_diff=1E-2)



			@unittest.skip(reason="Batching needs to be properly figured out first for this pipeline.")
			def 	_SCREAMING_SNAKE_CASE      (			self  :					List[Any]):
									pass



			@unittest.skip(reason="Batching needs to be properly figured out first for this pipeline.")
			def 	_SCREAMING_SNAKE_CASE      (			self  :					int):
									pass



			@unittest.skip(reason="`num_images_per_prompt` argument is not supported for this pipeline.")
			def 	_SCREAMING_SNAKE_CASE      (			self  :					List[Any]):
									pass




			def 	_SCREAMING_SNAKE_CASE      (			self  :					Any):
									return super().test_progress_bar()







@slow
@skip_mps
class   __lowercase    (			unittest.TestCase     ):




			"""simple docstring"""







			def 	_SCREAMING_SNAKE_CASE      (			self  :					List[str]):
									SCREAMING_SNAKE_CASE_:			int  		=	load_numpy(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video.npy")

									SCREAMING_SNAKE_CASE_:			Dict  		=	TextToVideoSDPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b")
									SCREAMING_SNAKE_CASE_:			List[str]  		=	DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
									SCREAMING_SNAKE_CASE_:			Dict  		=	pipe.to("cuda")

									SCREAMING_SNAKE_CASE_:			str  		=	"Spiderman is surfing"
									SCREAMING_SNAKE_CASE_:			str  		=	torch.Generator(device="cpu").manual_seed(0)

									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	pipe(lowerCAmelCase__       ,						generator=lowerCAmelCase__       ,						num_inference_steps=25       ,						output_type="pt").frames
									SCREAMING_SNAKE_CASE_:			Tuple  		=	video_frames.cpu().numpy()

									assert np.abs(expected_video - video).mean() < 5E-2




			def 	_SCREAMING_SNAKE_CASE      (			self  :					Tuple):
									SCREAMING_SNAKE_CASE_:			str  		=	load_numpy(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video_2step.npy")

									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	TextToVideoSDPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b")
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	pipe.to("cuda")

									SCREAMING_SNAKE_CASE_:			Dict  		=	"Spiderman is surfing"
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	torch.Generator(device="cpu").manual_seed(0)

									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	pipe(lowerCAmelCase__       ,						generator=lowerCAmelCase__       ,						num_inference_steps=2       ,						output_type="pt").frames
									SCREAMING_SNAKE_CASE_:			Tuple  		=	video_frames.cpu().numpy()

									assert np.abs(expected_video - video).mean() < 5E-2









# 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							:  List[str]            =						"""Run commands across TPU VMs for initial setup before running `accelerate launch`."""







def       A_				(   _UpperCAmelCase=None   ):
						if subparsers is not None:
												SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	subparsers.add_parser("tpu-config"    ,     description=_description   )
						else:
												SCREAMING_SNAKE_CASE_:			Optional[int]  		=	argparse.ArgumentParser("Accelerate tpu-config command"    ,     description=_description   )
						# Core arguments
						SCREAMING_SNAKE_CASE_:			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."    ,     )
						SCREAMING_SNAKE_CASE_:			Tuple  		=	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       A_				(   _UpperCAmelCase   ):
						SCREAMING_SNAKE_CASE_:			Tuple  		=	None

						# Get the default from the config file if it exists.
						if args.config_file is not None or os.path.isfile(_UpperCAmelCase   ):
												SCREAMING_SNAKE_CASE_:			List[Any]  		=	load_config_from_file(args.config_file   )
												if not args.command_file and defaults.command_file is not None and not args.command:
																		SCREAMING_SNAKE_CASE_:			Any  		=	defaults.command_file
												if not args.command and defaults.commands is not None:
																		SCREAMING_SNAKE_CASE_:			int  		=	defaults.commands
												if not args.tpu_name:
																		SCREAMING_SNAKE_CASE_:			List[str]  		=	defaults.tpu_name
												if not args.tpu_zone:
																		SCREAMING_SNAKE_CASE_:			Dict  		=	defaults.tpu_zone
						if args.accelerate_version == "dev":
												SCREAMING_SNAKE_CASE_:			Tuple  		=	"git+https://github.com/huggingface/accelerate.git"
						elif args.accelerate_version == "latest":
												SCREAMING_SNAKE_CASE_:			str  		=	"accelerate -U"
						elif isinstance(parse(args.accelerate_version   )    ,     _UpperCAmelCase   ):
												SCREAMING_SNAKE_CASE_:			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:
																		SCREAMING_SNAKE_CASE_:			Dict  		=	[f.read().splitlines()]

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

						# Then send it to gcloud
						# Eventually try to use google-api-core to do this instead of subprocess
						SCREAMING_SNAKE_CASE_:			int  		=	["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       A_				(   ):
						SCREAMING_SNAKE_CASE_:			List[str]  		=	tpu_command_parser()
						SCREAMING_SNAKE_CASE_:			List[Any]  		=	parser.parse_args()

						tpu_command_launcher(_UpperCAmelCase   )






