


import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizer

from diffusers import (
    AltDiffusionImgaImgPipeline,
    AutoencoderKL,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()





class  _a							(    unittest.TestCase):




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




	@property
	def 							UpperCAmelCase__( self						:		Optional[int]		)->	List[Any]:
			lowerCAmelCase__       :      Union[str, Any]   					= 1
			lowerCAmelCase__       :      Any   					= 3
			lowerCAmelCase__       :      str   					= (32, 32)

			lowerCAmelCase__       :      Dict   					= floats_tensor((batch_size, num_channels) + sizes		, rng=random.Random(0		)		).to(_SCREAMING_SNAKE_CASE		)
			return image




	@property
	def 							UpperCAmelCase__( self						:		Optional[Any]		)->	str:
			torch.manual_seed(0		)
			lowerCAmelCase__       :      List[Any]   					= 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		, )
			return model




	@property
	def 							UpperCAmelCase__( self						:		int		)->	List[str]:
			torch.manual_seed(0		)
			lowerCAmelCase__       :      Optional[int]   					= 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		, )
			return model




	@property
	def 							UpperCAmelCase__( self						:		Dict		)->	int:
			torch.manual_seed(0		)
			lowerCAmelCase__       :      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		, pad_token_id=1		, vocab_size=5006		, )
			return RobertaSeriesModelWithTransformation(_SCREAMING_SNAKE_CASE		)




	@property
	def 							UpperCAmelCase__( self						:		Optional[int]		)->	Tuple:
			def extract(*_SCREAMING_SNAKE_CASE						:		Tuple		, **_SCREAMING_SNAKE_CASE						:		Optional[Any]		):





					class  _a							:




						def __init__( self						:		List[Any]		)->	Optional[Any]:
								lowerCAmelCase__       :      Union[str, Any]   					= torch.ones([0]		)




						def 							UpperCAmelCase__( self						:		Optional[int]		, _SCREAMING_SNAKE_CASE						:		Tuple		)->	int:
								self.pixel_values.to(_SCREAMING_SNAKE_CASE		)
								return self
					return Out()

			return extract




	def 							UpperCAmelCase__( self						:		List[Any]		)->	str:
			lowerCAmelCase__       :      List[str]   					= '''cpu'''  # ensure determinism for the device-dependent torch.Generator
			lowerCAmelCase__       :      int   					= self.dummy_cond_unet
			lowerCAmelCase__       :      Union[str, Any]   					= PNDMScheduler(skip_prk_steps=_SCREAMING_SNAKE_CASE		)
			lowerCAmelCase__       :      Union[str, Any]   					= self.dummy_vae
			lowerCAmelCase__       :      Union[str, Any]   					= self.dummy_text_encoder
			lowerCAmelCase__       :      Union[str, Any]   					= XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''		)
			lowerCAmelCase__       :      List[str]   					= 77

			lowerCAmelCase__       :      Tuple   					= self.dummy_image.to(_SCREAMING_SNAKE_CASE		)
			lowerCAmelCase__       :      List[str]   					= init_image / 2 + 0.5

			# make sure here that pndm scheduler skips prk
			lowerCAmelCase__       :      List[str]   					= AltDiffusionImgaImgPipeline(
			    unet=_SCREAMING_SNAKE_CASE		, scheduler=_SCREAMING_SNAKE_CASE		, vae=_SCREAMING_SNAKE_CASE		, text_encoder=_SCREAMING_SNAKE_CASE		, tokenizer=_SCREAMING_SNAKE_CASE		, safety_checker=_SCREAMING_SNAKE_CASE		, feature_extractor=self.dummy_extractor		, )
			lowerCAmelCase__       :      str   					= VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor		, do_normalize=_SCREAMING_SNAKE_CASE		)
			lowerCAmelCase__       :      Dict   					= alt_pipe.to(_SCREAMING_SNAKE_CASE		)
			alt_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

			lowerCAmelCase__       :      str   					= '''A painting of a squirrel eating a burger'''
			lowerCAmelCase__       :      Any   					= torch.Generator(device=_SCREAMING_SNAKE_CASE		).manual_seed(0		)
			lowerCAmelCase__       :      List[str]   					= alt_pipe(
			    [prompt]		, generator=_SCREAMING_SNAKE_CASE		, guidance_scale=6.0		, num_inference_steps=2		, output_type='''np'''		, image=_SCREAMING_SNAKE_CASE		, )

			lowerCAmelCase__       :      str   					= output.images

			lowerCAmelCase__       :      Dict   					= torch.Generator(device=_SCREAMING_SNAKE_CASE		).manual_seed(0		)
			lowerCAmelCase__       :      Any   					= alt_pipe(
			    [prompt]		, generator=_SCREAMING_SNAKE_CASE		, guidance_scale=6.0		, num_inference_steps=2		, output_type='''np'''		, image=_SCREAMING_SNAKE_CASE		, return_dict=_SCREAMING_SNAKE_CASE		, )[0]

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

			assert image.shape == (1, 32, 32, 3)
			lowerCAmelCase__       :      Union[str, Any]   					= np.array([0.4427, 0.3731, 0.4249, 0.4941, 0.4546, 0.4148, 0.4193, 0.4666, 0.4499]		)

			assert np.abs(image_slice.flatten() - expected_slice		).max() < 5E-3
			assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 5E-3




	@unittest.skipIf(torch_device != '''cuda'''		, '''This test requires a GPU'''		)
	def 							UpperCAmelCase__( self						:		List[Any]		)->	Any:
			lowerCAmelCase__       :      Union[str, Any]   					= self.dummy_cond_unet
			lowerCAmelCase__       :      Tuple   					= PNDMScheduler(skip_prk_steps=_SCREAMING_SNAKE_CASE		)
			lowerCAmelCase__       :      Tuple   					= self.dummy_vae
			lowerCAmelCase__       :      Any   					= self.dummy_text_encoder
			lowerCAmelCase__       :      Union[str, Any]   					= XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''		)
			lowerCAmelCase__       :      Tuple   					= 77

			lowerCAmelCase__       :      int   					= self.dummy_image.to(_SCREAMING_SNAKE_CASE		)

			# put models in fp16
			lowerCAmelCase__       :      Optional[Any]   					= unet.half()
			lowerCAmelCase__       :      Optional[int]   					= vae.half()
			lowerCAmelCase__       :      int   					= bert.half()

			# make sure here that pndm scheduler skips prk
			lowerCAmelCase__       :      List[Any]   					= AltDiffusionImgaImgPipeline(
			    unet=_SCREAMING_SNAKE_CASE		, scheduler=_SCREAMING_SNAKE_CASE		, vae=_SCREAMING_SNAKE_CASE		, text_encoder=_SCREAMING_SNAKE_CASE		, tokenizer=_SCREAMING_SNAKE_CASE		, safety_checker=_SCREAMING_SNAKE_CASE		, feature_extractor=self.dummy_extractor		, )
			lowerCAmelCase__       :      Optional[Any]   					= VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor		, do_normalize=_SCREAMING_SNAKE_CASE		)
			lowerCAmelCase__       :      Union[str, Any]   					= alt_pipe.to(_SCREAMING_SNAKE_CASE		)
			alt_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

			lowerCAmelCase__       :      Dict   					= '''A painting of a squirrel eating a burger'''
			lowerCAmelCase__       :      Tuple   					= torch.manual_seed(0		)
			lowerCAmelCase__       :      List[Any]   					= alt_pipe(
			    [prompt]		, generator=_SCREAMING_SNAKE_CASE		, num_inference_steps=2		, output_type='''np'''		, image=_SCREAMING_SNAKE_CASE		, ).images

			assert image.shape == (1, 32, 32, 3)




	@unittest.skipIf(torch_device != '''cuda'''		, '''This test requires a GPU'''		)
	def 							UpperCAmelCase__( self						:		Tuple		)->	str:
			lowerCAmelCase__       :      Dict   					= load_image(
			    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
			    '''/img2img/sketch-mountains-input.jpg'''		)
			# resize to resolution that is divisible by 8 but not 16 or 32
			lowerCAmelCase__       :      Dict   					= init_image.resize((760, 504)		)

			lowerCAmelCase__       :      Optional[int]   					= '''BAAI/AltDiffusion'''
			lowerCAmelCase__       :      Optional[Any]   					= AltDiffusionImgaImgPipeline.from_pretrained(
			    _SCREAMING_SNAKE_CASE		, safety_checker=_SCREAMING_SNAKE_CASE		, )
			pipe.to(_SCREAMING_SNAKE_CASE		)
			pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)
			pipe.enable_attention_slicing()

			lowerCAmelCase__       :      str   					= '''A fantasy landscape, trending on artstation'''

			lowerCAmelCase__       :      Dict   					= torch.manual_seed(0		)
			lowerCAmelCase__       :      str   					= pipe(
			    prompt=_SCREAMING_SNAKE_CASE		, image=_SCREAMING_SNAKE_CASE		, strength=0.75		, guidance_scale=7.5		, generator=_SCREAMING_SNAKE_CASE		, output_type='''np'''		, )
			lowerCAmelCase__       :      Optional[int]   					= output.images[0]

			lowerCAmelCase__       :      Union[str, Any]   					= image[255:258, 383:386, -1]

			assert image.shape == (504, 760, 3)
			lowerCAmelCase__       :      List[str]   					= np.array([0.9358, 0.9397, 0.9599, 0.9901, 1.0000, 1.0000, 0.9882, 1.0000, 1.0000]		)

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





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




	def 							UpperCAmelCase__( self						:		Union[str, Any]		)->	Tuple:
			# clean up the VRAM after each test
			super().tearDown()
			gc.collect()
			torch.cuda.empty_cache()




	def 							UpperCAmelCase__( self						:		Union[str, Any]		)->	str:
			lowerCAmelCase__       :      Union[str, Any]   					= load_image(
			    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
			    '''/img2img/sketch-mountains-input.jpg'''		)
			lowerCAmelCase__       :      Tuple   					= init_image.resize((768, 512)		)
			lowerCAmelCase__       :      Tuple   					= load_numpy(
			    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy'''		)

			lowerCAmelCase__       :      Union[str, Any]   					= '''BAAI/AltDiffusion'''
			lowerCAmelCase__       :      Any   					= AltDiffusionImgaImgPipeline.from_pretrained(
			    _SCREAMING_SNAKE_CASE		, safety_checker=_SCREAMING_SNAKE_CASE		, )
			pipe.to(_SCREAMING_SNAKE_CASE		)
			pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)
			pipe.enable_attention_slicing()

			lowerCAmelCase__       :      List[str]   					= '''A fantasy landscape, trending on artstation'''

			lowerCAmelCase__       :      int   					= torch.manual_seed(0		)
			lowerCAmelCase__       :      Dict   					= pipe(
			    prompt=_SCREAMING_SNAKE_CASE		, image=_SCREAMING_SNAKE_CASE		, strength=0.75		, guidance_scale=7.5		, generator=_SCREAMING_SNAKE_CASE		, output_type='''np'''		, )
			lowerCAmelCase__       :      Optional[int]   					= output.images[0]

			assert image.shape == (512, 768, 3)
			# img2img is flaky across GPUs even in fp32, so using MAE here
			assert np.abs(expected_image - image		).max() < 1E-2








import logging
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.utilities import rank_zero_only
from utils_rag import save_json






def 					lowerCamelCase_   (	_a	):




		"""simple docstring"""






		lowerCAmelCase__       :      Any   					= filter(lambda _a	: p.requires_grad     ,       model.parameters()	)
		lowerCAmelCase__       :      str   					= sum([np.prod(p.size()	) for p in model_parameters]	)
		return params


lowerCamelCase						=		logging.getLogger(__name__)






def 					lowerCamelCase_   (	_a     ,       _a	):




		"""simple docstring"""






		if metric == "rouge2":
				lowerCAmelCase__       :      Optional[int]   					= '''{val_avg_rouge2:.4f}-{step_count}'''
		elif metric == "bleu":
				lowerCAmelCase__       :      Optional[int]   					= '''{val_avg_bleu:.4f}-{step_count}'''
		elif metric == "em":
				lowerCAmelCase__       :      List[Any]   					= '''{val_avg_em:.4f}-{step_count}'''
		else:
				raise NotImplementedError(
				    f'seq2seq callbacks only support rouge2 and bleu, got {metric}, You can make your own by adding to this'
				    ''' function.'''	)

		lowerCAmelCase__       :      Dict   					= ModelCheckpoint(
		    dirpath=_a     ,       filename=_a     ,       monitor=f'val_{metric}'     ,       mode='''max'''     ,       save_top_k=3     ,       every_n_epochs=1     ,       )
		return checkpoint_callback






def 					lowerCamelCase_   (	_a     ,       _a	):




		"""simple docstring"""






		return EarlyStopping(
		    monitor=f'val_{metric}'     ,       mode='''min''' if '''loss''' in metric else '''max'''     ,       patience=_a     ,       verbose=_a     ,       )





class  _a							(    pl.Callback):




	def 							UpperCAmelCase__( self						:		Optional[Any]		, _SCREAMING_SNAKE_CASE						:		int		, _SCREAMING_SNAKE_CASE						:		Any		)->	Optional[int]:
			lowerCAmelCase__       :      Dict   					= {F'lr_group_{i}': param['''lr'''] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups		)}
			pl_module.logger.log_metrics(_SCREAMING_SNAKE_CASE		)




	@rank_zero_only
	def 							UpperCAmelCase__( self						:		List[str]		, _SCREAMING_SNAKE_CASE						:		pl.Trainer		, _SCREAMING_SNAKE_CASE						:		pl.LightningModule		, _SCREAMING_SNAKE_CASE						:		str		, _SCREAMING_SNAKE_CASE						:		List[str]=True		)->	None:
			logger.info(F'***** {type_path} results at step {trainer.global_step:05d} *****'		)
			lowerCAmelCase__       :      List[Any]   					= trainer.callback_metrics
			trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ['''log''', '''progress_bar''', '''preds''']}		)
			# Log results
			lowerCAmelCase__       :      List[Any]   					= Path(pl_module.hparams.output_dir		)
			if type_path == "test":
					lowerCAmelCase__       :      Optional[int]   					= od / '''test_results.txt'''
					lowerCAmelCase__       :      Tuple   					= od / '''test_generations.txt'''
			else:
					# this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
					# If people want this it will be easy enough to add back.
					lowerCAmelCase__       :      int   					= od / F'{type_path}_results/{trainer.global_step:05d}.txt'
					lowerCAmelCase__       :      int   					= od / F'{type_path}_generations/{trainer.global_step:05d}.txt'
					results_file.parent.mkdir(exist_ok=_SCREAMING_SNAKE_CASE		)
					generations_file.parent.mkdir(exist_ok=_SCREAMING_SNAKE_CASE		)
			with open(_SCREAMING_SNAKE_CASE		, '''a+'''		) as writer:
					for key in sorted(_SCREAMING_SNAKE_CASE		):
							if key in ["log", "progress_bar", "preds"]:
									continue
							lowerCAmelCase__       :      Optional[int]   					= metrics[key]
							if isinstance(_SCREAMING_SNAKE_CASE		, torch.Tensor		):
									lowerCAmelCase__       :      List[str]   					= val.item()
							lowerCAmelCase__       :      List[str]   					= F'{key}: {val:.6f}\n'
							writer.write(_SCREAMING_SNAKE_CASE		)

			if not save_generations:
					return

			if "preds" in metrics:
					lowerCAmelCase__       :      Dict   					= '''\n'''.join(metrics['''preds''']		)
					generations_file.open('''w+'''		).write(_SCREAMING_SNAKE_CASE		)




	@rank_zero_only
	def 							UpperCAmelCase__( self						:		str		, _SCREAMING_SNAKE_CASE						:		Union[str, Any]		, _SCREAMING_SNAKE_CASE						:		Optional[int]		)->	Optional[int]:
			try:
					lowerCAmelCase__       :      Tuple   					= pl_module.model.model.num_parameters()
			except AttributeError:
					lowerCAmelCase__       :      Optional[Any]   					= pl_module.model.num_parameters()

			lowerCAmelCase__       :      Dict   					= count_trainable_parameters(_SCREAMING_SNAKE_CASE		)
			# mp stands for million parameters
			trainer.logger.log_metrics({'''n_params''': npars, '''mp''': npars / 1E6, '''grad_mp''': n_trainable_pars / 1E6}		)




	@rank_zero_only
	def 							UpperCAmelCase__( self						:		List[Any]		, _SCREAMING_SNAKE_CASE						:		pl.Trainer		, _SCREAMING_SNAKE_CASE						:		pl.LightningModule		)->	Optional[Any]:
			save_json(pl_module.metrics		, pl_module.metrics_save_path		)
			return self._write_logs(_SCREAMING_SNAKE_CASE		, _SCREAMING_SNAKE_CASE		, '''test'''		)




	@rank_zero_only
	def 							UpperCAmelCase__( self						:		Optional[int]		, _SCREAMING_SNAKE_CASE						:		pl.Trainer		, _SCREAMING_SNAKE_CASE						:		List[Any]		)->	List[Any]:
			save_json(pl_module.metrics		, pl_module.metrics_save_path		)
			# Uncommenting this will save val generations
			# return self._write_logs(trainer, pl_module, "valid")




