


"""simple docstring"""



import multiprocessing
import os
from typing import BinaryIO, Optional, Union

import fsspec

from .. import Dataset, Features, NamedSplit, config
from ..formatting import query_table
from ..packaged_modules.json.json import Json
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader







class _lowerCAmelCase    (				snake_case_     ):







   def __init__(     self						,     UpperCamelCase__						,     UpperCamelCase__ = None						,     UpperCamelCase__ = None						,     UpperCamelCase__ = None						,     UpperCamelCase__ = False						,     UpperCamelCase__ = False						,     UpperCamelCase__ = None						,     UpperCamelCase__ = None						,     **UpperCamelCase__						,     )     ->		int:

       '''simple docstring'''




       super().__init__(
           UpperCamelCase__						,     split=UpperCamelCase__						,     features=UpperCamelCase__						,     cache_dir=UpperCamelCase__						,     keep_in_memory=UpperCamelCase__						,     streaming=UpperCamelCase__						,     num_proc=UpperCamelCase__						,     **UpperCamelCase__						,     )
       snake_case   :				Tuple				 =   field
       snake_case   :				List[str]				 =   path_or_paths if isinstance(UpperCamelCase__						,     UpperCamelCase__						) else {self.split: path_or_paths}
       snake_case   :				Optional[int]				 =   Json(
           cache_dir=UpperCamelCase__						,     data_files=UpperCamelCase__						,     features=UpperCamelCase__						,     field=UpperCamelCase__						,     **UpperCamelCase__						,     )







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

       '''simple docstring'''




       if self.streaming:
           snake_case   :				Union[str, Any]				 =   self.builder.as_streaming_dataset(split=self.split						)
       # Build regular (map-style) dataset
       else:
           snake_case   :				int				 =   None
           snake_case   :				Any				 =   None
           snake_case   :				Union[str, Any]				 =   None
           snake_case   :				Optional[int]				 =   None

           self.builder.download_and_prepare(
               download_config=UpperCamelCase__						,     download_mode=UpperCamelCase__						,     verification_mode=UpperCamelCase__						,     base_path=UpperCamelCase__						,     num_proc=self.num_proc						,     )
           snake_case   :				List[str]				 =   self.builder.as_dataset(
               split=self.split						,     verification_mode=UpperCamelCase__						,     in_memory=self.keep_in_memory						)
       return dataset







class _lowerCAmelCase    :







   def __init__(     self						,     UpperCamelCase__						,     UpperCamelCase__						,     UpperCamelCase__ = None						,     UpperCamelCase__ = None						,     **UpperCamelCase__						,     )     ->		int:

       '''simple docstring'''




       if num_proc is not None and num_proc <= 0:
           raise ValueError(F'num_proc {num_proc} must be an integer > 0.'						)

       snake_case   :				Union[str, Any]				 =   dataset
       snake_case   :				Dict				 =   path_or_buf
       snake_case   :				Union[str, Any]				 =   batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
       snake_case   :				List[str]				 =   num_proc
       snake_case   :				Any				 =   "utf-8"
       snake_case   :				Dict				 =   to_json_kwargs







   def 			lowerCamelCase		(     self						)     ->		int:

       '''simple docstring'''




       snake_case   :				List[Any]				 =   self.to_json_kwargs.pop("path_or_buf"						,     UpperCamelCase__						)
       snake_case   :				List[Any]				 =   self.to_json_kwargs.pop("orient"						,     "records"						)
       snake_case   :				Optional[Any]				 =   self.to_json_kwargs.pop("lines"						,     True if orient == "records" else False						)
       snake_case   :				List[str]				 =   self.to_json_kwargs.pop("index"						,     False if orient in ["split", "table"] else True						)
       snake_case   :				List[Any]				 =   self.to_json_kwargs.pop("compression"						,     UpperCamelCase__						)

       if compression not in [None, "infer", "gzip", "bz2", "xz"]:
           raise NotImplementedError(F'`datasets` currently does not support {compression} compression'						)

       if isinstance(self.path_or_buf						,     (str, bytes, os.PathLike)						):
           with fsspec.open(self.path_or_buf						,     "wb"						,     compression=UpperCamelCase__						) as buffer:
               snake_case   :				int				 =   self._write(file_obj=UpperCamelCase__						,     orient=UpperCamelCase__						,     lines=UpperCamelCase__						,     index=UpperCamelCase__						,     **self.to_json_kwargs						)
       else:
           if compression:
               raise NotImplementedError(
                   F'The compression parameter is not supported when writing to a buffer, but compression={compression}'
                   " was passed. Please provide a local path instead."						)
           snake_case   :				int				 =   self._write(
               file_obj=self.path_or_buf						,     orient=UpperCamelCase__						,     lines=UpperCamelCase__						,     index=UpperCamelCase__						,     **self.to_json_kwargs						)
       return written







   def 			lowerCamelCase		(     self						,     UpperCamelCase__						)     ->		Any:

       '''simple docstring'''




       snake_case   :				List[Any]				 =   args

       snake_case   :				Tuple				 =   query_table(
           table=self.dataset.data						,     key=slice(UpperCamelCase__						,     offset + self.batch_size						)						,     indices=self.dataset._indices						,     )
       snake_case   :				Optional[Any]				 =   batch.to_pandas().to_json(
           path_or_buf=UpperCamelCase__						,     orient=UpperCamelCase__						,     lines=UpperCamelCase__						,     index=UpperCamelCase__						,     **UpperCamelCase__						)
       if not json_str.endswith("\n"						):
           json_str += "\n"
       return json_str.encode(self.encoding						)







   def 			lowerCamelCase		(     self						,     UpperCamelCase__						,     UpperCamelCase__						,     UpperCamelCase__						,     UpperCamelCase__						,     **UpperCamelCase__						,     )     ->		int:

       '''simple docstring'''




       snake_case   :				Optional[Any]				 =   0

       if self.num_proc is None or self.num_proc == 1:
           for offset in logging.tqdm(
               range(0						,     len(self.dataset						)						,     self.batch_size						)						,     unit="ba"						,     disable=not logging.is_progress_bar_enabled()						,     desc="Creating json from Arrow format"						,     ):
               snake_case   :				Dict				 =   self._batch_json((offset, orient, lines, index, to_json_kwargs)						)
               written += file_obj.write(UpperCamelCase__						)
       else:
           snake_case   :				Dict				 =   len(self.dataset						), self.batch_size
           with multiprocessing.Pool(self.num_proc						) as pool:
               for json_str in logging.tqdm(
                   pool.imap(
                       self._batch_json						,     [(offset, orient, lines, index, to_json_kwargs) for offset in range(0						,     UpperCamelCase__						,     UpperCamelCase__						)]						,     )						,     total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size						,     unit="ba"						,     disable=not logging.is_progress_bar_enabled()						,     desc="Creating json from Arrow format"						,     ):
                   written += file_obj.write(UpperCamelCase__						)

       return written









"""simple docstring"""



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

import numpy as np
import PIL
from PIL import Image

from ...utils import (
    BaseOutput,
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_k_diffusion_available,
    is_k_diffusion_version,
    is_onnx_available,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)







@dataclass
class _lowerCAmelCase    (				snake_case_     ):
   __UpperCAmelCase			:	Union[List[PIL.Image.Image], np.ndarray]
   __UpperCAmelCase			:	Optional[List[bool]]


try:
  if not (is_transformers_available() and is_torch_available()):
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
  from .pipeline_cycle_diffusion import CycleDiffusionPipeline
  from .pipeline_stable_diffusion import StableDiffusionPipeline
  from .pipeline_stable_diffusion_attend_and_excite import StableDiffusionAttendAndExcitePipeline
  from .pipeline_stable_diffusion_imgaimg import StableDiffusionImgaImgPipeline
  from .pipeline_stable_diffusion_inpaint import StableDiffusionInpaintPipeline
  from .pipeline_stable_diffusion_inpaint_legacy import StableDiffusionInpaintPipelineLegacy
  from .pipeline_stable_diffusion_instruct_pixapix import StableDiffusionInstructPixaPixPipeline
  from .pipeline_stable_diffusion_latent_upscale import StableDiffusionLatentUpscalePipeline
  from .pipeline_stable_diffusion_ldmad import StableDiffusionLDMaDPipeline
  from .pipeline_stable_diffusion_model_editing import StableDiffusionModelEditingPipeline
  from .pipeline_stable_diffusion_panorama import StableDiffusionPanoramaPipeline
  from .pipeline_stable_diffusion_paradigms import StableDiffusionParadigmsPipeline
  from .pipeline_stable_diffusion_sag import StableDiffusionSAGPipeline
  from .pipeline_stable_diffusion_upscale import StableDiffusionUpscalePipeline
  from .pipeline_stable_unclip import StableUnCLIPPipeline
  from .pipeline_stable_unclip_imgaimg import StableUnCLIPImgaImgPipeline
  from .safety_checker import StableDiffusionSafetyChecker
  from .stable_unclip_image_normalizer import StableUnCLIPImageNormalizer

try:
  if not (is_transformers_available() and is_torch_available() and is_transformers_version(""">=""", """4.25.0""")):
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  from ...utils.dummy_torch_and_transformers_objects import StableDiffusionImageVariationPipeline
else:
  from .pipeline_stable_diffusion_image_variation import StableDiffusionImageVariationPipeline


try:
  if not (is_transformers_available() and is_torch_available() and is_transformers_version(""">=""", """4.26.0""")):
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  from ...utils.dummy_torch_and_transformers_objects import (
      StableDiffusionDepthaImgPipeline,
      StableDiffusionDiffEditPipeline,
      StableDiffusionPixaPixZeroPipeline,
  )
else:
  from .pipeline_stable_diffusion_depthaimg import StableDiffusionDepthaImgPipeline
  from .pipeline_stable_diffusion_diffedit import StableDiffusionDiffEditPipeline
  from .pipeline_stable_diffusion_pixapix_zero import StableDiffusionPixaPixZeroPipeline


try:
  if not (
      is_torch_available()
      and is_transformers_available()
      and is_k_diffusion_available()
      and is_k_diffusion_version(""">=""", """0.0.12""")
  ):
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  from ...utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
else:
  from .pipeline_stable_diffusion_k_diffusion import StableDiffusionKDiffusionPipeline

try:
  if not (is_transformers_available() and is_onnx_available()):
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  from ...utils.dummy_onnx_objects import *  # noqa F403
else:
  from .pipeline_onnx_stable_diffusion import OnnxStableDiffusionPipeline, StableDiffusionOnnxPipeline
  from .pipeline_onnx_stable_diffusion_imgaimg import OnnxStableDiffusionImgaImgPipeline
  from .pipeline_onnx_stable_diffusion_inpaint import OnnxStableDiffusionInpaintPipeline
  from .pipeline_onnx_stable_diffusion_inpaint_legacy import OnnxStableDiffusionInpaintPipelineLegacy
  from .pipeline_onnx_stable_diffusion_upscale import OnnxStableDiffusionUpscalePipeline

if is_transformers_available() and is_flax_available():
  import flax







  @flax.struct.dataclass
  class _lowerCAmelCase    (				snake_case_     ):
     __UpperCAmelCase			:	np.ndarray
     __UpperCAmelCase			:	List[bool]



  from ...schedulers.scheduling_pndm_flax import PNDMSchedulerState
  from .pipeline_flax_stable_diffusion import FlaxStableDiffusionPipeline
  from .pipeline_flax_stable_diffusion_imgaimg import FlaxStableDiffusionImgaImgPipeline
  from .pipeline_flax_stable_diffusion_inpaint import FlaxStableDiffusionInpaintPipeline
  from .safety_checker_flax import FlaxStableDiffusionSafetyChecker





