




'''simple docstring'''





import gc
import unittest

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

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


_SCREAMING_SNAKE_CASE							:   Union[str, Any]      		= False




class  _snake_case		(	unittest.TestCase							):




			def    lowerCAmelCase__				(					self							)				->				int:

										'''simple docstring'''




										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()




			@property
			def    lowerCAmelCase__				(					self							)				->				List[Any]:

										'''simple docstring'''




										return 12




			@property
			def    lowerCAmelCase__				(					self							)				->				Tuple:

										'''simple docstring'''




										return 12




			@property
			def    lowerCAmelCase__				(					self							)				->				str:

										'''simple docstring'''




										return 32




			@property
			def    lowerCAmelCase__				(					self							)				->				Optional[Any]:

										'''simple docstring'''




										torch.manual_seed(0							)
										snake_case_						=      VQModel(
										    block_out_channels=[32, 64]				,							in_channels=3				,							out_channels=3				,							down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]				,							up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]				,							latent_channels=3				,							num_vq_embeddings=self.num_embed				,							vq_embed_dim=3				,							)
										return model




			@property
			def    lowerCAmelCase__				(					self							)				->				str:

										'''simple docstring'''




										snake_case_						=      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"							)
										return tokenizer




			@property
			def    lowerCAmelCase__				(					self							)				->				str:

										'''simple docstring'''




										torch.manual_seed(0							)
										snake_case_						=      CLIPTextConfig(
										    bos_token_id=0				,							eos_token_id=2				,							hidden_size=self.text_embedder_hidden_size				,							intermediate_size=37				,							layer_norm_eps=1e-05				,							num_attention_heads=4				,							num_hidden_layers=5				,							pad_token_id=1				,							vocab_size=1_000				,							)
										return CLIPTextModel(a__							)




			@property
			def    lowerCAmelCase__				(					self							)				->				Optional[int]:

										'''simple docstring'''




										torch.manual_seed(0							)

										snake_case_						=      12
										snake_case_						=      12

										snake_case_						=      {
										    "attention_bias": True,
										    "cross_attention_dim": 32,
										    "attention_head_dim": height * width,
										    "num_attention_heads": 1,
										    "num_vector_embeds": self.num_embed,
										    "num_embeds_ada_norm": self.num_embeds_ada_norm,
										    "norm_num_groups": 32,
										    "sample_size": width,
										    "activation_fn": "geglu-approximate",
										}

										snake_case_						=      TransformeraDModel(**a__							)
										return model




			def    lowerCAmelCase__				(					self							)				->				List[Any]:

										'''simple docstring'''




										snake_case_						=      "cpu"

										snake_case_						=      self.dummy_vqvae
										snake_case_						=      self.dummy_text_encoder
										snake_case_						=      self.dummy_tokenizer
										snake_case_						=      self.dummy_transformer
										snake_case_						=      VQDiffusionScheduler(self.num_embed							)
										snake_case_						=      LearnedClassifierFreeSamplingEmbeddings(learnable=a__							)

										snake_case_						=      VQDiffusionPipeline(
										    vqvae=a__				,							text_encoder=a__				,							tokenizer=a__				,							transformer=a__				,							scheduler=a__				,							learned_classifier_free_sampling_embeddings=a__				,							)
										snake_case_						=      pipe.to(a__							)
										pipe.set_progress_bar_config(disable=a__							)

										snake_case_						=      "teddy bear playing in the pool"

										snake_case_						=      torch.Generator(device=a__							).manual_seed(0							)
										snake_case_						=      pipe([prompt]				,							generator=a__				,							num_inference_steps=2				,							output_type="np"							)
										snake_case_						=      output.images

										snake_case_						=      torch.Generator(device=a__							).manual_seed(0							)
										snake_case_						=      pipe(
										    [prompt]				,							generator=a__				,							output_type="np"				,							return_dict=a__				,							num_inference_steps=2							)[0]

										snake_case_						=      image[0, -3:, -3:, -1]
										snake_case_						=      image_from_tuple[0, -3:, -3:, -1]

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

										snake_case_						=      np.array([0.6_5_5_1, 0.6_1_6_8, 0.5_0_0_8, 0.5_6_7_6, 0.5_6_5_9, 0.4_2_9_5, 0.6_0_7_3, 0.5_5_9_9, 0.4_9_9_2]							)

										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							)				->				Dict:

										'''simple docstring'''




										snake_case_						=      "cpu"

										snake_case_						=      self.dummy_vqvae
										snake_case_						=      self.dummy_text_encoder
										snake_case_						=      self.dummy_tokenizer
										snake_case_						=      self.dummy_transformer
										snake_case_						=      VQDiffusionScheduler(self.num_embed							)
										snake_case_						=      LearnedClassifierFreeSamplingEmbeddings(
										    learnable=a__				,							hidden_size=self.text_embedder_hidden_size				,							length=tokenizer.model_max_length							)

										snake_case_						=      VQDiffusionPipeline(
										    vqvae=a__				,							text_encoder=a__				,							tokenizer=a__				,							transformer=a__				,							scheduler=a__				,							learned_classifier_free_sampling_embeddings=a__				,							)
										snake_case_						=      pipe.to(a__							)
										pipe.set_progress_bar_config(disable=a__							)

										snake_case_						=      "teddy bear playing in the pool"

										snake_case_						=      torch.Generator(device=a__							).manual_seed(0							)
										snake_case_						=      pipe([prompt]				,							generator=a__				,							num_inference_steps=2				,							output_type="np"							)
										snake_case_						=      output.images

										snake_case_						=      torch.Generator(device=a__							).manual_seed(0							)
										snake_case_						=      pipe(
										    [prompt]				,							generator=a__				,							output_type="np"				,							return_dict=a__				,							num_inference_steps=2							)[0]

										snake_case_						=      image[0, -3:, -3:, -1]
										snake_case_						=      image_from_tuple[0, -3:, -3:, -1]

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

										snake_case_						=      np.array([0.6_6_9_3, 0.6_0_7_5, 0.4_9_5_9, 0.5_7_0_1, 0.5_5_8_3, 0.4_3_3_3, 0.6_1_7_1, 0.5_6_8_4, 0.4_9_8_8]							)

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




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




			def    lowerCAmelCase__				(					self							)				->				Optional[int]:

										'''simple docstring'''




										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()




			def    lowerCAmelCase__				(					self							)				->				List[Any]:

										'''simple docstring'''




										snake_case_						=      load_numpy(
										    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
										    "/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy"							)

										snake_case_						=      VQDiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq"							)
										snake_case_						=      pipeline.to(a__							)
										pipeline.set_progress_bar_config(disable=a__							)

										# requires GPU generator for gumbel softmax
										# don't use GPU generator in tests though
										snake_case_						=      torch.Generator(device=a__							).manual_seed(0							)
										snake_case_						=      pipeline(
										    "teddy bear playing in the pool"				,							num_images_per_prompt=1				,							generator=a__				,							output_type="np"				,							)

										snake_case_						=      output.images[0]

										assert image.shape == (256, 256, 3)
										assert np.abs(expected_image - image							).max() < 2.0











'''simple docstring'''







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				( lowercase						:      Dict						)	->	Any:
				_a								= filter(lambda lowercase						: p.requires_grad    ,							model.parameters()						)
				_a								= sum([np.prod(p.size()						) for p in model_parameters]						)
				return params


lowerCAmelCase_					:    int      	=			logging.getLogger(__name__)






def  _lowerCamelCase				( lowercase						:      List[Any]    ,							lowercase						:      Any						)	->	Any:
				if metric == "rouge2":
								_a								= "{val_avg_rouge2:.4f}-{step_count}"
				elif metric == "bleu":
								_a								= "{val_avg_bleu:.4f}-{step_count}"
				elif metric == "em":
								_a								= "{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."						)

				_a								= ModelCheckpoint(
				    dirpath=lowercase    ,							filename=lowercase    ,							monitor=F'val_{metric}'    ,							mode="max"    ,							save_top_k=3    ,							every_n_epochs=1    ,							)
				return checkpoint_callback






def  _lowerCamelCase				( lowercase						:      Optional[int]    ,							lowercase						:      Optional[int]						)	->	Union[str, Any]:
				return EarlyStopping(
				    monitor=F'val_{metric}'    ,							mode="min" if "loss" in metric else "max"    ,							patience=lowercase    ,							verbose=lowercase    ,							)



class 				__SCREAMING_SNAKE_CASE  (pl.Callback      ):






				"""simple docstring"""







				def   UpperCamelCase__   (  self     :					Optional[int]					,				__a     :					str					,				__a     :					List[Any]					):
								_a								= {f'lr_group_{i}': param["lr"] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups					)}
								pl_module.logger.log_metrics(__a					)




				@rank_zero_only
				def   UpperCamelCase__   (  self     :					Optional[int]					,				__a     :					pl.Trainer					,				__a     :					pl.LightningModule					,				__a     :					str					,				__a     :					Tuple=True					):
								logger.info(f'***** {type_path} results at step {trainer.global_step:05d} *****'					)
								_a								= 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
								_a								= Path(pl_module.hparams.output_dir					)
								if type_path == "test":
												_a								= od / "test_results.txt"
												_a								= 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.
												_a								= od / f'{type_path}_results/{trainer.global_step:05d}.txt'
												_a								= od / f'{type_path}_generations/{trainer.global_step:05d}.txt'
												results_file.parent.mkdir(exist_ok=__a					)
												generations_file.parent.mkdir(exist_ok=__a					)
								with open(__a					,				"a+"					) as writer:
												for key in sorted(__a					):
																if key in ["log", "progress_bar", "preds"]:
																				continue
																_a								= metrics[key]
																if isinstance(__a					,				torch.Tensor					):
																				_a								= val.item()
																_a								= f'{key}: {val:.6f}\n'
																writer.write(__a					)

								if not save_generations:
												return

								if "preds" in metrics:
												_a								= "\n".join(metrics["preds"]					)
												generations_file.open("w+"					).write(__a					)




				@rank_zero_only
				def   UpperCamelCase__   (  self     :					int					,				__a     :					List[Any]					,				__a     :					Union[str, Any]					):
								try:
												_a								= pl_module.model.model.num_parameters()
								except AttributeError:
												_a								= pl_module.model.num_parameters()

								_a								= count_trainable_parameters(__a					)
								# 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     :					Union[str, Any]					,				__a     :					pl.Trainer					,				__a     :					pl.LightningModule					):
								save_json(pl_module.metrics					,				pl_module.metrics_save_path					)
								return self._write_logs(__a					,				__a					,				"test"					)




				@rank_zero_only
				def   UpperCamelCase__   (  self     :					Any					,				__a     :					pl.Trainer					,				__a     :					int					):
								save_json(pl_module.metrics					,				pl_module.metrics_save_path					)
								# Uncommenting this will save val generations
								# return self._write_logs(trainer, pl_module, "valid")
