





"""simple docstring"""





import argparse

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType


########################################################################
# This is a fully working simple example to use Accelerate
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


lowerCAmelCase__   	=	16
lowerCAmelCase__   	=	32


def 					a__    (    _SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE = 16			):




		"""simple docstring"""


		UpperCamelCase      =		AutoTokenizer.from_pretrained("bert-base-cased"			)
		UpperCamelCase      =		load_dataset("glue"     ,		"mrpc"			)

		def tokenize_function(_SCREAMING_SNAKE_CASE			):
				# max_length=None => use the model max length (it's actually the default)
				UpperCamelCase      =		tokenizer(examples["sentence1"]     ,		examples["sentence2"]     ,		truncation=_SCREAMING_SNAKE_CASE     ,		max_length=_SCREAMING_SNAKE_CASE			)
				return outputs

		# Apply the method we just defined to all the examples in all the splits of the dataset
		# starting with the main process first:
		with accelerator.main_process_first():
				UpperCamelCase      =		datasets.map(
				    _SCREAMING_SNAKE_CASE     ,		batched=_SCREAMING_SNAKE_CASE     ,		remove_columns=["idx", "sentence1", "sentence2"]     ,		)

		# We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
		# transformers library
		UpperCamelCase      =		tokenized_datasets.rename_column("label"     ,		"labels"			)

		def collate_fn(_SCREAMING_SNAKE_CASE			):
				# On TPU it's best to pad everything to the same length or training will be very slow.
				UpperCamelCase      =		128 if accelerator.distributed_type == DistributedType.TPU else None
				# When using mixed precision we want round multiples of 8/16
				if accelerator.mixed_precision == "fp8":
						UpperCamelCase      =		16
				elif accelerator.mixed_precision != "no":
						UpperCamelCase      =		8
				else:
						UpperCamelCase      =		None

				return tokenizer.pad(
				    _SCREAMING_SNAKE_CASE     ,		padding="longest"     ,		max_length=_SCREAMING_SNAKE_CASE     ,		pad_to_multiple_of=_SCREAMING_SNAKE_CASE     ,		return_tensors="pt"     ,		)

		# Instantiate dataloaders.
		UpperCamelCase      =		DataLoader(
		    tokenized_datasets["train"]     ,		shuffle=_SCREAMING_SNAKE_CASE     ,		collate_fn=_SCREAMING_SNAKE_CASE     ,		batch_size=_SCREAMING_SNAKE_CASE     ,		drop_last=_SCREAMING_SNAKE_CASE			)
		UpperCamelCase      =		DataLoader(
		    tokenized_datasets["validation"]     ,		shuffle=_SCREAMING_SNAKE_CASE     ,		collate_fn=_SCREAMING_SNAKE_CASE     ,		batch_size=_SCREAMING_SNAKE_CASE     ,		drop_last=(accelerator.mixed_precision == "fp8")     ,		)

		return train_dataloader, eval_dataloader


def 					a__    (    _SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE			):




		"""simple docstring"""


		UpperCamelCase      =		Accelerator(cpu=args.cpu     ,		mixed_precision=args.mixed_precision			)
		# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
		UpperCamelCase      =		config["lr"]
		UpperCamelCase      =		int(config["num_epochs"]			)
		UpperCamelCase      =		int(config["seed"]			)
		UpperCamelCase      =		int(config["batch_size"]			)

		UpperCamelCase      =		evaluate.load("glue"     ,		"mrpc"			)

		# If the batch size is too big we use gradient accumulation
		UpperCamelCase      =		1
		if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
				UpperCamelCase      =		batch_size // MAX_GPU_BATCH_SIZE
				UpperCamelCase      =		MAX_GPU_BATCH_SIZE

		set_seed(_SCREAMING_SNAKE_CASE			)
		UpperCamelCase      , UpperCamelCase      =		get_dataloaders(_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE			)
		# Instantiate the model (we build the model here so that the seed also control new weights initialization)
		UpperCamelCase      =		AutoModelForSequenceClassification.from_pretrained("bert-base-cased"     ,		return_dict=_SCREAMING_SNAKE_CASE			)

		# We could avoid this line since the accelerator is set with `device_placement=True` (default value).
		# Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
		# creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
		UpperCamelCase      =		model.to(accelerator.device			)
		# Instantiate optimizer
		UpperCamelCase      =		AdamW(params=model.parameters()     ,		lr=_SCREAMING_SNAKE_CASE			)

		# Instantiate scheduler
		UpperCamelCase      =		get_linear_schedule_with_warmup(
		    optimizer=_SCREAMING_SNAKE_CASE     ,		num_warmup_steps=100     ,		num_training_steps=(len(_SCREAMING_SNAKE_CASE			) * num_epochs) // gradient_accumulation_steps     ,		)

		# Prepare everything
		# There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
		# prepare method.

		UpperCamelCase      , UpperCamelCase      , UpperCamelCase      , UpperCamelCase      , UpperCamelCase      =		accelerator.prepare(
		    _SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE			)

		# Now we train the model
		for epoch in range(_SCREAMING_SNAKE_CASE			):
				model.train()
				for step, batch in enumerate(_SCREAMING_SNAKE_CASE			):
						# We could avoid this line since we set the accelerator with `device_placement=True`.
						batch.to(accelerator.device			)
						UpperCamelCase      =		model(**_SCREAMING_SNAKE_CASE			)
						UpperCamelCase      =		outputs.loss
						UpperCamelCase      =		loss / gradient_accumulation_steps
						accelerator.backward(_SCREAMING_SNAKE_CASE			)
						if step % gradient_accumulation_steps == 0:
								optimizer.step()
								lr_scheduler.step()
								optimizer.zero_grad()

				model.eval()
				for step, batch in enumerate(_SCREAMING_SNAKE_CASE			):
						# We could avoid this line since we set the accelerator with `device_placement=True`.
						batch.to(accelerator.device			)
						with torch.no_grad():
								UpperCamelCase      =		model(**_SCREAMING_SNAKE_CASE			)
						UpperCamelCase      =		outputs.logits.argmax(dim=-1			)
						UpperCamelCase      , UpperCamelCase      =		accelerator.gather_for_metrics((predictions, batch["labels"])			)
						metric.add_batch(
						    predictions=_SCREAMING_SNAKE_CASE     ,		references=_SCREAMING_SNAKE_CASE     ,		)

				UpperCamelCase      =		metric.compute()
				# Use accelerator.print to print only on the main process.
				accelerator.print(F"epoch {epoch}:"     ,		_SCREAMING_SNAKE_CASE			)


def 					a__    (    ):




		"""simple docstring"""


		UpperCamelCase      =		argparse.ArgumentParser(description="Simple example of training script."			)
		parser.add_argument(
		    "--mixed_precision"     ,		type=_SCREAMING_SNAKE_CASE     ,		default=_SCREAMING_SNAKE_CASE     ,		choices=["no", "fp16", "bf16", "fp8"]     ,		help="Whether to use mixed precision. Choose"
		    "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
		    "and an Nvidia Ampere GPU."     ,		)
		parser.add_argument("--cpu"     ,		action="store_true"     ,		help="If passed, will train on the CPU."			)
		UpperCamelCase      =		parser.parse_args()
		UpperCamelCase      =		{"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
		training_function(_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE			)


if __name__ == "__main__":
						main()










import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaImgaImgPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
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

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()


class        _a							(			UpperCamelCase__     ,	unittest.TestCase			):
							_lowercase	:			Union[str, Any]        = KandinskyVaaImgaImgPipeline
							_lowercase	:			Tuple        = ['''image_embeds''', '''negative_image_embeds''', '''image''']
							_lowercase	:			Any        = [
							    '''image_embeds''',
							    '''negative_image_embeds''',
							    '''image''',
							]
							_lowercase	:			Union[str, Any]        = [
							    '''generator''',
							    '''height''',
							    '''width''',
							    '''strength''',
							    '''guidance_scale''',
							    '''num_inference_steps''',
							    '''return_dict''',
							    '''guidance_scale''',
							    '''num_images_per_prompt''',
							    '''output_type''',
							    '''return_dict''',
							]
							_lowercase	:			Optional[Any]        = False

							@property
							def       lowerCamelCase_ (							self:	Union[str, Any]							)			->   Dict:





									"""simple docstring"""







									return 32

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





									"""simple docstring"""







									return 32

							@property
							def       lowerCamelCase_ (							self:	Any							)			->   Any:





									"""simple docstring"""







									return self.time_input_dim

							@property
							def       lowerCamelCase_ (							self:	Tuple							)			->   Any:





									"""simple docstring"""







									return self.time_input_dim * 4

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





									"""simple docstring"""







									return 100

							@property
							def       lowerCamelCase_ (							self:	int							)			->   int:





									"""simple docstring"""







									torch.manual_seed(0							)

									lowercase__            =      {
									    '''in_channels''': 4,
									    # Out channels is double in channels because predicts mean and variance
									    '''out_channels''': 8,
									    '''addition_embed_type''': '''image''',
									    '''down_block_types''': ('''ResnetDownsampleBlock2D''', '''SimpleCrossAttnDownBlock2D'''),
									    '''up_block_types''': ('''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D'''),
									    '''mid_block_type''': '''UNetMidBlock2DSimpleCrossAttn''',
									    '''block_out_channels''': (self.block_out_channels_a, self.block_out_channels_a * 2),
									    '''layers_per_block''': 1,
									    '''encoder_hid_dim''': self.text_embedder_hidden_size,
									    '''encoder_hid_dim_type''': '''image_proj''',
									    '''cross_attention_dim''': self.cross_attention_dim,
									    '''attention_head_dim''': 4,
									    '''resnet_time_scale_shift''': '''scale_shift''',
									    '''class_embed_type''': None,
									}

									lowercase__            =      UNetaDConditionModel(**UpperCamelCase_							)
									return model

							@property
							def       lowerCamelCase_ (							self:	Optional[int]							)			->   Union[str, Any]:





									"""simple docstring"""







									return {
									    "block_out_channels": [32, 64],
									    "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
									    "in_channels": 3,
									    "latent_channels": 4,
									    "layers_per_block": 1,
									    "norm_num_groups": 8,
									    "norm_type": "spatial",
									    "num_vq_embeddings": 12,
									    "out_channels": 3,
									    "up_block_types": [
									        "AttnUpDecoderBlock2D",
									        "UpDecoderBlock2D",
									    ],
									    "vq_embed_dim": 4,
									}

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





									"""simple docstring"""







									torch.manual_seed(0							)
									lowercase__            =      VQModel(**self.dummy_movq_kwargs							)
									return model

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





									"""simple docstring"""







									lowercase__            =      self.dummy_unet
									lowercase__            =      self.dummy_movq

									lowercase__            =      {
									    '''num_train_timesteps''': 1_000,
									    '''beta_schedule''': '''linear''',
									    '''beta_start''': 0.00085,
									    '''beta_end''': 0.012,
									    '''clip_sample''': False,
									    '''set_alpha_to_one''': False,
									    '''steps_offset''': 0,
									    '''prediction_type''': '''epsilon''',
									    '''thresholding''': False,
									}

									lowercase__            =      DDIMScheduler(**UpperCamelCase_							)

									lowercase__            =      {
									    '''unet''': unet,
									    '''scheduler''': scheduler,
									    '''movq''': movq,
									}

									return components

							def       lowerCamelCase_ (							self:	Optional[int]    ,						UpperCamelCase_:	Optional[Any]    ,						UpperCamelCase_:	Optional[int]=0							)			->   Optional[int]:





									"""simple docstring"""







									lowercase__            =      floats_tensor((1, self.text_embedder_hidden_size)    ,						rng=random.Random(UpperCamelCase_							)							).to(UpperCamelCase_							)
									lowercase__            =      floats_tensor((1, self.text_embedder_hidden_size)    ,						rng=random.Random(seed + 1							)							).to(
									    UpperCamelCase_							)
									# create init_image
									lowercase__            =      floats_tensor((1, 3, 64, 64)    ,						rng=random.Random(UpperCamelCase_							)							).to(UpperCamelCase_							)
									lowercase__            =      image.cpu().permute(0    ,						2    ,						3    ,						1							)[0]
									lowercase__            =      Image.fromarray(np.uinta(UpperCamelCase_							)							).convert('''RGB'''							).resize((256, 256)							)

									if str(UpperCamelCase_							).startswith('''mps'''							):
											lowercase__            =      torch.manual_seed(UpperCamelCase_							)
									else:
											lowercase__            =      torch.Generator(device=UpperCamelCase_							).manual_seed(UpperCamelCase_							)
									lowercase__            =      {
									    '''image''': init_image,
									    '''image_embeds''': image_embeds,
									    '''negative_image_embeds''': negative_image_embeds,
									    '''generator''': generator,
									    '''height''': 64,
									    '''width''': 64,
									    '''num_inference_steps''': 10,
									    '''guidance_scale''': 7.0,
									    '''strength''': 0.2,
									    '''output_type''': '''np''',
									}
									return inputs


							def       lowerCamelCase_ (							self:	Optional[int]							)			->   Dict:





									"""simple docstring"""







									lowercase__            =      '''cpu'''

									lowercase__            =      self.get_dummy_components()

									lowercase__            =      self.pipeline_class(**UpperCamelCase_							)
									lowercase__            =      pipe.to(UpperCamelCase_							)

									pipe.set_progress_bar_config(disable=UpperCamelCase_							)

									lowercase__            =      pipe(**self.get_dummy_inputs(UpperCamelCase_							)							)
									lowercase__            =      output.images

									lowercase__            =      pipe(
									    **self.get_dummy_inputs(UpperCamelCase_							)    ,						return_dict=UpperCamelCase_    ,						)[0]

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

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

									lowercase__            =      np.array(
									    [0.6199778, 0.63984406, 0.46145785, 0.62944984, 0.5622215, 0.47306132, 0.47441456, 0.4607606, 0.48719263]							)
									assert (
									    np.abs(image_slice.flatten() - expected_slice							).max() < 1E-2
									), f' expected_slice {expected_slice}, but got {image_slice.flatten()}'
									assert (
									    np.abs(image_from_tuple_slice.flatten() - expected_slice							).max() < 1E-2
									), f' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'


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

							def       lowerCamelCase_ (							self:	str							)			->   List[Any]:





									"""simple docstring"""







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


							def       lowerCamelCase_ (							self:	List[str]							)			->   Union[str, Any]:





									"""simple docstring"""







									lowercase__            =      load_numpy(
									    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
									    '''/kandinskyv22/kandinskyv22_img2img_frog.npy'''							)

									lowercase__            =      load_image(
									    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/kandinsky/cat.png'''							)
									lowercase__            =      '''A red cartoon frog, 4k'''

									lowercase__            =      KandinskyVaaPriorPipeline.from_pretrained(
									    '''kandinsky-community/kandinsky-2-2-prior'''    ,						torch_dtype=torch.floataa							)
									pipe_prior.to(UpperCamelCase_							)

									lowercase__            =      KandinskyVaaImgaImgPipeline.from_pretrained(
									    '''kandinsky-community/kandinsky-2-2-decoder'''    ,						torch_dtype=torch.floataa							)
									lowercase__            =      pipeline.to(UpperCamelCase_							)

									pipeline.set_progress_bar_config(disable=UpperCamelCase_							)

									lowercase__            =      torch.Generator(device='''cpu'''							).manual_seed(0							)
									lowercase__     ,		lowercase__            =      pipe_prior(
									    UpperCamelCase_    ,						generator=UpperCamelCase_    ,						num_inference_steps=5    ,						negative_prompt=''''''    ,						).to_tuple()

									lowercase__            =      pipeline(
									    image=UpperCamelCase_    ,						image_embeds=UpperCamelCase_    ,						negative_image_embeds=UpperCamelCase_    ,						generator=UpperCamelCase_    ,						num_inference_steps=100    ,						height=768    ,						width=768    ,						strength=0.2    ,						output_type='''np'''    ,						)

									lowercase__            =      output.images[0]

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

									assert_mean_pixel_difference(UpperCamelCase_    ,						UpperCamelCase_							)


