


import inspect
from typing import Optional, Union

import numpy as np
import PIL
import torch
from torch.nn import functional as F
from torchvision import transforms
from transformers import CLIPFeatureExtractor, CLIPModel, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    DPMSolverMultistepScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput
from diffusers.utils import (
    PIL_INTERPOLATION,
    randn_tensor,
)

def 			lowerCAmelCase__					(     a__:						Union[str, Any]     ,    a__:						int     ,    a__:						Dict			)    ->						Optional[int]:


							'''simple docstring'''





							if isinstance(a__     ,    torch.Tensor			):
														return image
							elif isinstance(a__     ,    PIL.Image.Image			):
														_UpperCAmelCase											=     [image]

							if isinstance(image[0]     ,    PIL.Image.Image			):
														_UpperCAmelCase											=     [np.array(i.resize((w, h)     ,    resample=PIL_INTERPOLATION['lanczos']			)			)[None, :] for i in image]
														_UpperCAmelCase											=     np.concatenate(a__     ,    axis=0			)
														_UpperCAmelCase											=     np.array(a__			).astype(np.floataa			) / 255.0
														_UpperCAmelCase											=     image.transpose(0     ,    3     ,    1     ,    2			)
														_UpperCAmelCase											=     2.0 * image - 1.0
														_UpperCAmelCase											=     torch.from_numpy(a__			)
							elif isinstance(image[0]     ,    torch.Tensor			):
														_UpperCAmelCase											=     torch.cat(a__     ,    dim=0			)
							return image

def 			lowerCAmelCase__					(     a__:						Any     ,    a__:						Union[str, Any]     ,    a__:						Optional[int]     ,    a__:						Optional[Any]=0.9_995			)    ->						int:


							'''simple docstring'''





							if not isinstance(a__     ,    np.ndarray			):
														_UpperCAmelCase											=     True
														_UpperCAmelCase											=     va.device
														_UpperCAmelCase											=     va.cpu().numpy()
														_UpperCAmelCase											=     va.cpu().numpy()

							_UpperCAmelCase											=     np.sum(va * va / (np.linalg.norm(a__			) * np.linalg.norm(a__			))			)
							if np.abs(a__			) > DOT_THRESHOLD:
														_UpperCAmelCase											=     (1 - t) * va + t * va
							else:
														_UpperCAmelCase											=     np.arccos(a__			)
														_UpperCAmelCase											=     np.sin(a__			)
														_UpperCAmelCase											=     theta_a * t
														_UpperCAmelCase											=     np.sin(a__			)
														_UpperCAmelCase											=     np.sin(theta_a - theta_t			) / sin_theta_a
														_UpperCAmelCase											=     sin_theta_t / sin_theta_a
														_UpperCAmelCase											=     sa * va + sa * va

							if inputs_are_torch:
														_UpperCAmelCase											=     torch.from_numpy(a__			).to(a__			)

							return va

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


							'''simple docstring'''





							_UpperCAmelCase											=     F.normalize(a__     ,    dim=-1			)
							_UpperCAmelCase											=     F.normalize(a__     ,    dim=-1			)
							return (x - y).norm(dim=-1			).div(2			).arcsin().pow(2			).mul(2			)

def 			lowerCAmelCase__					(     a__:						Union[str, Any]     ,    a__:						List[Any]			)    ->						Tuple:


							'''simple docstring'''





							for param in model.parameters():
														_UpperCAmelCase											=     value







class 						__a (							UpperCAmelCase     ):
		def __init__(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE=None     ,       _SCREAMING_SNAKE_CASE=None     ,       _SCREAMING_SNAKE_CASE=None     ,       )					->			Union[str, Any]:




									"""simple docstring"""





									super().__init__()
									self.register_modules(
									    vae=_SCREAMING_SNAKE_CASE     ,       text_encoder=_SCREAMING_SNAKE_CASE     ,       clip_model=_SCREAMING_SNAKE_CASE     ,       tokenizer=_SCREAMING_SNAKE_CASE     ,       unet=_SCREAMING_SNAKE_CASE     ,       scheduler=_SCREAMING_SNAKE_CASE     ,       feature_extractor=_SCREAMING_SNAKE_CASE     ,       coca_model=_SCREAMING_SNAKE_CASE     ,       coca_tokenizer=_SCREAMING_SNAKE_CASE     ,       coca_transform=_SCREAMING_SNAKE_CASE     ,       )
									_UpperCAmelCase											=     (
									    feature_extractor.size
									    if isinstance(feature_extractor.size     ,       _SCREAMING_SNAKE_CASE		)
									    else feature_extractor.size['shortest_edge']
									)
									_UpperCAmelCase											=     transforms.Normalize(mean=feature_extractor.image_mean     ,       std=feature_extractor.image_std		)
									set_requires_grad(self.text_encoder     ,       _SCREAMING_SNAKE_CASE		)
									set_requires_grad(self.clip_model     ,       _SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE = "auto"		)					->			Tuple:




									"""simple docstring"""





									if slice_size == "auto":
																# half the attention head size is usually a good trade-off between
																# speed and memory
																_UpperCAmelCase											=     self.unet.config.attention_head_dim // 2
									self.unet.set_attention_slice(_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			Any:




									"""simple docstring"""





									self.enable_attention_slicing(_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			Union[str, Any]:




									"""simple docstring"""





									set_requires_grad(self.vae     ,       _SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			Tuple:




									"""simple docstring"""





									set_requires_grad(self.vae     ,       _SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			Optional[int]:




									"""simple docstring"""





									set_requires_grad(self.unet     ,       _SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			Optional[int]:




									"""simple docstring"""





									set_requires_grad(self.unet     ,       _SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			List[str]:




									"""simple docstring"""





									_UpperCAmelCase											=     min(int(num_inference_steps * strength		)     ,       _SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     max(num_inference_steps - init_timestep     ,       0		)
									_UpperCAmelCase											=     self.scheduler.timesteps[t_start:]

									return timesteps, num_inference_steps - t_start
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE=None		)					->			Optional[int]:




									"""simple docstring"""





									if not isinstance(_SCREAMING_SNAKE_CASE     ,       torch.Tensor		):
																raise ValueError(f'''`image` has to be of type `torch.Tensor` but is {type(_SCREAMING_SNAKE_CASE		)}'''		)

									_UpperCAmelCase											=     image.to(device=_SCREAMING_SNAKE_CASE     ,       dtype=_SCREAMING_SNAKE_CASE		)

									if isinstance(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		):
																_UpperCAmelCase											=     [
																    self.vae.encode(image[i : i + 1]		).latent_dist.sample(generator[i]		) for i in range(_SCREAMING_SNAKE_CASE		)
																]
																_UpperCAmelCase											=     torch.cat(_SCREAMING_SNAKE_CASE     ,       dim=0		)
									else:
																_UpperCAmelCase											=     self.vae.encode(_SCREAMING_SNAKE_CASE		).latent_dist.sample(_SCREAMING_SNAKE_CASE		)

									# Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
									_UpperCAmelCase											=     0.18215 * init_latents
									_UpperCAmelCase											=     init_latents.repeat_interleave(_SCREAMING_SNAKE_CASE     ,       dim=0		)

									_UpperCAmelCase											=     randn_tensor(init_latents.shape     ,       generator=_SCREAMING_SNAKE_CASE     ,       device=_SCREAMING_SNAKE_CASE     ,       dtype=_SCREAMING_SNAKE_CASE		)

									# get latents
									_UpperCAmelCase											=     self.scheduler.add_noise(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     init_latents

									return latents
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE		)					->			str:




									"""simple docstring"""





									_UpperCAmelCase											=     self.coca_transform(_SCREAMING_SNAKE_CASE		).unsqueeze(0		)
									with torch.no_grad(), torch.cuda.amp.autocast():
																_UpperCAmelCase											=     self.coca_model.generate(transformed_image.to(device=self.device     ,       dtype=self.coca_model.dtype		)		)
									_UpperCAmelCase											=     self.coca_tokenizer.decode(generated[0].cpu().numpy()		)
									return generated.split('<end_of_text>'		)[0].replace('<start_of_text>'     ,       ''		).rstrip(' .,'		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			Union[str, Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     self.feature_extractor.preprocess(_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     torch.from_numpy(clip_image_input['pixel_values'][0]		).unsqueeze(0		).to(self.device		).half()
									_UpperCAmelCase											=     self.clip_model.get_image_features(_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     image_embeddings_clip / image_embeddings_clip.norm(p=2     ,       dim=-1     ,       keepdim=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     image_embeddings_clip.repeat_interleave(_SCREAMING_SNAKE_CASE     ,       dim=0		)
									return image_embeddings_clip
		@torch.enable_grad()
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       )					->			Tuple:




									"""simple docstring"""





									_UpperCAmelCase											=     latents.detach().requires_grad_()

									_UpperCAmelCase											=     self.scheduler.scale_model_input(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

									# predict the noise residual
									_UpperCAmelCase											=     self.unet(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       encoder_hidden_states=_SCREAMING_SNAKE_CASE		).sample

									if isinstance(self.scheduler     ,       (PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)		):
																_UpperCAmelCase											=     self.scheduler.alphas_cumprod[timestep]
																_UpperCAmelCase											=     1 - alpha_prod_t
																# compute predicted original sample from predicted noise also called
																# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
																_UpperCAmelCase											=     (latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5

																_UpperCAmelCase											=     torch.sqrt(_SCREAMING_SNAKE_CASE		)
																_UpperCAmelCase											=     pred_original_sample * (fac) + latents * (1 - fac)
									elif isinstance(self.scheduler     ,       _SCREAMING_SNAKE_CASE		):
																_UpperCAmelCase											=     self.scheduler.sigmas[index]
																_UpperCAmelCase											=     latents - sigma * noise_pred
									else:
																raise ValueError(f'''scheduler type {type(self.scheduler		)} not supported'''		)

									# Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
									_UpperCAmelCase											=     1 / 0.18215 * sample
									_UpperCAmelCase											=     self.vae.decode(_SCREAMING_SNAKE_CASE		).sample
									_UpperCAmelCase											=     (image / 2 + 0.5).clamp(0     ,       1		)

									_UpperCAmelCase											=     transforms.Resize(self.feature_extractor_size		)(_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     self.normalize(_SCREAMING_SNAKE_CASE		).to(latents.dtype		)

									_UpperCAmelCase											=     self.clip_model.get_image_features(_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     image_embeddings_clip / image_embeddings_clip.norm(p=2     ,       dim=-1     ,       keepdim=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     spherical_dist_loss(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		).mean() * clip_guidance_scale

									_UpperCAmelCase											=     -torch.autograd.grad(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)[0]

									if isinstance(self.scheduler     ,       _SCREAMING_SNAKE_CASE		):
																_UpperCAmelCase											=     latents.detach() + grads * (sigma**2)
																_UpperCAmelCase											=     noise_pred_original
									else:
																_UpperCAmelCase											=     noise_pred_original - torch.sqrt(_SCREAMING_SNAKE_CASE		) * grads
									return noise_pred, latents






		@torch.no_grad()
		def __call__(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE = None     ,       _SCREAMING_SNAKE_CASE = None     ,       _SCREAMING_SNAKE_CASE = 512     ,       _SCREAMING_SNAKE_CASE = 512     ,       _SCREAMING_SNAKE_CASE = 0.6     ,       _SCREAMING_SNAKE_CASE = 50     ,       _SCREAMING_SNAKE_CASE = 7.5     ,       _SCREAMING_SNAKE_CASE = 1     ,       _SCREAMING_SNAKE_CASE = 0.0     ,       _SCREAMING_SNAKE_CASE = 100     ,       _SCREAMING_SNAKE_CASE = None     ,       _SCREAMING_SNAKE_CASE = "pil"     ,       _SCREAMING_SNAKE_CASE = True     ,       _SCREAMING_SNAKE_CASE = 0.8     ,       _SCREAMING_SNAKE_CASE = 0.1     ,       _SCREAMING_SNAKE_CASE = 0.1     ,       )					->			int:




									"""simple docstring"""





									if isinstance(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		) and len(_SCREAMING_SNAKE_CASE		) != batch_size:
																raise ValueError(f'''You have passed {batch_size} batch_size, but only {len(_SCREAMING_SNAKE_CASE		)} generators.'''		)

									if height % 8 != 0 or width % 8 != 0:
																raise ValueError(f'''`height` and `width` have to be divisible by 8 but are {height} and {width}.'''		)

									if isinstance(_SCREAMING_SNAKE_CASE     ,       torch.Generator		) and batch_size > 1:
																_UpperCAmelCase											=     [generator] + [None] * (batch_size - 1)

									_UpperCAmelCase											=     [
									    ('model', self.coca_model is None),
									    ('tokenizer', self.coca_tokenizer is None),
									    ('transform', self.coca_transform is None),
									]
									_UpperCAmelCase											=     [x[0] for x in coca_is_none if x[1]]
									_UpperCAmelCase											=     ', '.join(_SCREAMING_SNAKE_CASE		)
									# generate prompts with coca model if prompt is None
									if content_prompt is None:
																if len(_SCREAMING_SNAKE_CASE		):
																							raise ValueError(
																							    f'''Content prompt is None and CoCa [{coca_is_none_str}] is None.'''
																							    f'''Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline.'''		)
																_UpperCAmelCase											=     self.get_image_description(_SCREAMING_SNAKE_CASE		)
									if style_prompt is None:
																if len(_SCREAMING_SNAKE_CASE		):
																							raise ValueError(
																							    f'''Style prompt is None and CoCa [{coca_is_none_str}] is None.'''
																							    f''' Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline.'''		)
																_UpperCAmelCase											=     self.get_image_description(_SCREAMING_SNAKE_CASE		)

									# get prompt text embeddings for content and style
									_UpperCAmelCase											=     self.tokenizer(
									    _SCREAMING_SNAKE_CASE     ,       padding='max_length'     ,       max_length=self.tokenizer.model_max_length     ,       truncation=_SCREAMING_SNAKE_CASE     ,       return_tensors='pt'     ,       )
									_UpperCAmelCase											=     self.text_encoder(content_text_input.input_ids.to(self.device		)		)[0]

									_UpperCAmelCase											=     self.tokenizer(
									    _SCREAMING_SNAKE_CASE     ,       padding='max_length'     ,       max_length=self.tokenizer.model_max_length     ,       truncation=_SCREAMING_SNAKE_CASE     ,       return_tensors='pt'     ,       )
									_UpperCAmelCase											=     self.text_encoder(style_text_input.input_ids.to(self.device		)		)[0]

									_UpperCAmelCase											=     slerp(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

									# duplicate text embeddings for each generation per prompt
									_UpperCAmelCase											=     text_embeddings.repeat_interleave(_SCREAMING_SNAKE_CASE     ,       dim=0		)

									# set timesteps
									_UpperCAmelCase											=     'offset' in set(inspect.signature(self.scheduler.set_timesteps		).parameters.keys()		)
									_UpperCAmelCase											=     {}
									if accepts_offset:
																_UpperCAmelCase											=     1

									self.scheduler.set_timesteps(_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		)
									# Some schedulers like PNDM have timesteps as arrays
									# It's more optimized to move all timesteps to correct device beforehand
									self.scheduler.timesteps.to(self.device		)

									_UpperCAmelCase						,			_UpperCAmelCase											=     self.get_timesteps(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       self.device		)
									_UpperCAmelCase											=     timesteps[:1].repeat(_SCREAMING_SNAKE_CASE		)

									# Preprocess image
									_UpperCAmelCase											=     preprocess(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     self.prepare_latents(
									    _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       text_embeddings.dtype     ,       self.device     ,       _SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     preprocess(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     self.prepare_latents(
									    _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       text_embeddings.dtype     ,       self.device     ,       _SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     slerp(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

									if clip_guidance_scale > 0:
																_UpperCAmelCase											=     self.get_clip_image_embeddings(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
																_UpperCAmelCase											=     self.get_clip_image_embeddings(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
																_UpperCAmelCase											=     slerp(
																    _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

									# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
									# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
									# corresponds to doing no classifier free guidance.
									_UpperCAmelCase											=     guidance_scale > 1.0
									# get unconditional embeddings for classifier free guidance
									if do_classifier_free_guidance:
																_UpperCAmelCase											=     content_text_input.input_ids.shape[-1]
																_UpperCAmelCase											=     self.tokenizer(['']     ,       padding='max_length'     ,       max_length=_SCREAMING_SNAKE_CASE     ,       return_tensors='pt'		)
																_UpperCAmelCase											=     self.text_encoder(uncond_input.input_ids.to(self.device		)		)[0]
																# duplicate unconditional embeddings for each generation per prompt
																_UpperCAmelCase											=     uncond_embeddings.repeat_interleave(_SCREAMING_SNAKE_CASE     ,       dim=0		)

																# For classifier free guidance, we need to do two forward passes.
																# Here we concatenate the unconditional and text embeddings into a single batch
																# to avoid doing two forward passes
																_UpperCAmelCase											=     torch.cat([uncond_embeddings, text_embeddings]		)

									# get the initial random noise unless the user supplied it

									# Unlike in other pipelines, latents need to be generated in the target device
									# for 1-to-1 results reproducibility with the CompVis implementation.
									# However this currently doesn't work in `mps`.
									_UpperCAmelCase											=     (batch_size, self.unet.config.in_channels, height // 8, width // 8)
									_UpperCAmelCase											=     text_embeddings.dtype
									if latents is None:
																if self.device.type == "mps":
																							# randn does not work reproducibly on mps
																							_UpperCAmelCase											=     torch.randn(_SCREAMING_SNAKE_CASE     ,       generator=_SCREAMING_SNAKE_CASE     ,       device='cpu'     ,       dtype=_SCREAMING_SNAKE_CASE		).to(
																							    self.device		)
																else:
																							_UpperCAmelCase											=     torch.randn(_SCREAMING_SNAKE_CASE     ,       generator=_SCREAMING_SNAKE_CASE     ,       device=self.device     ,       dtype=_SCREAMING_SNAKE_CASE		)
									else:
																if latents.shape != latents_shape:
																							raise ValueError(f'''Unexpected latents shape, got {latents.shape}, expected {latents_shape}'''		)
																_UpperCAmelCase											=     latents.to(self.device		)

									# scale the initial noise by the standard deviation required by the scheduler
									_UpperCAmelCase											=     latents * self.scheduler.init_noise_sigma

									# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
									# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
									# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
									# and should be between [0, 1]
									_UpperCAmelCase											=     'eta' in set(inspect.signature(self.scheduler.step		).parameters.keys()		)
									_UpperCAmelCase											=     {}
									if accepts_eta:
																_UpperCAmelCase											=     eta

									# check if the scheduler accepts generator
									_UpperCAmelCase											=     'generator' in set(inspect.signature(self.scheduler.step		).parameters.keys()		)
									if accepts_generator:
																_UpperCAmelCase											=     generator

									with self.progress_bar(total=_SCREAMING_SNAKE_CASE		):
																for i, t in enumerate(_SCREAMING_SNAKE_CASE		):
																							# expand the latents if we are doing classifier free guidance
																							_UpperCAmelCase											=     torch.cat([latents] * 2		) if do_classifier_free_guidance else latents
																							_UpperCAmelCase											=     self.scheduler.scale_model_input(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

																							# predict the noise residual
																							_UpperCAmelCase											=     self.unet(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       encoder_hidden_states=_SCREAMING_SNAKE_CASE		).sample

																							# perform classifier free guidance
																							if do_classifier_free_guidance:
																														_UpperCAmelCase						,			_UpperCAmelCase											=     noise_pred.chunk(2		)
																														_UpperCAmelCase											=     noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

																							# perform clip guidance
																							if clip_guidance_scale > 0:
																														_UpperCAmelCase											=     (
																														    text_embeddings.chunk(2		)[1] if do_classifier_free_guidance else text_embeddings
																														)
																														_UpperCAmelCase						,			_UpperCAmelCase											=     self.cond_fn(
																														    _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       )

																							# compute the previous noisy sample x_t -> x_t-1
																							_UpperCAmelCase											=     self.scheduler.step(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		).prev_sample

        # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
									_UpperCAmelCase											=     1 / 0.18215 * latents
									_UpperCAmelCase											=     self.vae.decode(_SCREAMING_SNAKE_CASE		).sample

									_UpperCAmelCase											=     (image / 2 + 0.5).clamp(0     ,       1		)
									_UpperCAmelCase											=     image.cpu().permute(0     ,       2     ,       3     ,       1		).numpy()

									if output_type == "pil":
																_UpperCAmelCase											=     self.numpy_to_pil(_SCREAMING_SNAKE_CASE		)

									if not return_dict:
																return (image, None)

									return StableDiffusionPipelineOutput(images=_SCREAMING_SNAKE_CASE     ,       nsfw_content_detected=_SCREAMING_SNAKE_CASE		)




import random
import unittest

from torch.utils.data import BatchSampler, DataLoader, IterableDataset

from accelerate import Accelerator
from accelerate.data_loader import (
    BatchSamplerShard,
    DataLoaderDispatcher,
    DataLoaderShard,
    IterableDatasetShard,
    SkipBatchSampler,
    SkipDataLoader,
    skip_first_batches,
)

class 						__a (							UpperCAmelCase     ):
		def __init__(     self     ,       _SCREAMING_SNAKE_CASE=0.01     ,       _SCREAMING_SNAKE_CASE=1000		)					->			str:




									"""simple docstring"""





									_UpperCAmelCase											=     p_stop
									_UpperCAmelCase											=     max_length






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




									"""simple docstring"""





									_UpperCAmelCase											=     0
									_UpperCAmelCase											=     False
									while not stop and count < self.max_length:
																yield count
																count += 1
																_UpperCAmelCase											=     random.random() < self.p_stop



class 						__a (							unittest.TestCase     ):
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE=False     ,       _SCREAMING_SNAKE_CASE=True		)					->			int:




									"""simple docstring"""





									_UpperCAmelCase											=     [
									    BatchSamplerShard(_SCREAMING_SNAKE_CASE     ,       2     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)
									    for i in range(2		)
									]
									_UpperCAmelCase											=     [list(_SCREAMING_SNAKE_CASE		) for batch_sampler_shard in batch_sampler_shards]
									if not split_batches:
																self.assertListEqual([len(_SCREAMING_SNAKE_CASE		) for shard in batch_sampler_shards]     ,       [len(_SCREAMING_SNAKE_CASE		) for e in expected]		)
									self.assertListEqual(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     BatchSampler(range(24		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(24		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									# Expected shouldn't change
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

									# Check the shards when the dataset is a round multiple of batch size but not total batch size.
									_UpperCAmelCase											=     BatchSampler(range(21		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [0, 1, 2]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(21		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

									# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
									# num_processes batch.
									_UpperCAmelCase											=     BatchSampler(range(22		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 0, 1]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(22		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

									# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
									# num_processes batch.
									_UpperCAmelCase											=     BatchSampler(range(20		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 0]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [1, 2, 3]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(20		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

									# Check the shards when the dataset is very small.
									_UpperCAmelCase											=     BatchSampler(range(2		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [[[0, 1, 0]], [[1, 0, 1]]]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(2		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [[], []]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			Union[str, Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     BatchSampler(range(24		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(24		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									# Expected shouldn't change
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)

									# Check the shards when the dataset is not a round multiple of batch size.
									_UpperCAmelCase											=     BatchSampler(range(22		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [0, 1]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(22		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)

									# Check the shards when the dataset is not a round multiple of batch size or num_processes.
									_UpperCAmelCase											=     BatchSampler(range(21		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 0]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [1, 2]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(21		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)

									# Check the shards when the dataset is very small.
									_UpperCAmelCase											=     BatchSampler(range(2		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [[[0, 1]], [[0, 1]]]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(2		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [[], []]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			Tuple:




									"""simple docstring"""





									_UpperCAmelCase											=     BatchSampler(range(24		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(24		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									# Expected shouldn't change
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

									# Check the shards when the dataset is a round multiple of batch size but not total batch size.
									_UpperCAmelCase											=     BatchSampler(range(21		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(21		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

									# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
									# num_processes batch.
									_UpperCAmelCase											=     BatchSampler(range(22		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(22		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

									# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
									# num_processes batch.
									_UpperCAmelCase											=     BatchSampler(range(20		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(20		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

									# Check the shards when the dataset is very small.
									_UpperCAmelCase											=     BatchSampler(range(2		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [[[0, 1]], []]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(2		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [[], []]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     BatchSampler(range(24		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(24		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									# Expected shouldn't change
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

									# Check the shards when the dataset is not a round multiple of batch size.
									_UpperCAmelCase											=     BatchSampler(range(22		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(22		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

									# Check the shards when the dataset is not a round multiple of batch size or num_processes.
									_UpperCAmelCase											=     BatchSampler(range(21		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(21		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

									# Check the shards when the dataset is very small.
									_UpperCAmelCase											=     BatchSampler(range(2		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [[[0, 1]], []]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(2		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [[], []]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
									_UpperCAmelCase											=     [BatchSamplerShard(_SCREAMING_SNAKE_CASE     ,       2     ,       _SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		) for i in range(2		)]

									self.assertEqual(len(batch_sampler_shards[0]		)     ,       3		)
									self.assertEqual(len(batch_sampler_shards[1]		)     ,       2		)

									self.assertListEqual(list(batch_sampler_shards[0]		)     ,       [[0, 1, 2], [5, 6, 7, 8], [12, 13]]		)
									self.assertListEqual(list(batch_sampler_shards[1]		)     ,       [[3, 4], [9, 10, 11]]		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE=False     ,       _SCREAMING_SNAKE_CASE=2     ,       _SCREAMING_SNAKE_CASE=False		)					->			Dict:




									"""simple docstring"""





									random.seed(_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     list(_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     [
									    IterableDatasetShard(
									        _SCREAMING_SNAKE_CASE     ,       batch_size=_SCREAMING_SNAKE_CASE     ,       drop_last=_SCREAMING_SNAKE_CASE     ,       num_processes=_SCREAMING_SNAKE_CASE     ,       process_index=_SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE     ,       )
									    for i in range(_SCREAMING_SNAKE_CASE		)
									]
									_UpperCAmelCase											=     []
									for iterable_dataset_shard in iterable_dataset_shards:
																# Since our random iterable dataset will be... random... we need to use a seed to get reproducible results.
																random.seed(_SCREAMING_SNAKE_CASE		)
																iterable_dataset_lists.append(list(_SCREAMING_SNAKE_CASE		)		)

									_UpperCAmelCase											=     batch_size // num_processes if split_batches else batch_size
									# All iterable dataset shard should have the same length, a round multiple of shard_batch_size
									_UpperCAmelCase											=     iterable_dataset_lists[0]
									for l in iterable_dataset_lists[1:]:
																self.assertEqual(len(_SCREAMING_SNAKE_CASE		)     ,       len(_SCREAMING_SNAKE_CASE		)		)
																self.assertTrue(len(_SCREAMING_SNAKE_CASE		) % shard_batch_size == 0		)

									_UpperCAmelCase											=     []
									for idx in range(0     ,       len(_SCREAMING_SNAKE_CASE		)     ,       _SCREAMING_SNAKE_CASE		):
																for l in iterable_dataset_lists:
																							observed += l[idx : idx + shard_batch_size]

									if not drop_last:
																while len(_SCREAMING_SNAKE_CASE		) < len(_SCREAMING_SNAKE_CASE		):
																							reference += reference
									self.assertListEqual(_SCREAMING_SNAKE_CASE     ,       reference[: len(_SCREAMING_SNAKE_CASE		)]		)
		def       UpperCAmelCase__							(     self		)					->			List[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     42
									_UpperCAmelCase											=     RandomIterableDataset()

									self.check_iterable_dataset_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)
									self.check_iterable_dataset_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)
									self.check_iterable_dataset_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)
									self.check_iterable_dataset_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)

									# Edge case with a very small dataset
									_UpperCAmelCase											=     RandomIterableDataset(max_length=2		)

									self.check_iterable_dataset_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)
									self.check_iterable_dataset_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)
									self.check_iterable_dataset_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)
									self.check_iterable_dataset_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			str:




									"""simple docstring"""





									_UpperCAmelCase											=     BatchSampler(range(16		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     SkipBatchSampler(_SCREAMING_SNAKE_CASE     ,       2		)
									self.assertListEqual(list(_SCREAMING_SNAKE_CASE		)     ,       [[8, 9, 10, 11], [12, 13, 14, 15]]		)
		def       UpperCAmelCase__							(     self		)					->			Tuple:




									"""simple docstring"""





									_UpperCAmelCase											=     SkipDataLoader(list(range(16		)		)     ,       batch_size=4     ,       skip_batches=2		)
									self.assertListEqual([t.tolist() for t in dataloader]     ,       [[8, 9, 10, 11], [12, 13, 14, 15]]		)
		def       UpperCAmelCase__							(     self		)					->			int:




									"""simple docstring"""





									_UpperCAmelCase											=     DataLoader(list(range(16		)		)     ,       batch_size=4		)
									_UpperCAmelCase											=     skip_first_batches(_SCREAMING_SNAKE_CASE     ,       num_batches=2		)
									self.assertListEqual([t.tolist() for t in new_dataloader]     ,       [[8, 9, 10, 11], [12, 13, 14, 15]]		)
		def       UpperCAmelCase__							(     self		)					->			int:




									"""simple docstring"""





									_UpperCAmelCase											=     DataLoaderShard(list(range(16		)		)     ,       batch_size=4		)
									for idx, _ in enumerate(_SCREAMING_SNAKE_CASE		):
																self.assertEqual(dataloader.end_of_dataloader     ,       idx == 3		)

									# Test it also works on the second iteration
									for idx, _ in enumerate(_SCREAMING_SNAKE_CASE		):
																self.assertEqual(dataloader.end_of_dataloader     ,       idx == 3		)






		def       UpperCAmelCase__							(     self		)					->			Any:




									"""simple docstring"""





									Accelerator()
									_UpperCAmelCase											=     DataLoaderDispatcher(range(16		)     ,       batch_size=4		)
									for idx, _ in enumerate(_SCREAMING_SNAKE_CASE		):
																self.assertEqual(dataloader.end_of_dataloader     ,       idx == 3		)

									# Test it also works on the second iteration
									for idx, _ in enumerate(_SCREAMING_SNAKE_CASE		):
																self.assertEqual(dataloader.end_of_dataloader     ,       idx == 3		)
