
import gc
import random
import tempfile
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion_safe import StableDiffusionPipelineSafe as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import require_torch_gpu




class  lowerCamelCase__							(				unittest.TestCase):




	'''simple docstring'''




	def    _lowerCamelCase				(					self		:str						)		-> Optional[int]:
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()







	@property
	def    _lowerCamelCase				(					self		:Optional[int]						)		-> Optional[Any]:
								__UpperCamelCase	:   List[Any]			   =					1
								__UpperCamelCase	:   int			   =					3
								__UpperCamelCase	:   str			   =					(3_2, 3_2)

								__UpperCamelCase	:   Optional[Any]			   =					floats_tensor((batch_size, num_channels) + sizes     ,   rng=random.Random(0						)						).to(a						)
								return image







	@property
	def    _lowerCamelCase				(					self		:Optional[int]						)		-> Tuple:
								torch.manual_seed(0						)
								__UpperCamelCase	:   int			   =					UNetaDConditionModel(
								    block_out_channels=(3_2, 6_4)     ,   layers_per_block=2     ,   sample_size=3_2     ,   in_channels=4     ,   out_channels=4     ,   down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")     ,   up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")     ,   cross_attention_dim=3_2     ,   )
								return model







	@property
	def    _lowerCamelCase				(					self		:Optional[int]						)		-> int:
								torch.manual_seed(0						)
								__UpperCamelCase	:   Any			   =					AutoencoderKL(
								    block_out_channels=[3_2, 6_4]     ,   in_channels=3     ,   out_channels=3     ,   down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]     ,   up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]     ,   latent_channels=4     ,   )
								return model







	@property
	def    _lowerCamelCase				(					self		:Dict						)		-> Union[str, Any]:
								torch.manual_seed(0						)
								__UpperCamelCase	:   Optional[int]			   =					CLIPTextConfig(
								    bos_token_id=0     ,   eos_token_id=2     ,   hidden_size=3_2     ,   intermediate_size=3_7     ,   layer_norm_eps=1E-0_5     ,   num_attention_heads=4     ,   num_hidden_layers=5     ,   pad_token_id=1     ,   vocab_size=1_0_0_0     ,   )
								return CLIPTextModel(a						)







	@property
	def    _lowerCamelCase				(					self		:int						)		-> Optional[int]:
								def extract(*a		:List[Any]     ,   **a		:str						):




															class  lowerCamelCase__							:




																'''simple docstring'''




																def __init__(					self		:Tuple						)		-> Optional[int]:
																							__UpperCamelCase	:   int			   =					torch.ones([0]						)







																def    _lowerCamelCase				(					self		:Dict     ,   a		:Optional[Any]						)		-> Any:
																							self.pixel_values.to(a						)
																							return self
															return Out()

								return extract







	def    _lowerCamelCase				(					self		:Dict						)		-> List[Any]:
								__UpperCamelCase	:   int			   =					"cpu"  # ensure determinism for the device-dependent torch.Generator
								__UpperCamelCase	:   str			   =					self.dummy_cond_unet
								__UpperCamelCase	:   int			   =					DDIMScheduler(
								    beta_start=0.00085     ,   beta_end=0.012     ,   beta_schedule="scaled_linear"     ,   clip_sample=a     ,   set_alpha_to_one=a     ,   )

								__UpperCamelCase	:   int			   =					self.dummy_vae
								__UpperCamelCase	:   Optional[Any]			   =					self.dummy_text_encoder
								__UpperCamelCase	:   Any			   =					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"						)

								# make sure here that pndm scheduler skips prk
								__UpperCamelCase	:   Any			   =					StableDiffusionPipeline(
								    unet=a     ,   scheduler=a     ,   vae=a     ,   text_encoder=a     ,   tokenizer=a     ,   safety_checker=a     ,   feature_extractor=self.dummy_extractor     ,   )
								__UpperCamelCase	:   Optional[int]			   =					sd_pipe.to(a						)
								sd_pipe.set_progress_bar_config(disable=a						)

								__UpperCamelCase	:   str			   =					"A painting of a squirrel eating a burger"

								__UpperCamelCase	:   Optional[int]			   =					torch.Generator(device=a						).manual_seed(0						)
								__UpperCamelCase	:   int			   =					sd_pipe([prompt]     ,   generator=a     ,   guidance_scale=6.0     ,   num_inference_steps=2     ,   output_type="np"						)
								__UpperCamelCase	:   Optional[int]			   =					output.images

								__UpperCamelCase	:   int			   =					torch.Generator(device=a						).manual_seed(0						)
								__UpperCamelCase	:   List[str]			   =					sd_pipe(
								    [prompt]     ,   generator=a     ,   guidance_scale=6.0     ,   num_inference_steps=2     ,   output_type="np"     ,   return_dict=a     ,   )[0]

								__UpperCamelCase	:   List[Any]			   =					image[0, -3:, -3:, -1]
								__UpperCamelCase	:   Dict			   =					image_from_tuple[0, -3:, -3:, -1]

								assert image.shape == (1, 6_4, 6_4, 3)
								__UpperCamelCase	:   str			   =					np.array([0.5756, 0.6118, 0.5005, 0.5041, 0.5471, 0.4726, 0.4976, 0.4865, 0.4864]						)

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







	def    _lowerCamelCase				(					self		:int						)		-> Optional[int]:
								__UpperCamelCase	:   Union[str, Any]			   =					"cpu"  # ensure determinism for the device-dependent torch.Generator
								__UpperCamelCase	:   List[Any]			   =					self.dummy_cond_unet
								__UpperCamelCase	:   int			   =					PNDMScheduler(skip_prk_steps=a						)
								__UpperCamelCase	:   Union[str, Any]			   =					self.dummy_vae
								__UpperCamelCase	:   Dict			   =					self.dummy_text_encoder
								__UpperCamelCase	:   str			   =					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"						)

								# make sure here that pndm scheduler skips prk
								__UpperCamelCase	:   str			   =					StableDiffusionPipeline(
								    unet=a     ,   scheduler=a     ,   vae=a     ,   text_encoder=a     ,   tokenizer=a     ,   safety_checker=a     ,   feature_extractor=self.dummy_extractor     ,   )
								__UpperCamelCase	:   str			   =					sd_pipe.to(a						)
								sd_pipe.set_progress_bar_config(disable=a						)

								__UpperCamelCase	:   Optional[Any]			   =					"A painting of a squirrel eating a burger"
								__UpperCamelCase	:   Dict			   =					torch.Generator(device=a						).manual_seed(0						)
								__UpperCamelCase	:   Dict			   =					sd_pipe([prompt]     ,   generator=a     ,   guidance_scale=6.0     ,   num_inference_steps=2     ,   output_type="np"						)

								__UpperCamelCase	:   Tuple			   =					output.images

								__UpperCamelCase	:   Optional[Any]			   =					torch.Generator(device=a						).manual_seed(0						)
								__UpperCamelCase	:   Optional[int]			   =					sd_pipe(
								    [prompt]     ,   generator=a     ,   guidance_scale=6.0     ,   num_inference_steps=2     ,   output_type="np"     ,   return_dict=a     ,   )[0]

								__UpperCamelCase	:   Dict			   =					image[0, -3:, -3:, -1]
								__UpperCamelCase	:   Optional[Any]			   =					image_from_tuple[0, -3:, -3:, -1]

								assert image.shape == (1, 6_4, 6_4, 3)
								__UpperCamelCase	:   Tuple			   =					np.array([0.5125, 0.5716, 0.4828, 0.5060, 0.5650, 0.4768, 0.5185, 0.4895, 0.4993]						)

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







	def    _lowerCamelCase				(					self		:Optional[Any]						)		-> Any:
								__UpperCamelCase	:   Any			   =					StableDiffusionPipeline.from_pretrained(
								    "hf-internal-testing/tiny-stable-diffusion-lms-pipe"     ,   safety_checker=a						)
								assert isinstance(a     ,   a						)
								assert isinstance(pipe.scheduler     ,   a						)
								assert pipe.safety_checker is None

								__UpperCamelCase	:   List[Any]			   =					pipe("example prompt"     ,   num_inference_steps=2						).images[0]
								assert image is not None

								# check that there's no error when saving a pipeline with one of the models being None
								with tempfile.TemporaryDirectory() as tmpdirname:
															pipe.save_pretrained(a						)
															__UpperCamelCase	:   Optional[Any]			   =					StableDiffusionPipeline.from_pretrained(a						)

								# sanity check that the pipeline still works
								assert pipe.safety_checker is None
								__UpperCamelCase	:   Tuple			   =					pipe("example prompt"     ,   num_inference_steps=2						).images[0]
								assert image is not None







	@unittest.skipIf(torch_device != "cuda"     ,   "This test requires a GPU"						)
	def    _lowerCamelCase				(					self		:Any						)		-> List[Any]:
								__UpperCamelCase	:   Union[str, Any]			   =					self.dummy_cond_unet
								__UpperCamelCase	:   Tuple			   =					PNDMScheduler(skip_prk_steps=a						)
								__UpperCamelCase	:   Any			   =					self.dummy_vae
								__UpperCamelCase	:   Optional[int]			   =					self.dummy_text_encoder
								__UpperCamelCase	:   Optional[Any]			   =					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"						)

								# put models in fp16
								__UpperCamelCase	:   List[Any]			   =					unet.half()
								__UpperCamelCase	:   List[Any]			   =					vae.half()
								__UpperCamelCase	:   str			   =					bert.half()

								# make sure here that pndm scheduler skips prk
								__UpperCamelCase	:   Any			   =					StableDiffusionPipeline(
								    unet=a     ,   scheduler=a     ,   vae=a     ,   text_encoder=a     ,   tokenizer=a     ,   safety_checker=a     ,   feature_extractor=self.dummy_extractor     ,   )
								__UpperCamelCase	:   Optional[Any]			   =					sd_pipe.to(a						)
								sd_pipe.set_progress_bar_config(disable=a						)

								__UpperCamelCase	:   str			   =					"A painting of a squirrel eating a burger"
								__UpperCamelCase	:   Union[str, Any]			   =					sd_pipe([prompt]     ,   num_inference_steps=2     ,   output_type="np"						).images

								assert image.shape == (1, 6_4, 6_4, 3)




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




	'''simple docstring'''




	def    _lowerCamelCase				(					self		:int						)		-> Dict:
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()







	def    _lowerCamelCase				(					self		:str						)		-> Tuple:
								__UpperCamelCase	:   int			   =					StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"     ,   safety_checker=a						)
								__UpperCamelCase	:   int			   =					LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config						)
								__UpperCamelCase	:   str			   =					sd_pipe.to(a						)
								sd_pipe.set_progress_bar_config(disable=a						)

								__UpperCamelCase	:   List[Any]			   =					(
								    "portrait of girl with smokey eyes makeup in abandoned hotel, grange clothes, redshift, wide high angle"
								    " coloured polaroid photograph with flash, kodak film, hyper real, stunning moody cinematography, with"
								    " anamorphic lenses, by maripol, fallen angels by wong kar - wai, style of suspiria and neon demon and"
								    " children from bahnhof zoo, detailed "
								)
								__UpperCamelCase	:   Tuple			   =					4_0_0_3_6_6_0_3_4_6
								__UpperCamelCase	:   Any			   =					7

								# without safety guidance (sld_guidance_scale = 0)
								__UpperCamelCase	:   Dict			   =					torch.manual_seed(a						)
								__UpperCamelCase	:   Any			   =					sd_pipe(
								    [prompt]     ,   generator=a     ,   guidance_scale=a     ,   num_inference_steps=5_0     ,   output_type="np"     ,   width=5_1_2     ,   height=5_1_2     ,   sld_guidance_scale=0     ,   )

								__UpperCamelCase	:   Dict			   =					output.images
								__UpperCamelCase	:   str			   =					image[0, -3:, -3:, -1]
								__UpperCamelCase	:   Dict			   =					[0.2278, 0.2231, 0.2249, 0.2333, 0.2303, 0.1885, 0.2273, 0.2144, 0.2176]

								assert image.shape == (1, 5_1_2, 5_1_2, 3)

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

								# without safety guidance (strong configuration)
								__UpperCamelCase	:   Any			   =					torch.manual_seed(a						)
								__UpperCamelCase	:   Any			   =					sd_pipe(
								    [prompt]     ,   generator=a     ,   guidance_scale=a     ,   num_inference_steps=5_0     ,   output_type="np"     ,   width=5_1_2     ,   height=5_1_2     ,   sld_guidance_scale=2_0_0_0     ,   sld_warmup_steps=7     ,   sld_threshold=0.025     ,   sld_momentum_scale=0.5     ,   sld_mom_beta=0.7     ,   )

								__UpperCamelCase	:   Union[str, Any]			   =					output.images
								__UpperCamelCase	:   List[str]			   =					image[0, -3:, -3:, -1]
								__UpperCamelCase	:   Optional[int]			   =					[0.2383, 0.2276, 0.236, 0.2192, 0.2186, 0.2053, 0.1971, 0.1901, 0.1719]

								assert image.shape == (1, 5_1_2, 5_1_2, 3)

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







	def    _lowerCamelCase				(					self		:Union[str, Any]						)		-> Union[str, Any]:
								__UpperCamelCase	:   Optional[int]			   =					StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"     ,   safety_checker=a						)
								__UpperCamelCase	:   List[str]			   =					LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config						)
								__UpperCamelCase	:   Any			   =					sd_pipe.to(a						)
								sd_pipe.set_progress_bar_config(disable=a						)

								__UpperCamelCase	:   Union[str, Any]			   =					"padme amidala taking a bath artwork, safe for work, no nudity"
								__UpperCamelCase	:   Any			   =					2_7_3_4_9_7_1_7_5_5
								__UpperCamelCase	:   Optional[Any]			   =					7

								__UpperCamelCase	:   Dict			   =					torch.manual_seed(a						)
								__UpperCamelCase	:   int			   =					sd_pipe(
								    [prompt]     ,   generator=a     ,   guidance_scale=a     ,   num_inference_steps=5_0     ,   output_type="np"     ,   width=5_1_2     ,   height=5_1_2     ,   sld_guidance_scale=0     ,   )

								__UpperCamelCase	:   List[str]			   =					output.images
								__UpperCamelCase	:   List[Any]			   =					image[0, -3:, -3:, -1]
								__UpperCamelCase	:   Tuple			   =					[0.3502, 0.3622, 0.3396, 0.3642, 0.3478, 0.3318, 0.35, 0.3348, 0.3297]

								assert image.shape == (1, 5_1_2, 5_1_2, 3)

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

								__UpperCamelCase	:   int			   =					torch.manual_seed(a						)
								__UpperCamelCase	:   Dict			   =					sd_pipe(
								    [prompt]     ,   generator=a     ,   guidance_scale=a     ,   num_inference_steps=5_0     ,   output_type="np"     ,   width=5_1_2     ,   height=5_1_2     ,   sld_guidance_scale=2_0_0_0     ,   sld_warmup_steps=7     ,   sld_threshold=0.025     ,   sld_momentum_scale=0.5     ,   sld_mom_beta=0.7     ,   )

								__UpperCamelCase	:   List[str]			   =					output.images
								__UpperCamelCase	:   List[str]			   =					image[0, -3:, -3:, -1]
								__UpperCamelCase	:   Optional[int]			   =					[0.5531, 0.5206, 0.4895, 0.5156, 0.5182, 0.4751, 0.4802, 0.4803, 0.4443]

								assert image.shape == (1, 5_1_2, 5_1_2, 3)

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







	def    _lowerCamelCase				(					self		:Optional[Any]						)		-> List[str]:
								__UpperCamelCase	:   Dict			   =					StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"						)
								__UpperCamelCase	:   Any			   =					sd_pipe.to(a						)
								sd_pipe.set_progress_bar_config(disable=a						)

								__UpperCamelCase	:   Any			   =					(
								    "the four horsewomen of the apocalypse, painting by tom of finland, gaston bussiere, craig mullins, j. c."
								    " leyendecker"
								)
								__UpperCamelCase	:   List[str]			   =					1_0_4_4_3_5_5_2_3_4
								__UpperCamelCase	:   str			   =					1_2

								__UpperCamelCase	:   Dict			   =					torch.manual_seed(a						)
								__UpperCamelCase	:   str			   =					sd_pipe(
								    [prompt]     ,   generator=a     ,   guidance_scale=a     ,   num_inference_steps=5_0     ,   output_type="np"     ,   width=5_1_2     ,   height=5_1_2     ,   sld_guidance_scale=0     ,   )

								__UpperCamelCase	:   Tuple			   =					output.images
								__UpperCamelCase	:   Any			   =					image[0, -3:, -3:, -1]
								__UpperCamelCase	:   Optional[int]			   =					np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]						)

								assert image.shape == (1, 5_1_2, 5_1_2, 3)

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

								__UpperCamelCase	:   Dict			   =					torch.manual_seed(a						)
								__UpperCamelCase	:   Optional[Any]			   =					sd_pipe(
								    [prompt]     ,   generator=a     ,   guidance_scale=a     ,   num_inference_steps=5_0     ,   output_type="np"     ,   width=5_1_2     ,   height=5_1_2     ,   sld_guidance_scale=2_0_0_0     ,   sld_warmup_steps=7     ,   sld_threshold=0.025     ,   sld_momentum_scale=0.5     ,   sld_mom_beta=0.7     ,   )

								__UpperCamelCase	:   Optional[Any]			   =					output.images
								__UpperCamelCase	:   List[Any]			   =					image[0, -3:, -3:, -1]
								__UpperCamelCase	:   List[str]			   =					np.array([0.5818, 0.6285, 0.6835, 0.6019, 0.625, 0.6754, 0.6096, 0.6334, 0.6561]						)
								assert image.shape == (1, 5_1_2, 5_1_2, 3)

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

import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized, parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
		from sagemaker import Session, TrainingJobAnalytics
		from sagemaker.huggingface import HuggingFace




@pytest.mark.skipif(
    literal_eval(os.getenv('TEST_SAGEMAKER'     ,						'False')) is not True     ,						reason='Skipping test because should only be run when releasing minor transformers version'     ,						)
@pytest.mark.usefixtures('sm_env')
@parameterized_class(
    [
        {
            'framework': 'pytorch',
            'script': 'run_glue.py',
            'model_name_or_path': 'distilbert-base-cased',
            'instance_type': 'ml.p3.16xlarge',
            'results': {'train_runtime': 6_5_0, 'eval_accuracy': 0.7, 'eval_loss': 0.6},
        },
        {
            'framework': 'pytorch',
            'script': 'run_ddp.py',
            'model_name_or_path': 'distilbert-base-cased',
            'instance_type': 'ml.p3.16xlarge',
            'results': {'train_runtime': 6_0_0, 'eval_accuracy': 0.7, 'eval_loss': 0.6},
        },
        {
            'framework': 'tensorflow',
            'script': 'run_tf_dist.py',
            'model_name_or_path': 'distilbert-base-cased',
            'instance_type': 'ml.p3.16xlarge',
            'results': {'train_runtime': 6_0_0, 'eval_accuracy': 0.6, 'eval_loss': 0.7},
        },
    ])
class  lowerCamelCase__							(				unittest.TestCase):




	'''simple docstring'''




	def    _lowerCamelCase				(					self		:List[Any]						)		-> Any:
								if self.framework == "pytorch":
															subprocess.run(
															    f'cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py'.split()     ,   encoding="utf-8"     ,   check=a     ,   )
								assert hasattr(self     ,   "env"						)







	def    _lowerCamelCase				(					self		:Any     ,   a		:Optional[Any]						)		-> Dict:
								__UpperCamelCase	:   str			   =					f'{self.env.base_job_name}-{instance_count}-{"ddp" if "ddp" in self.script else "smd"}'
								# distributed data settings
								__UpperCamelCase	:   Optional[int]			   =					{"smdistributed": {"dataparallel": {"enabled": True}}} if self.script != "run_ddp.py" else None

								# creates estimator
								return HuggingFace(
								    entry_point=self.script     ,   source_dir=self.env.test_path     ,   role=self.env.role     ,   image_uri=self.env.image_uri     ,   base_job_name=a     ,   instance_count=a     ,   instance_type=self.instance_type     ,   debugger_hook_config=a     ,   hyperparameters={**self.env.distributed_hyperparameters, "model_name_or_path": self.model_name_or_path}     ,   metric_definitions=self.env.metric_definitions     ,   distribution=a     ,   py_version="py36"     ,   )







	def    _lowerCamelCase				(					self		:Dict     ,   a		:Dict						)		-> Optional[int]:
								TrainingJobAnalytics(a						).export_csv(f'{self.env.test_path}/{job_name}_metrics.csv'						)







	@parameterized.expand([(2,)]						)
	def    _lowerCamelCase				(					self		:Dict     ,   a		:Tuple						)		-> List[Any]:
								# create estimator
								__UpperCamelCase	:   int			   =					self.create_estimator(a						)

								# run training
								estimator.fit()

								# result dataframe
								__UpperCamelCase	:   Optional[int]			   =					TrainingJobAnalytics(estimator.latest_training_job.name						).dataframe()

								# extract kpis
								__UpperCamelCase	:   Optional[Any]			   =					list(result_metrics_df[result_metrics_df.metric_name == "eval_accuracy"]["value"]						)
								__UpperCamelCase	:   Tuple			   =					list(result_metrics_df[result_metrics_df.metric_name == "eval_loss"]["value"]						)
								# get train time from SageMaker job, this includes starting, preprocessing, stopping
								__UpperCamelCase	:   int			   =					(
								    Session().describe_training_job(estimator.latest_training_job.name						).get("TrainingTimeInSeconds"     ,   9_9_9_9_9_9						)
								)

								# assert kpis
								assert train_runtime <= self.results["train_runtime"]
								assert all(t >= self.results["eval_accuracy"] for t in eval_accuracy						)
								assert all(t <= self.results["eval_loss"] for t in eval_loss						)

								# dump tests result into json file to share in PR
								with open(f'{estimator.latest_training_job.name}.json'     ,   "w"						) as outfile:
															json.dump({"train_time": train_runtime, "eval_accuracy": eval_accuracy, "eval_loss": eval_loss}     ,   a						)