
import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    PNDMScheduler,
    StableDiffusionLDMaDPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS


enable_full_determinism()







class        A_   (			unittest.TestCase   ):




	'''simple docstring'''


	_UpperCamelCase     :     Tuple							 =       StableDiffusionLDMaDPipeline
	_UpperCamelCase     :     List[Any]							 =       TEXT_TO_IMAGE_PARAMS
	_UpperCamelCase     :     str							 =       TEXT_TO_IMAGE_BATCH_PARAMS
	_UpperCamelCase     :     int							 =       TEXT_TO_IMAGE_IMAGE_PARAMS





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		torch.manual_seed(0       )
		lowercase       =    UNetaDConditionModel(
		    block_out_channels=(32, 64)		,    layers_per_block=2		,    sample_size=32		,    in_channels=4		,    out_channels=4		,    down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')		,    up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')		,    cross_attention_dim=32		,    )
		lowercase       =    DDIMScheduler(
		    beta_start=0.00_085		,    beta_end=0.012		,    beta_schedule='scaled_linear'		,    clip_sample=snake_case		,    set_alpha_to_one=snake_case		,    )
		torch.manual_seed(0       )
		lowercase       =    AutoencoderKL(
		    block_out_channels=[32, 64]		,    in_channels=6		,    out_channels=6		,    down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']		,    up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']		,    latent_channels=4		,    )
		torch.manual_seed(0       )
		lowercase       =    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		,    )
		lowercase       =    CLIPTextModel(snake_case       )
		lowercase       =    CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )

		lowercase       =    {
		    'unet': unet,
		    'scheduler': scheduler,
		    'vae': vae,
		    'text_encoder': text_encoder,
		    'tokenizer': tokenizer,
		    'safety_checker': None,
		    'feature_extractor': None,
		}
		return components





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case=0       ):
		if str(snake_case       ).startswith('mps'       ):
			lowercase       =    torch.manual_seed(snake_case       )
		else:
			lowercase       =    torch.Generator(device=snake_case       ).manual_seed(snake_case       )
		lowercase       =    {
		    'prompt': 'A painting of a squirrel eating a burger',
		    'generator': generator,
		    'num_inference_steps': 2,
		    'guidance_scale': 6.0,
		    'output_type': 'numpy',
		}
		return inputs





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase       =    'cpu'  # ensure determinism for the device-dependent torch.Generator

		lowercase       =    self.get_dummy_components()
		lowercase       =    StableDiffusionLDMaDPipeline(**snake_case       )
		lowercase       =    ldmad_pipe.to(snake_case       )
		ldmad_pipe.set_progress_bar_config(disable=snake_case       )

		lowercase       =    self.get_dummy_inputs(snake_case       )
		lowercase       =    ldmad_pipe(**snake_case       )
		lowercase  ,					lowercase       =    output.rgb, output.depth

		lowercase       =    rgb[0, -3:, -3:, -1]
		lowercase       =    depth[0, -3:, -1]

		assert rgb.shape == (1, 64, 64, 3)
		assert depth.shape == (1, 64, 64)

		lowercase       =    np.array(
		    [0.37_338_176, 0.70_247, 0.74_203_193, 0.51_643_604, 0.58_256_793, 0.60_932_136, 0.4_181_095, 0.48_355_877, 0.46_535_262]       )
		lowercase       =    np.array([103.46_727, 85.812_004, 87.849_236]       )

		assert np.abs(image_slice_rgb.flatten() - expected_slice_rgb       ).max() < 1E-2
		assert np.abs(image_slice_depth.flatten() - expected_slice_depth       ).max() < 1E-2





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase       =    self.get_dummy_components()
		lowercase       =    StableDiffusionLDMaDPipeline(**snake_case       )
		lowercase       =    ldmad_pipe.to(snake_case       )
		ldmad_pipe.set_progress_bar_config(disable=snake_case       )

		lowercase       =    self.get_dummy_inputs(snake_case       )
		lowercase       =    3 * [inputs['prompt']]

		# forward
		lowercase       =    ldmad_pipe(**snake_case       )
		lowercase  ,					lowercase       =    output.rgb, output.depth
		lowercase       =    rgb_slice_a[0, -3:, -3:, -1]
		lowercase       =    depth_slice_a[0, -3:, -1]

		lowercase       =    self.get_dummy_inputs(snake_case       )
		lowercase       =    3 * [inputs.pop('prompt'       )]

		lowercase       =    ldmad_pipe.tokenizer(
		    snake_case		,    padding='max_length'		,    max_length=ldmad_pipe.tokenizer.model_max_length		,    truncation=snake_case		,    return_tensors='pt'		,    )
		lowercase       =    text_inputs['input_ids'].to(snake_case       )

		lowercase       =    ldmad_pipe.text_encoder(snake_case       )[0]

		lowercase       =    prompt_embeds

		# forward
		lowercase       =    ldmad_pipe(**snake_case       )
		lowercase  ,					lowercase       =    output.rgb, output.depth
		lowercase       =    rgb_slice_a[0, -3:, -3:, -1]
		lowercase       =    depth_slice_a[0, -3:, -1]

		assert np.abs(rgb_slice_a.flatten() - rgb_slice_a.flatten()       ).max() < 1E-4
		assert np.abs(depth_slice_a.flatten() - depth_slice_a.flatten()       ).max() < 1E-4





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase       =    'cpu'  # ensure determinism for the device-dependent torch.Generator
		lowercase       =    self.get_dummy_components()
		lowercase       =    PNDMScheduler(skip_prk_steps=snake_case       )
		lowercase       =    StableDiffusionLDMaDPipeline(**snake_case       )
		lowercase       =    ldmad_pipe.to(snake_case       )
		ldmad_pipe.set_progress_bar_config(disable=snake_case       )

		lowercase       =    self.get_dummy_inputs(snake_case       )
		lowercase       =    'french fries'
		lowercase       =    ldmad_pipe(**snake_case		,    negative_prompt=snake_case       )

		lowercase  ,					lowercase       =    output.rgb, output.depth
		lowercase       =    rgb[0, -3:, -3:, -1]
		lowercase       =    depth[0, -3:, -1]

		assert rgb.shape == (1, 64, 64, 3)
		assert depth.shape == (1, 64, 64)

		lowercase       =    np.array(
		    [0.37_044, 0.71_811_503, 0.7_223_251, 0.48_603_675, 0.5_638_391, 0.6_364_948, 0.42_833_704, 0.4_901_315, 0.47_926_217]       )
		lowercase       =    np.array([107.84_738, 84.62_802, 89.962_135]       )
		assert np.abs(rgb_slice.flatten() - expected_slice_rgb       ).max() < 1E-2
		assert np.abs(depth_slice.flatten() - expected_slice_depth       ).max() < 1E-2







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




	'''simple docstring'''


	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		super().tearDown()
		gc.collect()
		torch.cuda.empty_cache()





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case="cpu"		,    snake_case=torch.floataa		,    snake_case=0       ):
		lowercase       =    torch.Generator(device=snake_case       ).manual_seed(snake_case       )
		lowercase       =    np.random.RandomState(snake_case       ).standard_normal((1, 4, 64, 64)       )
		lowercase       =    torch.from_numpy(snake_case       ).to(device=snake_case		,    dtype=snake_case       )
		lowercase       =    {
		    'prompt': 'a photograph of an astronaut riding a horse',
		    'latents': latents,
		    'generator': generator,
		    'num_inference_steps': 3,
		    'guidance_scale': 7.5,
		    'output_type': 'numpy',
		}
		return inputs





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase       =    StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'       )
		lowercase       =    ldmad_pipe.to(snake_case       )
		ldmad_pipe.set_progress_bar_config(disable=snake_case       )

		lowercase       =    self.get_inputs(snake_case       )
		lowercase       =    ldmad_pipe(**snake_case       )
		lowercase  ,					lowercase       =    output.rgb, output.depth
		lowercase       =    rgb[0, -3:, -3:, -1].flatten()
		lowercase       =    rgb[0, -3:, -1].flatten()

		assert rgb.shape == (1, 512, 512, 3)
		assert depth.shape == (1, 512, 512)

		lowercase       =    np.array(
		    [0.53_805_465, 0.56_707_305, 0.5_486_515, 0.57_012_236, 0.5_814_511, 0.56_253_487, 0.54_843_014, 0.55_092_263, 0.6_459_706]       )
		lowercase       =    np.array(
		    [0.9_263_781, 0.6_678_672, 0.5_486_515, 0.92_202_145, 0.67_831_135, 0.56_253_487, 0.9_241_694, 0.7_551_478, 0.6_459_706]       )
		assert np.abs(rgb_slice - expected_slice_rgb       ).max() < 3E-3
		assert np.abs(depth_slice - expected_slice_depth       ).max() < 3E-3







@nightly
@require_torch_gpu
class        A_   (			unittest.TestCase   ):




	'''simple docstring'''


	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		super().tearDown()
		gc.collect()
		torch.cuda.empty_cache()





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case="cpu"		,    snake_case=torch.floataa		,    snake_case=0       ):
		lowercase       =    torch.Generator(device=snake_case       ).manual_seed(snake_case       )
		lowercase       =    np.random.RandomState(snake_case       ).standard_normal((1, 4, 64, 64)       )
		lowercase       =    torch.from_numpy(snake_case       ).to(device=snake_case		,    dtype=snake_case       )
		lowercase       =    {
		    'prompt': 'a photograph of an astronaut riding a horse',
		    'latents': latents,
		    'generator': generator,
		    'num_inference_steps': 50,
		    'guidance_scale': 7.5,
		    'output_type': 'numpy',
		}
		return inputs





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase       =    StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'       ).to(snake_case       )
		ldmad_pipe.set_progress_bar_config(disable=snake_case       )

		lowercase       =    self.get_inputs(snake_case       )
		lowercase       =    ldmad_pipe(**snake_case       )
		lowercase  ,					lowercase       =    output.rgb, output.depth

		lowercase       =    0.495_586
		lowercase       =    0.33_795_515
		lowercase       =    112.48_518
		lowercase       =    98.489_746
		assert np.abs(expected_rgb_mean - rgb.mean()       ) < 1E-3
		assert np.abs(expected_rgb_std - rgb.std()       ) < 1E-3
		assert np.abs(expected_depth_mean - depth.mean()       ) < 1E-3
		assert np.abs(expected_depth_std - depth.std()       ) < 1E-3





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase       =    StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d-4c'       ).to(snake_case       )
		ldmad_pipe.set_progress_bar_config(disable=snake_case       )

		lowercase       =    self.get_inputs(snake_case       )
		lowercase       =    ldmad_pipe(**snake_case       )
		lowercase  ,					lowercase       =    output.rgb, output.depth

		lowercase       =    0.4_194_127
		lowercase       =    0.35_375_586
		lowercase       =    0.5_638_502
		lowercase       =    0.34_686_103

		assert rgb.shape == (1, 512, 512, 3)
		assert depth.shape == (1, 512, 512, 1)
		assert np.abs(expected_rgb_mean - rgb.mean()       ) < 1E-3
		assert np.abs(expected_rgb_std - rgb.std()       ) < 1E-3
		assert np.abs(expected_depth_mean - depth.mean()       ) < 1E-3
		assert np.abs(expected_depth_std - depth.std()       ) < 1E-3



# 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 pathlib import Path

import torch

from ...utils import is_npu_available, is_xpu_available
from .config_args import ClusterConfig, default_json_config_file
from .config_utils import SubcommandHelpFormatter


UpperCAmelCase	      = '''Create a default config file for Accelerate with only a few flags set.'''






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE="no" ,							__SCREAMING_SNAKE_CASE = default_json_config_file ,							__SCREAMING_SNAKE_CASE = False							):
	lowercase       =    Path(__SCREAMING_SNAKE_CASE							)
	path.parent.mkdir(parents=__SCREAMING_SNAKE_CASE ,							exist_ok=__SCREAMING_SNAKE_CASE							)
	if path.exists():
		print(
		    F'''Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`.'''							)
		return False
	lowercase       =    mixed_precision.lower()
	if mixed_precision not in ["no", "fp16", "bf16", "fp8"]:
		raise ValueError(
		    F'''`mixed_precision` should be one of \'no\', \'fp16\', \'bf16\', or \'fp8\'. Received {mixed_precision}'''							)
	lowercase       =    {
	    'compute_environment': 'LOCAL_MACHINE',
	    'mixed_precision': mixed_precision,
	}
	if torch.cuda.is_available():
		lowercase       =    torch.cuda.device_count()
		lowercase       =    num_gpus
		lowercase       =    False
		if num_gpus > 1:
			lowercase       =    'MULTI_GPU'
		else:
			lowercase       =    'NO'
	elif is_xpu_available() and use_xpu:
		lowercase       =    torch.xpu.device_count()
		lowercase       =    num_xpus
		lowercase       =    False
		if num_xpus > 1:
			lowercase       =    'MULTI_XPU'
		else:
			lowercase       =    'NO'
	elif is_npu_available():
		lowercase       =    torch.npu.device_count()
		lowercase       =    num_npus
		lowercase       =    False
		if num_npus > 1:
			lowercase       =    'MULTI_NPU'
		else:
			lowercase       =    'NO'
	else:
		lowercase       =    0
		lowercase       =    True
		lowercase       =    1
		lowercase       =    'NO'
	lowercase       =    ClusterConfig(**__SCREAMING_SNAKE_CASE							)
	config.to_json_file(__SCREAMING_SNAKE_CASE							)
	return path






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							):
	lowercase       =    parser.add_parser('default' ,							parents=__SCREAMING_SNAKE_CASE ,							help=__SCREAMING_SNAKE_CASE ,							formatter_class=__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\'.'
	    ) ,							dest='save_location' ,							)

	parser.add_argument(
	    '--mixed_precision' ,							choices=['no', 'fp16', 'bf16'] ,							type=__SCREAMING_SNAKE_CASE ,							help='Whether or not to use mixed precision training. '
	    'Choose between FP16 and BF16 (bfloat16) training. '
	    'BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.' ,							default='no' ,							)
	parser.set_defaults(func=__SCREAMING_SNAKE_CASE							)
	return parser






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE							):
	lowercase       =    write_basic_config(args.mixed_precision ,							args.save_location							)
	if config_file:
		print(F'''accelerate configuration saved at {config_file}'''							)

