

from typing import List, Optional, Union

import torch

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


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

UpperCAmelCase_						='\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline\n        >>> import torch\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-2-prior")\n        >>> pipe_prior.to("cuda")\n        >>> prompt = "red cat, 4k photo"\n        >>> out = pipe_prior(prompt)\n        >>> image_emb = out.image_embeds\n        >>> zero_image_emb = out.negative_image_embeds\n        >>> pipe = KandinskyV22Pipeline.from_pretrained("kandinsky-community/kandinsky-2-2-decoder")\n        >>> pipe.to("cuda")\n        >>> image = pipe(\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=50,\n        ... ).images\n        >>> image[0].save("cat.png")\n        ```\n'



def       UpperCAmelCase       (     _snake_case						, _snake_case						, _snake_case=8							):
							lowerCAmelCase   =		height // scale_factor**2
							if height % scale_factor**2 != 0:
														new_height += 1
							lowerCAmelCase   =		width // scale_factor**2
							if width % scale_factor**2 != 0:
														new_width += 1
							return new_height * scale_factor, new_width * scale_factor


class 							__UpperCamelCase			( __lowercase  ):





				'''simple docstring'''



				def __init__( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , ):
											super().__init__()

											self.register_modules(
											    unet=__a       , scheduler=__a       , movq=__a       , )
											lowerCAmelCase   =		2 ** (len(self.movq.config.block_out_channels						) - 1)



				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						):
											if latents is None:
																		lowerCAmelCase   =		randn_tensor(__a       , generator=__a       , device=__a       , dtype=__a						)
											else:
																		if latents.shape != shape:
																									raise ValueError(F"""Unexpected latents shape, got {latents.shape}, expected {shape}"""						)
																		lowerCAmelCase   =		latents.to(__a						)

											lowerCAmelCase   =		latents * scheduler.init_noise_sigma
											return latents



				def       __snake_case		( self       , UpperCAmelCase_=0						):
											if is_accelerate_available():
																		from accelerate import cpu_offload
											else:
																		raise ImportError('''Please install accelerate via `pip install accelerate`'''						)

											lowerCAmelCase   =		torch.device(F"""cuda:{gpu_id}"""						)

											lowerCAmelCase   =		[
											    self.unet,
											    self.movq,
											]
											for cpu_offloaded_model in models:
																		if cpu_offloaded_model is not None:
																									cpu_offload(__a       , __a						)



				def       __snake_case		( self       , UpperCAmelCase_=0						):
											if is_accelerate_available() and is_accelerate_version('''>='''       , '''0.17.0.dev0'''						):
																		from accelerate import cpu_offload_with_hook
											else:
																		raise ImportError('''`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'''						)

											lowerCAmelCase   =		torch.device(F"""cuda:{gpu_id}"""						)

											if self.device.type != "cpu":
																		self.to('''cpu'''       , silence_dtype_warnings=__a						)
																		torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

											lowerCAmelCase   =		None
											for cpu_offloaded_model in [self.unet, self.movq]:
																		lowerCAmelCase   =		cpu_offload_with_hook(__a       , __a       , prev_module_hook=__a						)

											# We'll offload the last model manually.
											lowerCAmelCase   =		hook



				@property
				# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
				def       __snake_case		( self						):
											if not hasattr(self.unet       , '''_hf_hook'''						):
																		return self.device
											for module in self.unet.modules():
																		if (
																		    hasattr(__a       , '''_hf_hook'''						)
																		    and hasattr(module._hf_hook       , '''execution_device'''						)
																		    and module._hf_hook.execution_device is not None
																		):
																									return torch.device(module._hf_hook.execution_device						)
											return self.device






				@torch.no_grad()
				@replace_example_docstring(__a						)
				def __call__( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_ = 5_12       , UpperCAmelCase_ = 5_12       , UpperCAmelCase_ = 1_00       , UpperCAmelCase_ = 4.0       , UpperCAmelCase_ = 1       , UpperCAmelCase_ = None       , UpperCAmelCase_ = None       , UpperCAmelCase_ = "pil"       , UpperCAmelCase_ = True       , ):
											lowerCAmelCase   =		self._execution_device

											lowerCAmelCase   =		guidance_scale > 1.0

											if isinstance(__a       , __a						):
																		lowerCAmelCase   =		torch.cat(__a       , dim=0						)
											lowerCAmelCase   =		image_embeds.shape[0] * num_images_per_prompt
											if isinstance(__a       , __a						):
																		lowerCAmelCase   =		torch.cat(__a       , dim=0						)

											if do_classifier_free_guidance:
																		lowerCAmelCase   =		image_embeds.repeat_interleave(__a       , dim=0						)
																		lowerCAmelCase   =		negative_image_embeds.repeat_interleave(__a       , dim=0						)

											lowerCAmelCase   =		torch.cat([negative_image_embeds, image_embeds]       , dim=0						).to(dtype=self.unet.dtype       , device=__a						)

											self.scheduler.set_timesteps(__a       , device=__a						)
											lowerCAmelCase   =		self.scheduler.timesteps

											lowerCAmelCase   =		self.unet.config.in_channels

											lowerCAmelCase   =		downscale_height_and_width(__a       , __a       , self.movq_scale_factor						)

											# create initial latent
											lowerCAmelCase   =		self.prepare_latents(
											    (batch_size, num_channels_latents, height, width)       , image_embeds.dtype       , __a       , __a       , __a       , self.scheduler       , )

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

																		lowerCAmelCase   =		{"""image_embeds""": image_embeds}
																		lowerCAmelCase   =		self.unet(
																		    sample=__a       , timestep=__a       , encoder_hidden_states=__a       , added_cond_kwargs=__a       , return_dict=__a       , )[0]

																		if do_classifier_free_guidance:
																									lowerCAmelCase   =		noise_pred.split(latents.shape[1]       , dim=1						)
																									lowerCAmelCase   =		noise_pred.chunk(2						)
																									lowerCAmelCase   =		variance_pred.chunk(2						)
																									lowerCAmelCase   =		noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
																									lowerCAmelCase   =		torch.cat([noise_pred, variance_pred_text]       , dim=1						)

																		if not (
																		    hasattr(self.scheduler.config       , '''variance_type'''						)
																		    and self.scheduler.config.variance_type in ["learned", "learned_range"]
																		):
																									lowerCAmelCase   =		noise_pred.split(latents.shape[1]       , dim=1						)

																		# compute the previous noisy sample x_t -> x_t-1
																		lowerCAmelCase   =		self.scheduler.step(
																		    __a       , __a       , __a       , generator=__a       , )[0]
											# post-processing
											lowerCAmelCase   =		self.movq.decode(__a       , force_not_quantize=__a						)["""sample"""]

											if output_type not in ["pt", "np", "pil"]:
																		raise ValueError(F"""Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}"""						)

											if output_type in ["np", "pil"]:
																		lowerCAmelCase   =		image * 0.5 + 0.5
																		lowerCAmelCase   =		image.clamp(0       , 1						)
																		lowerCAmelCase   =		image.cpu().permute(0       , 2       , 3       , 1						).float().numpy()

											if output_type == "pil":
																		lowerCAmelCase   =		self.numpy_to_pil(__a						)

											if not return_dict:
																		return (image,)

											return ImagePipelineOutput(images=__a						)









from collections import Counter
from pathlib import Path
from typing import Optional, Tuple

import yaml


class 							__UpperCamelCase			( yaml.SafeLoader  ):





				'''simple docstring'''



				def       __snake_case		( self       , UpperCAmelCase_						):
											lowerCAmelCase   =		[self.constructed_objects[key_node] for key_node, _ in node.value]
											lowerCAmelCase   =		[tuple(UpperCAmelCase_						) if isinstance(UpperCAmelCase_       , UpperCAmelCase_						) else key for key in keys]
											lowerCAmelCase   =		Counter(UpperCAmelCase_						)
											lowerCAmelCase   =		[key for key in counter if counter[key] > 1]
											if duplicate_keys:
																		raise TypeError(F"""Got duplicate yaml keys: {duplicate_keys}"""						)






				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_=False						):
											lowerCAmelCase   =		super().construct_mapping(UpperCAmelCase_       , deep=UpperCAmelCase_						)
											self._check_no_duplicates_on_constructed_node(UpperCAmelCase_						)
											return mapping




def       UpperCAmelCase       (     _snake_case							):
							lowerCAmelCase   =		list(readme_content.splitlines()							)
							if full_content and full_content[0] == "---" and "---" in full_content[1:]:
														lowerCAmelCase   =		full_content[1:].index('''---'''							) + 1
														lowerCAmelCase   =		'''\n'''.join(full_content[1:sep_idx]							)
														return yamlblock, "\n".join(full_content[sep_idx + 1 :]							)

							return None, "\n".join(_snake_case							)


class 							__UpperCamelCase			( __UpperCAmelCase  ):





				'''simple docstring'''



				__a							: Any										={"""train_eval_index"""}  # train-eval-index in the YAML metadata



				@classmethod
				def       __snake_case		( cls       , UpperCAmelCase_						):
											with open(UpperCAmelCase_       , encoding='''utf-8'''						) as readme_file:
																		lowerCAmelCase   ,     lowerCAmelCase   =		_split_yaml_from_readme(readme_file.read()						)
											if yaml_string is not None:
																		return cls.from_yaml_string(UpperCAmelCase_						)
											else:
																		return cls()



				def       __snake_case		( self       , UpperCAmelCase_						):
											if path.exists():
																		with open(UpperCAmelCase_       , encoding='''utf-8'''						) as readme_file:
																									lowerCAmelCase   =		readme_file.read()
											else:
																		lowerCAmelCase   =		None
											lowerCAmelCase   =		self._to_readme(UpperCAmelCase_						)
											with open(UpperCAmelCase_       , '''w'''       , encoding='''utf-8'''						) as readme_file:
																		readme_file.write(UpperCAmelCase_						)



				def       __snake_case		( self       , UpperCAmelCase_ = None						):
											if readme_content is not None:
																		lowerCAmelCase   ,     lowerCAmelCase   =		_split_yaml_from_readme(UpperCAmelCase_						)
																		lowerCAmelCase   =		'''---\n''' + self.to_yaml_string() + '''---\n''' + content
											else:
																		lowerCAmelCase   =		'''---\n''' + self.to_yaml_string() + '''---\n'''
											return full_content



				@classmethod
				def       __snake_case		( cls       , UpperCAmelCase_						):
											lowerCAmelCase   =		yaml.load(UpperCAmelCase_       , Loader=_NoDuplicateSafeLoader						) or {}

											# Convert the YAML keys to DatasetMetadata fields
											lowerCAmelCase   =		{
											    (key.replace('''-'''       , '''_'''						) if key.replace('''-'''       , '''_'''						) in cls._FIELDS_WITH_DASHES else key): value
											    for key, value in metadata_dict.items()
											}
											return cls(**UpperCAmelCase_						)






				def       __snake_case		( self						):
											return yaml.safe_dump(
											    {
											        (key.replace('''_'''       , '''-'''						) if key in self._FIELDS_WITH_DASHES else key): value
											        for key, value in self.items()
											    }       , sort_keys=UpperCAmelCase_       , allow_unicode=UpperCAmelCase_       , encoding='''utf-8'''       , ).decode('''utf-8'''						)




UpperCAmelCase_						={
    """image-classification""": [],
    """translation""": [],
    """image-segmentation""": [],
    """fill-mask""": [],
    """automatic-speech-recognition""": [],
    """token-classification""": [],
    """sentence-similarity""": [],
    """audio-classification""": [],
    """question-answering""": [],
    """summarization""": [],
    """zero-shot-classification""": [],
    """table-to-text""": [],
    """feature-extraction""": [],
    """other""": [],
    """multiple-choice""": [],
    """text-classification""": [],
    """text-to-image""": [],
    """text2text-generation""": [],
    """zero-shot-image-classification""": [],
    """tabular-classification""": [],
    """tabular-regression""": [],
    """image-to-image""": [],
    """tabular-to-text""": [],
    """unconditional-image-generation""": [],
    """text-retrieval""": [],
    """text-to-speech""": [],
    """object-detection""": [],
    """audio-to-audio""": [],
    """text-generation""": [],
    """conversational""": [],
    """table-question-answering""": [],
    """visual-question-answering""": [],
    """image-to-text""": [],
    """reinforcement-learning""": [],
    """voice-activity-detection""": [],
    """time-series-forecasting""": [],
    """document-question-answering""": [],
}


if __name__ == "__main__":
		from argparse import ArgumentParser

		UpperCAmelCase_						=ArgumentParser(usage="""Validate the yaml metadata block of a README.md file.""")
		ap.add_argument("""readme_filepath""")
		UpperCAmelCase_						=ap.parse_args()

		UpperCAmelCase_						=Path(args.readme_filepath)
		UpperCAmelCase_						=DatasetMetadata.from_readme(readme_filepath)
		print(dataset_metadata)
		dataset_metadata.to_readme(readme_filepath)



