



import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import 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
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






class      UpperCamelCase    (						a__       ,    a__       ,    unittest.TestCase     ):
	lowerCAmelCase	:   Tuple				=	StableDiffusionSAGPipeline
	lowerCAmelCase	:   Dict				=	TEXT_TO_IMAGE_PARAMS
	lowerCAmelCase	:   Tuple				=	TEXT_TO_IMAGE_BATCH_PARAMS
	lowerCAmelCase	:   List[Any]				=	TEXT_TO_IMAGE_IMAGE_PARAMS
	lowerCAmelCase	:   Union[str, Any]				=	TEXT_TO_IMAGE_IMAGE_PARAMS
	lowerCAmelCase	:   int				=	False



	def 			__A							(       self		):
							torch.manual_seed(0		)
							A__							=       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			,	)
							A__							=       DDIMScheduler(
							    beta_start=0.00_085			,	beta_end=0.012			,	beta_schedule="scaled_linear"			,	clip_sample=_lowerCamelCase			,	set_alpha_to_one=_lowerCamelCase			,	)
							torch.manual_seed(0		)
							A__							=       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			,	)
							torch.manual_seed(0		)
							A__							=       CLIPTextConfig(
							    bos_token_id=0			,	eos_token_id=2			,	hidden_size=32			,	intermediate_size=37			,	layer_norm_eps=1e-0_5			,	num_attention_heads=4			,	num_hidden_layers=5			,	pad_token_id=1			,	vocab_size=1_000			,	)
							A__							=       CLIPTextModel(_lowerCamelCase		)
							A__							=       CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)

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



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__=0		):
							if str(_lowerCamelCase		).startswith("mps"		):
													A__							=       torch.manual_seed(_lowerCamelCase		)
							else:
													A__							=       torch.Generator(device=_lowerCamelCase		).manual_seed(_lowerCamelCase		)
							A__							=       {
							    '''prompt''': '''.''',
							    '''generator''': generator,
							    '''num_inference_steps''': 2,
							    '''guidance_scale''': 1.0,
							    '''sag_scale''': 1.0,
							    '''output_type''': '''numpy''',
							}
							return inputs



	def 			__A							(       self		):
							super().test_inference_batch_single_identical(expected_max_diff=3e-3		)






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



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



	def 			__A							(       self		):
							A__							=       StableDiffusionSAGPipeline.from_pretrained("CompVis/stable-diffusion-v1-4"		)
							A__							=       sag_pipe.to(_lowerCamelCase		)
							sag_pipe.set_progress_bar_config(disable=_lowerCamelCase		)

							A__							=       '''.'''
							A__							=       torch.manual_seed(0		)
							A__							=       sag_pipe(
							    [prompt]			,	generator=_lowerCamelCase			,	guidance_scale=7.5			,	sag_scale=1.0			,	num_inference_steps=20			,	output_type="np"		)

							A__							=       output.images

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

							assert image.shape == (1, 512, 512, 3)
							A__							=       np.array([0.1_568, 0.1_738, 0.1_695, 0.1_693, 0.1_507, 0.1_705, 0.1_547, 0.1_751, 0.1_949]		)

							assert np.abs(image_slice.flatten() - expected_slice		).max() < 5e-2



	def 			__A							(       self		):
							A__							=       StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"		)
							A__							=       sag_pipe.to(_lowerCamelCase		)
							sag_pipe.set_progress_bar_config(disable=_lowerCamelCase		)

							A__							=       '''.'''
							A__							=       torch.manual_seed(0		)
							A__							=       sag_pipe(
							    [prompt]			,	generator=_lowerCamelCase			,	guidance_scale=7.5			,	sag_scale=1.0			,	num_inference_steps=20			,	output_type="np"		)

							A__							=       output.images

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

							assert image.shape == (1, 512, 512, 3)
							A__							=       np.array([0.3_459, 0.2_876, 0.2_537, 0.3_002, 0.2_671, 0.2_160, 0.3_026, 0.2_262, 0.2_371]		)

							assert np.abs(image_slice.flatten() - expected_slice		).max() < 5e-2



	def 			__A							(       self		):
							A__							=       StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"		)
							A__							=       sag_pipe.to(_lowerCamelCase		)
							sag_pipe.set_progress_bar_config(disable=_lowerCamelCase		)

							A__							=       '''.'''
							A__							=       torch.manual_seed(0		)
							A__							=       sag_pipe(
							    [prompt]			,	width=768			,	height=512			,	generator=_lowerCamelCase			,	guidance_scale=7.5			,	sag_scale=1.0			,	num_inference_steps=20			,	output_type="np"			,	)

							A__							=       output.images

							assert image.shape == (1, 512, 768, 3)










import os
import time

import pytest

from datasets.utils.filelock import FileLock, Timeout
def  UpperCamelCase						(       _A			:				Optional[int]    )->       List[Any]:





						"""simple docstring"""





						A__							=       FileLock(str(tmpdir / "foo.lock"    )    )
						A__							=       FileLock(str(tmpdir / "foo.lock"    )    )
						A__							=       0.01
						with locka.acquire():
												with pytest.raises(_A    ):
																		A__							=       time.time()
																		locka.acquire(_A    )
																		assert time.time() - _start > timeout
def  UpperCamelCase						(       _A			:				str    )->       List[Any]:





						"""simple docstring"""





						A__							=       "a" * 1000 + ".lock"
						A__							=       FileLock(str(tmpdir / filename    )    )
						assert locka._lock_file.endswith(".lock"    )
						assert not locka._lock_file.endswith(_A    )
						assert len(os.path.basename(locka._lock_file    )    ) <= 255
						A__							=       FileLock(tmpdir / filename    )
						with locka.acquire():
												with pytest.raises(_A    ):
																		locka.acquire(0    )





