





from copy import deepcopy

import torch
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader

from accelerate.accelerator import Accelerator
from accelerate.state import GradientState
from accelerate.test_utils import RegressionDataset, RegressionModel
from accelerate.utils import DistributedType, is_torch_version, set_seed





def lowerCAmelCase__(__snake_case       ,__snake_case       ,__snake_case       ,__snake_case							)				->      Union[str, Any]:







	'''simple docstring'''







	for param, grad_param in zip(model_a.parameters()       ,model_b.parameters()							):
		if not param.requires_grad:
			continue
		if not did_step:
			# Grads should not be in sync
			assert (
			    torch.allclose(param.grad       ,grad_param.grad							) is False
			), F'Gradients in sync when they should not be at iteration {iteration}:\nmodel_a grad ({param.grad}) == model_b grad ({grad_param.grad})'
		else:
			# Grads should be in sync
			assert (
			    torch.allclose(param.grad       ,grad_param.grad							) is True
			), F'Gradients not in sync when they should be at iteration {iteration}:\nmodel_a grad ({param.grad}) != model_b grad ({grad_param.grad})'





def lowerCAmelCase__(__snake_case       ,__snake_case       ,__snake_case       ,__snake_case       ,__snake_case=True							)				->      Any:







	'''simple docstring'''







	model.train()
	lowerCamelCase__								=		model(_lowerCAmelCase							)
	lowerCamelCase__								=		F.mse_loss(_lowerCAmelCase       ,target.to(output.device							)							)
	if not do_backward:
		loss /= accelerator.gradient_accumulation_steps
		loss.backward()
	else:
		accelerator.backward(_lowerCAmelCase							)





def lowerCAmelCase__(__snake_case       ,__snake_case=False							)				->      str:







	'''simple docstring'''







	set_seed(42							)
	lowerCamelCase__								=		RegressionModel()
	lowerCamelCase__								=		deepcopy(_lowerCAmelCase							)
	lowerCamelCase__								=		RegressionDataset(length=80							)
	lowerCamelCase__								=		DataLoader(_lowerCAmelCase       ,batch_size=16							)
	model.to(accelerator.device							)
	if sched:
		lowerCamelCase__								=		AdamW(params=model.parameters()       ,lr=1E-3							)
		lowerCamelCase__								=		AdamW(params=ddp_model.parameters()       ,lr=1E-3							)
		lowerCamelCase__								=		LambdaLR(_lowerCAmelCase       ,lr_lambda=lambda __snake_case							: epoch**0.6_5							)
		lowerCamelCase__								=		LambdaLR(_lowerCAmelCase       ,lr_lambda=lambda __snake_case							: epoch**0.6_5							)
	# Make a copy of `model`
	if sched:
		lowerCamelCase__								=		accelerator.prepare(_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase							)
	else:
		lowerCamelCase__								=		accelerator.prepare(_lowerCAmelCase       ,_lowerCAmelCase							)
	if sched:
		return (model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched)
	return model, ddp_model, dataloader





def lowerCAmelCase__(__snake_case							)				->      Optional[Any]:







	'''simple docstring'''







	lowerCamelCase__								=		get_training_setup(_lowerCAmelCase							)
	# Use a single batch
	lowerCamelCase__								=		next(iter(_lowerCAmelCase							)							).values()
	for iteration in range(3							):
		# Gather the distributed inputs and targs for the base model
		lowerCamelCase__								=		accelerator.gather((ddp_input, ddp_target)							)
		lowerCamelCase__								=		input.to(accelerator.device							), target.to(accelerator.device							)
		# Perform our initial ground truth step in non "DDP"
		step_model(_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase							)
		# Do "gradient accumulation" (noop)
		if iteration % 2 == 0:
			# Accumulate grads locally
			with accelerator.no_sync(_lowerCAmelCase							):
				step_model(_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase							)
		else:
			# Sync grads
			step_model(_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase							)

		# Since `no_sync` is a noop, `ddp_model` and `model` grads should always be in sync
		check_model_parameters(_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase							)
		for param, ddp_param in zip(model.parameters()       ,ddp_model.parameters()							):
			if not param.requires_grad:
				continue
			assert torch.allclose(
			    param.grad       ,ddp_param.grad							), F'Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})'

		# Shuffle ddp_input on each iteration
		torch.manual_seed(1337 + iteration							)
		lowerCamelCase__								=		ddp_input[torch.randperm(len(_lowerCAmelCase							)							)]





def lowerCAmelCase__(__snake_case							)				->      int:







	'''simple docstring'''







	lowerCamelCase__								=		get_training_setup(_lowerCAmelCase							)
	# Use a single batch
	lowerCamelCase__								=		next(iter(_lowerCAmelCase							)							).values()
	for iteration in range(3							):
		# Gather the distributed inputs and targs for the base model
		lowerCamelCase__								=		accelerator.gather((ddp_input, ddp_target)							)
		lowerCamelCase__								=		input.to(accelerator.device							), target.to(accelerator.device							)
		# Perform our initial ground truth step in non "DDP"
		step_model(_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase							)
		# Do "gradient accumulation" (noop)
		if iteration % 2 == 0:
			# Accumulate grads locally
			with accelerator.no_sync(_lowerCAmelCase							):
				step_model(_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase							)
		else:
			# Sync grads
			step_model(_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase							)

		# DDP model and model should only be in sync when not (iteration % 2 == 0)
		for param, ddp_param in zip(model.parameters()       ,ddp_model.parameters()							):
			if not param.requires_grad:
				continue
			if iteration % 2 == 0:
				# Grads should not be in sync
				assert (
				    torch.allclose(param.grad       ,ddp_param.grad							) is False
				), F'Gradients in sync when they should not be:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})'
			else:
				# Grads should be in sync
				assert (
				    torch.allclose(param.grad       ,ddp_param.grad							) is True
				), F'Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})'

        # Shuffle ddp_input on each iteration
		torch.manual_seed(1337 + iteration							)
		lowerCamelCase__								=		ddp_input[torch.randperm(len(_lowerCAmelCase							)							)]





def lowerCAmelCase__(__snake_case=False       ,__snake_case=False							)				->      Dict:







	'''simple docstring'''







	lowerCamelCase__								=		Accelerator(
	    split_batches=_lowerCAmelCase       ,dispatch_batches=_lowerCAmelCase       ,gradient_accumulation_steps=2							)
	# Test that context manager behaves properly
	lowerCamelCase__								=		get_training_setup(_lowerCAmelCase							)
	for iteration, batch in enumerate(_lowerCAmelCase							):
		lowerCamelCase__								=		batch.values()
		# Gather the distributed inputs and targs for the base model
		lowerCamelCase__								=		accelerator.gather((ddp_input, ddp_target)							)
		lowerCamelCase__								=		input.to(accelerator.device							), target.to(accelerator.device							)
		# Perform our initial ground truth step in non "DDP"
		step_model(_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase							)
		# Do "gradient accumulation" (noop)
		with accelerator.accumulate(_lowerCAmelCase							):
			step_model(_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase							)

		# DDP model and model should only be in sync when not (iteration % 2 == 0)
		for param, ddp_param in zip(model.parameters()       ,ddp_model.parameters()							):
			if not param.requires_grad:
				continue
			if ((iteration + 1) % 2 == 0) or (iteration == len(_lowerCAmelCase							) - 1):
				# Grads should be in sync
				assert (
				    torch.allclose(param.grad       ,ddp_param.grad							) is True
				), F'Gradients not in sync when they should be at iteration {iteration}:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})'
			else:
				# Grads should not be in sync
				assert (
				    torch.allclose(param.grad       ,ddp_param.grad							) is False
				), F'Gradients in sync when they should not be at iteration {iteration}:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})'

        # Shuffle ddp_input on each iteration
		torch.manual_seed(1337 + iteration							)
		lowerCamelCase__								=		ddp_input[torch.randperm(len(_lowerCAmelCase							)							)]
	GradientState._reset_state()





def lowerCAmelCase__(__snake_case=False       ,__snake_case=False							)				->      Optional[int]:







	'''simple docstring'''







	lowerCamelCase__								=		Accelerator(
	    split_batches=_lowerCAmelCase       ,dispatch_batches=_lowerCAmelCase       ,gradient_accumulation_steps=2							)
	# Test that context manager behaves properly
	lowerCamelCase__								=		get_training_setup(_lowerCAmelCase       ,_lowerCAmelCase							)
	for iteration, batch in enumerate(_lowerCAmelCase							):
		lowerCamelCase__								=		batch.values()
		# Gather the distributed inputs and targs for the base model
		lowerCamelCase__								=		accelerator.gather((ddp_input, ddp_target)							)
		lowerCamelCase__								=		input.to(accelerator.device							), target.to(accelerator.device							)
		# Perform our initial ground truth step in non "DDP"
		model.train()
		ddp_model.train()
		step_model(_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase							)
		opt.step()

		if ((iteration + 1) % 2 == 0) or ((iteration + 1) == len(_lowerCAmelCase							)):
			if split_batches:
				sched.step()
			else:
				for _ in range(accelerator.num_processes							):
					sched.step()
		opt.zero_grad()
		# Perform gradient accumulation under wrapper
		with accelerator.accumulate(_lowerCAmelCase							):
			step_model(_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase							)
			ddp_opt.step()
			ddp_sched.step()
			ddp_opt.zero_grad()

		# Learning rates should be the same
		assert (
		    opt.param_groups[0]["lr"] == ddp_opt.param_groups[0]["lr"]
		), F'Learning rates found in each optimizer did not align\nopt: {opt.param_groups[0]["lr"]}\nDDP opt: {ddp_opt.param_groups[0]["lr"]}\n'
		lowerCamelCase__								=		(((iteration + 1) % 2) == 0) or ((iteration + 1) == len(_lowerCAmelCase							))
		if accelerator.num_processes > 1:
			check_model_parameters(_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase							)
		# Shuffle ddp_input on each iteration
		torch.manual_seed(1337 + iteration							)
	GradientState._reset_state()





def lowerCAmelCase__()				->      Tuple:







	'''simple docstring'''







	lowerCamelCase__								=		Accelerator()

	lowerCamelCase__								=		RegressionDataset(length=80							)
	lowerCamelCase__								=		DataLoader(_lowerCAmelCase       ,batch_size=16							)
	lowerCamelCase__								=		RegressionDataset(length=96							)
	lowerCamelCase__								=		DataLoader(_lowerCAmelCase       ,batch_size=16							)
	lowerCamelCase__								=		accelerator.prepare(_lowerCAmelCase       ,_lowerCAmelCase							)
	assert accelerator.gradient_state.active_dataloader is None
	for iteration, _ in enumerate(_lowerCAmelCase							):
		assert id(accelerator.gradient_state.active_dataloader							) == id(_lowerCAmelCase							)
		if iteration < len(_lowerCAmelCase							) - 1:
			assert not accelerator.gradient_state.end_of_dataloader
			if iteration == 1:
				for batch_num, _ in enumerate(_lowerCAmelCase							):
					assert id(accelerator.gradient_state.active_dataloader							) == id(_lowerCAmelCase							)
					if batch_num < len(_lowerCAmelCase							) - 1:
						assert not accelerator.gradient_state.end_of_dataloader
					else:
						assert accelerator.gradient_state.end_of_dataloader
		else:
			assert accelerator.gradient_state.end_of_dataloader
	assert accelerator.gradient_state.active_dataloader is None





def lowerCAmelCase__()				->      Optional[int]:







	'''simple docstring'''







	lowerCamelCase__								=		Accelerator()
	lowerCamelCase__								=		accelerator.state
	if state.local_process_index == 0:
		print('''**Test `accumulate` gradient accumulation with dataloader break**'''							)
	test_dataloader_break()
	if state.distributed_type == DistributedType.NO:
		if state.local_process_index == 0:
			print('''**Test NOOP `no_sync` context manager**'''							)
		test_noop_sync(_lowerCAmelCase							)
	if state.distributed_type in (DistributedType.MULTI_GPU, DistributedType.MULTI_CPU):
		if state.local_process_index == 0:
			print('''**Test Distributed `no_sync` context manager**'''							)
		test_distributed_sync(_lowerCAmelCase							)
	if state.distributed_type == DistributedType.MULTI_GPU:
		for split_batch in [True, False]:
			for dispatch_batches in [True, False]:
				if state.local_process_index == 0:
					print(
					    '''**Test `accumulate` gradient accumulation, '''       ,F'`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**'       ,)
				test_gradient_accumulation(_lowerCAmelCase       ,_lowerCAmelCase							)

    # Currently will break on torch 2.0 +, need to investigate why
	if is_torch_version('''<'''       ,'''2.0'''							) or state.distributed_type == DistributedType.NO:
		if state.local_process_index == 0:
			print(
			    '''**Test `accumulate` gradient accumulation with optimizer and scheduler, '''       ,'''`split_batches=False`, `dispatch_batches=False`**'''       ,)
		test_gradient_accumulation_with_opt_and_scheduler()
		if state.distributed_type == DistributedType.MULTI_GPU:
			for split_batch in [True, False]:
				for dispatch_batches in [True, False]:
					if not split_batch and not dispatch_batches:
						continue
					if state.local_process_index == 0:
						print(
						    '''**Test `accumulate` gradient accumulation with optimizer and scheduler, '''       ,F'`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**'       ,)
					test_gradient_accumulation_with_opt_and_scheduler(_lowerCAmelCase       ,_lowerCAmelCase							)





def lowerCAmelCase__(__snake_case							)				->      Tuple:







	'''simple docstring'''







	main()


if __name__ == "__main__":
			main()








'''simple docstring'''






import gc
import unittest

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

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
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 PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class 				SCREAMING_SNAKE_CASE(							A__							,		A__							,		A__							,		unittest.TestCase			):







	"""simple docstring"""






	lowerCamelCase__			 =       AltDiffusionPipeline
	lowerCamelCase__			 =       TEXT_TO_IMAGE_PARAMS
	lowerCamelCase__			 =       TEXT_TO_IMAGE_BATCH_PARAMS
	lowerCamelCase__			 =       TEXT_TO_IMAGE_IMAGE_PARAMS
	lowerCamelCase__			 =       TEXT_TO_IMAGE_IMAGE_PARAMS







	def      A	(						self			:     Dict		)			->		int:
								torch.manual_seed(0		)
								UpperCAmelCase :       str 							=		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		,			)
								UpperCAmelCase :       Dict 							=		DDIMScheduler(
								    beta_start=0.0_00_85		,			beta_end=0.0_12		,			beta_schedule='''scaled_linear'''		,			clip_sample=__snake_case		,			set_alpha_to_one=__snake_case		,			)
								torch.manual_seed(0		)
								UpperCAmelCase :       Union[str, 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		,			)

								# TODO: address the non-deterministic text encoder (fails for save-load tests)
								# torch.manual_seed(0)
								# text_encoder_config = RobertaSeriesConfig(
								#     hidden_size=32,
								#     project_dim=32,
								#     intermediate_size=37,
								#     layer_norm_eps=1e-05,
								#     num_attention_heads=4,
								#     num_hidden_layers=5,
								#     vocab_size=5002,
								# )
								# text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

								torch.manual_seed(0		)
								UpperCAmelCase :       Dict 							=		CLIPTextConfig(
								    bos_token_id=0		,			eos_token_id=2		,			hidden_size=32		,			projection_dim=32		,			intermediate_size=37		,			layer_norm_eps=1E-05		,			num_attention_heads=4		,			num_hidden_layers=5		,			pad_token_id=1		,			vocab_size=5002		,			)
								UpperCAmelCase :       List[Any] 							=		CLIPTextModel(__snake_case		)

								UpperCAmelCase :       Optional[Any] 							=		XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''		)
								UpperCAmelCase :       Optional[int] 							=		77

								UpperCAmelCase :       Optional[int] 							=		{
								    '''unet''': unet,
								    '''scheduler''': scheduler,
								    '''vae''': vae,
								    '''text_encoder''': text_encoder,
								    '''tokenizer''': tokenizer,
								    '''safety_checker''': None,
								    '''feature_extractor''': None,
								}
								return components







	def      A	(						self			:     Optional[Any]		,			__snake_case			:     Dict		,			__snake_case			:     List[str]=0		)			->		Union[str, Any]:
								if str(__snake_case		).startswith('''mps'''		):
															UpperCAmelCase :       str 							=		torch.manual_seed(__snake_case		)
								else:
															UpperCAmelCase :       Tuple 							=		torch.Generator(device=__snake_case		).manual_seed(__snake_case		)
								UpperCAmelCase :       Dict 							=		{
								    '''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      A	(						self			:     Union[str, Any]		)			->		List[str]:
								super().test_attention_slicing_forward_pass(expected_max_diff=3E-3		)







	def      A	(						self			:     Tuple		)			->		List[str]:
								super().test_inference_batch_single_identical(expected_max_diff=3E-3		)







	def      A	(						self			:     Tuple		)			->		Optional[int]:
								UpperCAmelCase :       List[str] 							=		'''cpu'''  # ensure determinism for the device-dependent torch.Generator

								UpperCAmelCase :       Any 							=		self.get_dummy_components()
								torch.manual_seed(0		)
								UpperCAmelCase :       Optional[int] 							=		RobertaSeriesConfig(
								    hidden_size=32		,			project_dim=32		,			intermediate_size=37		,			layer_norm_eps=1E-05		,			num_attention_heads=4		,			num_hidden_layers=5		,			vocab_size=5002		,			)
								# TODO: remove after fixing the non-deterministic text encoder
								UpperCAmelCase :       List[str] 							=		RobertaSeriesModelWithTransformation(__snake_case		)
								UpperCAmelCase :       str 							=		text_encoder

								UpperCAmelCase :       Optional[int] 							=		AltDiffusionPipeline(**__snake_case		)
								UpperCAmelCase :       str 							=		alt_pipe.to(__snake_case		)
								alt_pipe.set_progress_bar_config(disable=__snake_case		)

								UpperCAmelCase :       Optional[int] 							=		self.get_dummy_inputs(__snake_case		)
								UpperCAmelCase :       Optional[int] 							=		'''A photo of an astronaut'''
								UpperCAmelCase :       List[Any] 							=		alt_pipe(**__snake_case		)
								UpperCAmelCase :       Optional[Any] 							=		output.images
								UpperCAmelCase :       str 							=		image[0, -3:, -3:, -1]

								assert image.shape == (1, 64, 64, 3)
								UpperCAmelCase :       List[str] 							=		np.array(
								    [0.5_74_81_62, 0.60_44_71_45, 0.48_82_12_17, 0.50_10_06_36, 0.5_43_11_85, 0.45_76_36_83, 0.49_65_76_96, 0.48_13_27_33, 0.47_57_30_93]		)

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







	def      A	(						self			:     int		)			->		Union[str, Any]:
								UpperCAmelCase :       Optional[int] 							=		'''cpu'''  # ensure determinism for the device-dependent torch.Generator

								UpperCAmelCase :       int 							=		self.get_dummy_components()
								UpperCAmelCase :       int 							=		PNDMScheduler(skip_prk_steps=__snake_case		)
								torch.manual_seed(0		)
								UpperCAmelCase :       int 							=		RobertaSeriesConfig(
								    hidden_size=32		,			project_dim=32		,			intermediate_size=37		,			layer_norm_eps=1E-05		,			num_attention_heads=4		,			num_hidden_layers=5		,			vocab_size=5002		,			)
								# TODO: remove after fixing the non-deterministic text encoder
								UpperCAmelCase :       Union[str, Any] 							=		RobertaSeriesModelWithTransformation(__snake_case		)
								UpperCAmelCase :       Union[str, Any] 							=		text_encoder
								UpperCAmelCase :       Optional[int] 							=		AltDiffusionPipeline(**__snake_case		)
								UpperCAmelCase :       Dict 							=		alt_pipe.to(__snake_case		)
								alt_pipe.set_progress_bar_config(disable=__snake_case		)

								UpperCAmelCase :       int 							=		self.get_dummy_inputs(__snake_case		)
								UpperCAmelCase :       Optional[int] 							=		alt_pipe(**__snake_case		)
								UpperCAmelCase :       Optional[int] 							=		output.images
								UpperCAmelCase :       int 							=		image[0, -3:, -3:, -1]

								assert image.shape == (1, 64, 64, 3)
								UpperCAmelCase :       Optional[int] 							=		np.array(
								    [0.51_60_50_93, 0.5_70_72_41, 0.47_36_55_07, 0.50_57_88_86, 0.5_63_38_77, 0.4_64_25_03, 0.5_18_20_81, 0.48_76_34_84, 0.49_08_42_37]		)

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

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







	"""simple docstring"""






	def      A	(						self			:     str		)			->		Any:
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()







	def      A	(						self			:     List[Any]		)			->		Any:
								# make sure here that pndm scheduler skips prk
								UpperCAmelCase :       List[Any] 							=		AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''		,			safety_checker=__snake_case		)
								UpperCAmelCase :       Tuple 							=		alt_pipe.to(__snake_case		)
								alt_pipe.set_progress_bar_config(disable=__snake_case		)

								UpperCAmelCase :       List[Any] 							=		'''A painting of a squirrel eating a burger'''
								UpperCAmelCase :       Any 							=		torch.manual_seed(0		)
								UpperCAmelCase :       Optional[int] 							=		alt_pipe([prompt]		,			generator=__snake_case		,			guidance_scale=6.0		,			num_inference_steps=20		,			output_type='''np'''		)

								UpperCAmelCase :       Dict 							=		output.images

								UpperCAmelCase :       Optional[Any] 							=		image[0, -3:, -3:, -1]

								assert image.shape == (1, 512, 512, 3)
								UpperCAmelCase :       List[str] 							=		np.array([0.10_10, 0.08_00, 0.07_94, 0.08_85, 0.08_43, 0.07_62, 0.07_69, 0.07_29, 0.05_86]		)

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







	def      A	(						self			:     Tuple		)			->		int:
								UpperCAmelCase :       int 							=		DDIMScheduler.from_pretrained('''BAAI/AltDiffusion'''		,			subfolder='''scheduler'''		)

								UpperCAmelCase :       Tuple 							=		AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''		,			scheduler=__snake_case		,			safety_checker=__snake_case		)
								UpperCAmelCase :       Dict 							=		alt_pipe.to(__snake_case		)
								alt_pipe.set_progress_bar_config(disable=__snake_case		)

								UpperCAmelCase :       Tuple 							=		'''A painting of a squirrel eating a burger'''
								UpperCAmelCase :       Optional[int] 							=		torch.manual_seed(0		)

								UpperCAmelCase :       List[Any] 							=		alt_pipe([prompt]		,			generator=__snake_case		,			num_inference_steps=2		,			output_type='''numpy'''		)
								UpperCAmelCase :       Dict 							=		output.images

								UpperCAmelCase :       Optional[Any] 							=		image[0, -3:, -3:, -1]

								assert image.shape == (1, 512, 512, 3)
								UpperCAmelCase :       Union[str, Any] 							=		np.array([0.40_19, 0.40_52, 0.38_10, 0.41_19, 0.39_16, 0.39_82, 0.46_51, 0.41_95, 0.53_23]		)

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





