



'''simple docstring'''



import inspect
from typing import Callable, List, Optional, Union

import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer

from diffusers import DiffusionPipeline
from diffusers.models import AutoencoderKL, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
from diffusers.utils import logging


lowerCamelCase	 =       logging.get_logger(__name__)  # pylint: disable=invalid-name





class _UpperCamelCase		(     A		):





				'''simple docstring'''





				def __init__(							self  :  Optional[Any]  ,       _lowerCAmelCase  :  AutoencoderKL  ,       _lowerCAmelCase  :  CLIPTextModel  ,       _lowerCAmelCase  :  CLIPTokenizer  ,       _lowerCAmelCase  :  UNetaDConditionModel  ,       _lowerCAmelCase  :  Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]  ,       _lowerCAmelCase  :  StableDiffusionSafetyChecker  ,       _lowerCAmelCase  :  CLIPImageProcessor  ,       ):







											'''simple docstring'''

											super().__init__()
											self.register_modules(
											    vae=_lowerCAmelCase  ,       text_encoder=_lowerCAmelCase  ,       tokenizer=_lowerCAmelCase  ,       unet=_lowerCAmelCase  ,       scheduler=_lowerCAmelCase  ,       safety_checker=_lowerCAmelCase  ,       feature_extractor=_lowerCAmelCase  ,       )

				def 				__lowerCamelCase     (							self  :  str  ,       _lowerCAmelCase  :  Optional[Union[str, int]] = "auto"):







											'''simple docstring'''

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

				def 				__lowerCamelCase     (							self  :  Dict):







											'''simple docstring'''

											self.enable_attention_slicing(_lowerCAmelCase)

				@torch.no_grad()
				def __call__(							self  :  List[Any]  ,       _lowerCAmelCase  :  Union[str, List[str]]  ,       _lowerCAmelCase  :  int = 5_1_2  ,       _lowerCAmelCase  :  int = 5_1_2  ,       _lowerCAmelCase  :  int = 5_0  ,       _lowerCAmelCase  :  float = 7.5  ,       _lowerCAmelCase  :  Optional[Union[str, List[str]]] = None  ,       _lowerCAmelCase  :  Optional[int] = 1  ,       _lowerCAmelCase  :  float = 0.0  ,       _lowerCAmelCase  :  Optional[torch.Generator] = None  ,       _lowerCAmelCase  :  Optional[torch.FloatTensor] = None  ,       _lowerCAmelCase  :  Optional[str] = "pil"  ,       _lowerCAmelCase  :  bool = True  ,       _lowerCAmelCase  :  Optional[Callable[[int, int, torch.FloatTensor], None]] = None  ,       _lowerCAmelCase  :  int = 1  ,       _lowerCAmelCase  :  Optional[torch.FloatTensor] = None  ,       **_lowerCAmelCase  :  Optional[int]  ,       ):







											'''simple docstring'''

											if isinstance(_lowerCAmelCase  ,       _lowerCAmelCase):
																		__lowercase       	=1
											elif isinstance(_lowerCAmelCase  ,       _lowerCAmelCase):
																		__lowercase       	=len(_lowerCAmelCase)
											else:
																		raise ValueError(f"""`prompt` has to be of type `str` or `list` but is {type(_lowerCAmelCase)}""")

											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 (callback_steps is None) or (
											    callback_steps is not None and (not isinstance(_lowerCAmelCase  ,       _lowerCAmelCase) or callback_steps <= 0)
											):
																		raise ValueError(
																		    f"""`callback_steps` has to be a positive integer but is {callback_steps} of type"""
																		    f""" {type(_lowerCAmelCase)}.""")

											# get prompt text embeddings
											__lowercase       	=self.tokenizer(
											    _lowerCAmelCase  ,       padding='max_length'  ,       max_length=self.tokenizer.model_max_length  ,       return_tensors='pt'  ,       )
											__lowercase       	=text_inputs.input_ids

											if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
																		__lowercase       	=self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :])
																		logger.warning(
																		    'The following part of your input was truncated because CLIP can only handle sequences up to'
																		    f""" {self.tokenizer.model_max_length} tokens: {removed_text}""")
																		__lowercase       	=text_input_ids[:, : self.tokenizer.model_max_length]

											if text_embeddings is None:
																		__lowercase       	=self.text_encoder(text_input_ids.to(self.device))[0]

											# duplicate text embeddings for each generation per prompt, using mps friendly method
											__lowercase	,				__lowercase	,				__lowercase       	=text_embeddings.shape
											__lowercase       	=text_embeddings.repeat(1  ,       _lowerCAmelCase  ,       1)
											__lowercase       	=text_embeddings.view(bs_embed * num_images_per_prompt  ,       _lowerCAmelCase  ,       -1)

											# 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.
											__lowercase       	=guidance_scale > 1.0
											# get unconditional embeddings for classifier free guidance
											if do_classifier_free_guidance:
																		__lowercase       	=42
																		if negative_prompt is None:
																									__lowercase       	=['']
																		elif type(_lowerCAmelCase) is not type(_lowerCAmelCase):
																									raise TypeError(
																									    f"""`negative_prompt` should be the same type to `prompt`, but got {type(_lowerCAmelCase)} !="""
																									    f""" {type(_lowerCAmelCase)}.""")
																		elif isinstance(_lowerCAmelCase  ,       _lowerCAmelCase):
																									__lowercase       	=[negative_prompt]
																		elif batch_size != len(_lowerCAmelCase):
																									raise ValueError(
																									    f"""`negative_prompt`: {negative_prompt} has batch size {len(_lowerCAmelCase)}, but `prompt`:"""
																									    f""" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"""
																									    ' the batch size of `prompt`.')
																		else:
																									__lowercase       	=negative_prompt

																		__lowercase       	=text_input_ids.shape[-1]
																		__lowercase       	=self.tokenizer(
																		    _lowerCAmelCase  ,       padding='max_length'  ,       max_length=_lowerCAmelCase  ,       truncation=_lowerCAmelCase  ,       return_tensors='pt'  ,       )
																		__lowercase       	=self.text_encoder(uncond_input.input_ids.to(self.device))[0]

																		# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
																		__lowercase       	=uncond_embeddings.shape[1]
																		__lowercase       	=uncond_embeddings.repeat(_lowerCAmelCase  ,       _lowerCAmelCase  ,       1)
																		__lowercase       	=uncond_embeddings.view(batch_size * num_images_per_prompt  ,       _lowerCAmelCase  ,       -1)

																		# 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
																		__lowercase       	=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`.
											__lowercase       	=(batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
											__lowercase       	=(batch_size * num_images_per_prompt, self.unet.config.in_channels, 6_4, 6_4)
											__lowercase       	=text_embeddings.dtype
											if latents is None:
																		if self.device.type == "mps":
																									# randn does not exist on mps
																									__lowercase       	=torch.randn(
																									    _lowerCAmelCase  ,       generator=_lowerCAmelCase  ,       device='cpu'  ,       dtype=_lowerCAmelCase).to(self.device)
																									__lowercase       	=torch.randn(_lowerCAmelCase  ,       generator=_lowerCAmelCase  ,       device='cpu'  ,       dtype=_lowerCAmelCase).to(
																									    self.device)
																		else:
																									__lowercase       	=torch.randn(
																									    _lowerCAmelCase  ,       generator=_lowerCAmelCase  ,       device=self.device  ,       dtype=_lowerCAmelCase)
																									__lowercase       	=torch.randn(_lowerCAmelCase  ,       generator=_lowerCAmelCase  ,       device=self.device  ,       dtype=_lowerCAmelCase)
											else:
																		if latents_reference.shape != latents_shape:
																									raise ValueError(f"""Unexpected latents shape, got {latents.shape}, expected {latents_shape}""")
																		__lowercase       	=latents_reference.to(self.device)
																		__lowercase       	=latents.to(self.device)

											# This is the key part of the pipeline where we
											# try to ensure that the generated images w/ the same seed
											# but different sizes actually result in similar images
											__lowercase       	=(latents_shape[3] - latents_shape_reference[3]) // 2
											__lowercase       	=(latents_shape[2] - latents_shape_reference[2]) // 2
											__lowercase       	=latents_shape_reference[3] if dx >= 0 else latents_shape_reference[3] + 2 * dx
											__lowercase       	=latents_shape_reference[2] if dy >= 0 else latents_shape_reference[2] + 2 * dy
											__lowercase       	=0 if dx < 0 else dx
											__lowercase       	=0 if dy < 0 else dy
											__lowercase       	=max(-dx  ,       0)
											__lowercase       	=max(-dy  ,       0)
											# import pdb
											# pdb.set_trace()
											__lowercase       	=latents_reference[:, :, dy : dy + h, dx : dx + w]

											# set timesteps
											self.scheduler.set_timesteps(_lowerCAmelCase)

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

											# scale the initial noise by the standard deviation required by the scheduler
											__lowercase       	=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]
											__lowercase       	='eta' in set(inspect.signature(self.scheduler.step).parameters.keys())
											__lowercase       	={}
											if accepts_eta:
																		__lowercase       	=eta

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

																		# predict the noise residual
																		__lowercase       	=self.unet(_lowerCAmelCase  ,       _lowerCAmelCase  ,       encoder_hidden_states=_lowerCAmelCase).sample

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

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

																		# call the callback, if provided
																		if callback is not None and i % callback_steps == 0:
																									callback(_lowerCAmelCase  ,       _lowerCAmelCase  ,       _lowerCAmelCase)

											__lowercase       	=1 / 0.1_8215 * latents
											__lowercase       	=self.vae.decode(_lowerCAmelCase).sample

											__lowercase       	=(image / 2 + 0.5).clamp(0  ,       1)

											# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
											__lowercase       	=image.cpu().permute(0  ,       2  ,       3  ,       1).float().numpy()

											if self.safety_checker is not None:
																		__lowercase       	=self.feature_extractor(self.numpy_to_pil(_lowerCAmelCase)  ,       return_tensors='pt').to(
																		    self.device)
																		__lowercase	,				__lowercase       	=self.safety_checker(
																		    images=_lowerCAmelCase  ,       clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype))
											else:
																		__lowercase       	=None

											if output_type == "pil":
																		__lowercase       	=self.numpy_to_pil(_lowerCAmelCase)

											if not return_dict:
																		return (image, has_nsfw_concept)

											return StableDiffusionPipelineOutput(images=_lowerCAmelCase  ,       nsfw_content_detected=_lowerCAmelCase)








'''simple docstring'''



import os

import pytest
import yaml

from datasets.features.features import Features, Value
from datasets.info import DatasetInfo, DatasetInfosDict



@pytest.mark.parametrize(
    'files'  ,   [
        ['full:README.md', 'dataset_infos.json'],
        ['empty:README.md', 'dataset_infos.json'],
        ['dataset_infos.json'],
        ['full:README.md'],
    ]  ,   )
def   _A       (   _lowerCAmelCase  ,   _lowerCAmelCase					):



							"""simple docstring"""



							__lowercase       	=tmp_path_factory.mktemp('dset_infos_dir'					)
							if "full:README.md" in files:
														with open(dataset_infos_dir / 'README.md'  ,   'w'					) as f:
																					f.write('---\ndataset_info:\n  dataset_size: 42\n---'					)
							if "empty:README.md" in files:
														with open(dataset_infos_dir / 'README.md'  ,   'w'					) as f:
																					f.write(''					)
    # we want to support dataset_infos.json for backward compatibility
							if "dataset_infos.json" in files:
														with open(dataset_infos_dir / 'dataset_infos.json'  ,   'w'					) as f:
																					f.write('{"default": {"dataset_size": 42}}'					)
							__lowercase       	=DatasetInfosDict.from_directory(_lowerCAmelCase					)
							assert dataset_infos
							assert dataset_infos["default"].dataset_size == 42



@pytest.mark.parametrize(
    'dataset_info'  ,   [
        DatasetInfo(),
        DatasetInfo(
            description='foo'  ,   features=Features({'a': Value('int32'					)}					)  ,   builder_name='builder'  ,   config_name='config'  ,   version='1.0.0'  ,   splits=[{'name': 'train'}]  ,   download_size=42  ,   ),
    ]  ,   )
def   _A       (   _lowerCAmelCase  ,   _lowerCAmelCase					):



							"""simple docstring"""



							__lowercase       	=str(_lowerCAmelCase					)
							dataset_info.write_to_directory(_lowerCAmelCase					)
							__lowercase       	=DatasetInfo.from_directory(_lowerCAmelCase					)
							assert dataset_info == reloaded
							assert os.path.exists(os.path.join(_lowerCAmelCase  ,   'dataset_info.json'					)					)



def   _A       (   ):



							"""simple docstring"""



							__lowercase       	=DatasetInfo(
							    description='foo'  ,   citation='bar'  ,   homepage='https://foo.bar'  ,   license='CC0'  ,   features=Features({'a': Value('int32'					)}					)  ,   post_processed={}  ,   supervised_keys=()  ,   task_templates=[]  ,   builder_name='builder'  ,   config_name='config'  ,   version='1.0.0'  ,   splits=[{'name': 'train', 'num_examples': 42}]  ,   download_checksums={}  ,   download_size=1_337  ,   post_processing_size=442  ,   dataset_size=1_234  ,   size_in_bytes=1_337 + 442 + 1_234  ,   )
							__lowercase       	=dataset_info._to_yaml_dict()
							assert sorted(_lowerCAmelCase					) == sorted(DatasetInfo._INCLUDED_INFO_IN_YAML					)
							for key in DatasetInfo._INCLUDED_INFO_IN_YAML:
														assert key in dataset_info_yaml_dict
														assert isinstance(dataset_info_yaml_dict[key]  ,   (list, dict, int, str)					)
							__lowercase       	=yaml.safe_dump(_lowerCAmelCase					)
							__lowercase       	=yaml.safe_load(_lowerCAmelCase					)
							assert dataset_info_yaml_dict == reloaded



def   _A       (   ):



							"""simple docstring"""



							__lowercase       	=DatasetInfo()
							__lowercase       	=dataset_info._to_yaml_dict()
							assert dataset_info_yaml_dict == {}



@pytest.mark.parametrize(
    'dataset_infos_dict'  ,   [
        DatasetInfosDict(),
        DatasetInfosDict({'default': DatasetInfo()}					),
        DatasetInfosDict({'my_config_name': DatasetInfo()}					),
        DatasetInfosDict(
            {
                'default': DatasetInfo(
                    description='foo'  ,   features=Features({'a': Value('int32'					)}					)  ,   builder_name='builder'  ,   config_name='config'  ,   version='1.0.0'  ,   splits=[{'name': 'train'}]  ,   download_size=42  ,   )
            }					),
        DatasetInfosDict(
            {
                'v1': DatasetInfo(dataset_size=42					),
                'v2': DatasetInfo(dataset_size=1_337					),
            }					),
    ]  ,   )
def   _A       (   _lowerCAmelCase  ,   _lowerCAmelCase					):



							"""simple docstring"""



							__lowercase       	=str(_lowerCAmelCase					)
							dataset_infos_dict.write_to_directory(_lowerCAmelCase					)
							__lowercase       	=DatasetInfosDict.from_directory(_lowerCAmelCase					)

							# the config_name of the dataset_infos_dict take over the attribute
							for config_name, dataset_info in dataset_infos_dict.items():
														__lowercase       	=config_name
														# the yaml representation doesn't include fields like description or citation
														# so we just test that we can recover what we can from the yaml
														__lowercase       	=DatasetInfo._from_yaml_dict(dataset_info._to_yaml_dict()					)
							assert dataset_infos_dict == reloaded

							if dataset_infos_dict:
														assert os.path.exists(os.path.join(_lowerCAmelCase  ,   'README.md'					)					)



