


"""simple docstring"""







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,
)


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

__lowercase    					=      """
    Examples:
        ```py
        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline
        >>> import torch

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


def      lowercase			(				A_  ,   A_  ,   A_=8       )->      List[Any]:
  '''simple docstring'''







  a						:		int      							=			height // scale_factor**2
  if height % scale_factor**2 != 0:
    new_height += 1
  a						:		List[str]      							=			width // scale_factor**2
  if width % scale_factor**2 != 0:
    new_width += 1
  return new_height * scale_factor, new_width * scale_factor







class 						_A  (  _a  ):







    """simple docstring"""




    def __init__(					self				:    Tuple				, __UpperCAmelCase				:    UNetaDConditionModel				, __UpperCAmelCase				:    DDPMScheduler				, __UpperCAmelCase				:    VQModel				, ):
      super().__init__()

      self.register_modules(
          unet=__UpperCAmelCase				, scheduler=__UpperCAmelCase				, movq=__UpperCAmelCase				, )
      a						:		int      							=			2 ** (len(self.movq.config.block_out_channels) - 1)






    def 	__snake_case							(					self				:    Optional[Any]				, __UpperCAmelCase				:    Union[str, Any]				, __UpperCAmelCase				:    Optional[int]				, __UpperCAmelCase				:    Union[str, Any]				, __UpperCAmelCase				:    Any				, __UpperCAmelCase				:    List[str]				, __UpperCAmelCase				:    List[str]):
      if latents is None:
        a						:		Any      							=			randn_tensor(__UpperCAmelCase				, generator=__UpperCAmelCase				, device=__UpperCAmelCase				, dtype=__UpperCAmelCase)
      else:
        if latents.shape != shape:
          raise ValueError(f'''Unexpected latents shape, got {latents.shape}, expected {shape}''')
        a						:		Dict      							=			latents.to(__UpperCAmelCase)

      a						:		Dict      							=			latents * scheduler.init_noise_sigma
      return latents






    def 	__snake_case							(					self				:    Any				, __UpperCAmelCase				:    List[Any]=0):
      if is_accelerate_available():
        from accelerate import cpu_offload
      else:
        raise ImportError("Please install accelerate via `pip install accelerate`")

      a						:		Dict      							=			torch.device(f'''cuda:{gpu_id}''')

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






    def 	__snake_case							(					self				:    List[Any]				, __UpperCAmelCase				:    Union[str, Any]=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.")

      a						:		List[Any]      							=			torch.device(f'''cuda:{gpu_id}''')

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

      a						:		Any      							=			None
      for cpu_offloaded_model in [self.unet, self.movq]:
        a							,		a						:		Tuple      							=			cpu_offload_with_hook(__UpperCAmelCase				, __UpperCAmelCase				, prev_module_hook=__UpperCAmelCase)

      # We'll offload the last model manually.
      a						:		Dict      							=			hook






    @property
    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
    def 	__snake_case							(					self				:    Optional[int]):
      if not hasattr(self.unet				, "_hf_hook"):
        return self.device
      for module in self.unet.modules():
        if (
            hasattr(__UpperCAmelCase				, "_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(__UpperCAmelCase)
    def __call__(					self				:    Optional[Any]				, __UpperCAmelCase				:    Union[torch.FloatTensor, List[torch.FloatTensor]]				, __UpperCAmelCase				:    Union[torch.FloatTensor, List[torch.FloatTensor]]				, __UpperCAmelCase				:    int = 512				, __UpperCAmelCase				:    int = 512				, __UpperCAmelCase				:    int = 100				, __UpperCAmelCase				:    float = 4.0				, __UpperCAmelCase				:    int = 1				, __UpperCAmelCase				:    Optional[Union[torch.Generator, List[torch.Generator]]] = None				, __UpperCAmelCase				:    Optional[torch.FloatTensor] = None				, __UpperCAmelCase				:    Optional[str] = "pil"				, __UpperCAmelCase				:    bool = True				, ):
      a						:		Optional[int]      							=			self._execution_device

      a						:		int      							=			guidance_scale > 1.0

      if isinstance(__UpperCAmelCase				, __UpperCAmelCase):
        a						:		Union[str, Any]      							=			torch.cat(__UpperCAmelCase				, dim=0)
      a						:		int      							=			image_embeds.shape[0] * num_images_per_prompt
      if isinstance(__UpperCAmelCase				, __UpperCAmelCase):
        a						:		List[str]      							=			torch.cat(__UpperCAmelCase				, dim=0)

      if do_classifier_free_guidance:
        a						:		Union[str, Any]      							=			image_embeds.repeat_interleave(__UpperCAmelCase				, dim=0)
        a						:		Tuple      							=			negative_image_embeds.repeat_interleave(__UpperCAmelCase				, dim=0)

      a						:		Union[str, Any]      							=			torch.cat([negative_image_embeds, image_embeds]				, dim=0).to(dtype=self.unet.dtype				, device=__UpperCAmelCase)

      self.scheduler.set_timesteps(__UpperCAmelCase				, device=__UpperCAmelCase)
      a						:		Union[str, Any]      							=			self.scheduler.timesteps

      a						:		Any      							=			self.unet.config.in_channels

      a							,		a						:		str      							=			downscale_height_and_width(__UpperCAmelCase				, __UpperCAmelCase				, self.movq_scale_factor)

      # create initial latent
      a						:		Any      							=			self.prepare_latents(
          (batch_size, num_channels_latents, height, width)				, image_embeds.dtype				, __UpperCAmelCase				, __UpperCAmelCase				, __UpperCAmelCase				, self.scheduler				, )

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

        a						:		Dict      							=			{"image_embeds": image_embeds}
        a						:		Tuple      							=			self.unet(
            sample=__UpperCAmelCase				, timestep=__UpperCAmelCase				, encoder_hidden_states=__UpperCAmelCase				, added_cond_kwargs=__UpperCAmelCase				, return_dict=__UpperCAmelCase				, )[0]

        if do_classifier_free_guidance:
          a							,		a						:		int      							=			noise_pred.split(latents.shape[1]				, dim=1)
          a							,		a						:		Any      							=			noise_pred.chunk(2)
          a							,		a						:		int      							=			variance_pred.chunk(2)
          a						:		Any      							=			noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
          a						:		str      							=			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"]
        ):
          a							,		a						:		Union[str, Any]      							=			noise_pred.split(latents.shape[1]				, dim=1)

        # compute the previous noisy sample x_t -> x_t-1
        a						:		Any      							=			self.scheduler.step(
            __UpperCAmelCase				, __UpperCAmelCase				, __UpperCAmelCase				, generator=__UpperCAmelCase				, )[0]
      # post-processing
      a						:		Tuple      							=			self.movq.decode(__UpperCAmelCase				, force_not_quantize=__UpperCAmelCase)["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"]:
        a						:		Optional[Any]      							=			image * 0.5 + 0.5
        a						:		Union[str, Any]      							=			image.clamp(0				, 1)
        a						:		Union[str, Any]      							=			image.cpu().permute(0				, 2				, 3				, 1).float().numpy()

      if output_type == "pil":
        a						:		Optional[int]      							=			self.numpy_to_pil(__UpperCAmelCase)

      if not return_dict:
        return (image,)

      return ImagePipelineOutput(images=__UpperCAmelCase)










"""simple docstring"""







import itertools
from dataclasses import dataclass
from typing import List, Optional

import pyarrow as pa
import pyarrow.parquet as pq

import datasets
from datasets.table import table_cast


__lowercase    					=      datasets.utils.logging.get_logger(__name__)






@dataclass
class 						_A  (  datasets.BuilderConfig  ):







    """simple docstring"""




    UpperCAmelCase       :							int	=       1_0_0_0_0
    UpperCAmelCase       :							Optional[List[str]]	=       None
    UpperCAmelCase       :							Optional[datasets.Features]	=       None






class 						_A  (  datasets.ArrowBasedBuilder  ):







    """simple docstring"""




    UpperCAmelCase       :							str	=       ParquetConfig






    def 	__snake_case							(					self				:    Tuple):
      return datasets.DatasetInfo(features=self.config.features)






    def 	__snake_case							(					self				:    List[Any]				, __UpperCAmelCase				:    str):
      if not self.config.data_files:
        raise ValueError(f'''At least one data file must be specified, but got data_files={self.config.data_files}''')
      a						:		str      							=			dl_manager.download_and_extract(self.config.data_files)
      if isinstance(__UpperCAmelCase				, (str, list, tuple)):
        a						:		Dict      							=			data_files
        if isinstance(__UpperCAmelCase				, __UpperCAmelCase):
          a						:		str      							=			[files]
        # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
        a						:		List[Any]      							=			[dl_manager.iter_files(__UpperCAmelCase) for file in files]
        return [datasets.SplitGenerator(name=datasets.Split.TRAIN				, gen_kwargs={"files": files})]
      a						:		Dict      							=			[]
      for split_name, files in data_files.items():
        if isinstance(__UpperCAmelCase				, __UpperCAmelCase):
          a						:		Optional[int]      							=			[files]
        # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
        a						:		Tuple      							=			[dl_manager.iter_files(__UpperCAmelCase) for file in files]
        # Infer features is they are stoed in the arrow schema
        if self.info.features is None:
          for file in itertools.chain.from_iterable(__UpperCAmelCase):
            with open(__UpperCAmelCase				, "rb") as f:
              a						:		Tuple      							=			datasets.Features.from_arrow_schema(pq.read_schema(__UpperCAmelCase))
            break
        splits.append(datasets.SplitGenerator(name=__UpperCAmelCase				, gen_kwargs={"files": files}))
      return splits






    def 	__snake_case							(					self				:    List[str]				, __UpperCAmelCase				:    pa.Table):
      if self.info.features is not None:
        # more expensive cast to support nested features with keys in a different order
        # allows str <-> int/float or str to Audio for example
        a						:		Optional[int]      							=			table_cast(__UpperCAmelCase				, self.info.features.arrow_schema)
      return pa_table






    def 	__snake_case							(					self				:    Union[str, Any]				, __UpperCAmelCase				:    int):
      a						:		Tuple      							=			self.info.features.arrow_schema if self.info.features is not None else None
      if self.info.features is not None and self.config.columns is not None:
        if sorted(field.name for field in schema) != sorted(self.config.columns):
          raise ValueError(
              f'''Tried to load parquet data with columns \'{self.config.columns}\' with mismatching features \'{self.info.features}\'''')
      for file_idx, file in enumerate(itertools.chain.from_iterable(__UpperCAmelCase)):
        with open(__UpperCAmelCase				, "rb") as f:
          a						:		Tuple      							=			pq.ParquetFile(__UpperCAmelCase)
          try:
            for batch_idx, record_batch in enumerate(
                parquet_file.iter_batches(batch_size=self.config.batch_size				, columns=self.config.columns)):
              a						:		Optional[Any]      							=			pa.Table.from_batches([record_batch])
              # Uncomment for debugging (will print the Arrow table size and elements)
              # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
              # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
              yield f'''{file_idx}_{batch_idx}''', self._cast_table(__UpperCAmelCase)
          except ValueError as e:
            logger.error(f'''Failed to read file \'{file}\' with error {type(__UpperCAmelCase)}: {e}''')
            raise






