import os
from argparse import ArgumentParser
from pathlib import Path
from shutil import copyfile
from typing import List

from datasets import config
from datasets.builder import DatasetBuilder
from datasets.commands import BaseDatasetsCLICommand
from datasets.download.download_config import DownloadConfig
from datasets.download.download_manager import DownloadMode
from datasets.load import dataset_module_factory, import_main_class
from datasets.utils.info_utils import VerificationMode



def       __SCREAMING_SNAKE_CASE    (   __UpperCamelCase       :			Union[str, Any]					,	**__UpperCamelCase       :			str      )				->   int:







  """simple docstring"""





  return RunBeamCommand(
      args.dataset					,	args.name					,	args.cache_dir					,	args.beam_pipeline_options					,	args.data_dir					,	args.all_configs					,	args.save_info or args.save_infos					,	args.ignore_verifications					,	args.force_redownload					,	**__UpperCamelCase					,	)







class __snake_case							(		lowerCamelCase_						):




   @staticmethod
   def       __a		(       _lowercase     :							ArgumentParser ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= parser.add_parser("""run_beam"""			,      help="""Run a Beam dataset processing pipeline""" )
     run_beam_parser.add_argument("""dataset"""			,      type=_lowercase			,      help="""Name of the dataset to download""" )
     run_beam_parser.add_argument("""--name"""			,      type=_lowercase			,      default=_lowercase			,      help="""Dataset config name""" )
     run_beam_parser.add_argument(
         """--cache_dir"""			,      type=_lowercase			,      default=_lowercase			,      help="""Cache directory where the datasets are stored"""			,      )
     run_beam_parser.add_argument(
         """--beam_pipeline_options"""			,      type=_lowercase			,      default=""""""			,      help="""Beam pipeline options, separated by commas. Example:: `--beam_pipeline_options=job_name=my-job,project=my-project`"""			,      )
     run_beam_parser.add_argument(
         """--data_dir"""			,      type=_lowercase			,      default=_lowercase			,      help="""Can be used to specify a manual directory to get the files from"""			,      )
     run_beam_parser.add_argument("""--all_configs"""			,      action="""store_true"""			,      help="""Test all dataset configurations""" )
     run_beam_parser.add_argument("""--save_info"""			,      action="""store_true"""			,      help="""Save the dataset infos file""" )
     run_beam_parser.add_argument(
         """--ignore_verifications"""			,      action="""store_true"""			,      help="""Run the test without checksums and splits checks""" )
     run_beam_parser.add_argument("""--force_redownload"""			,      action="""store_true"""			,      help="""Force dataset redownload""" )
     # aliases
     run_beam_parser.add_argument("""--save_infos"""			,      action="""store_true"""			,      help="""alias for save_info""" )
     run_beam_parser.set_defaults(func=_lowercase )




   def __init__(       self     :							List[str]			,      _lowercase     :							str			,      _lowercase     :							str			,      _lowercase     :							str			,      _lowercase     :							str			,      _lowercase     :							str			,      _lowercase     :							bool			,      _lowercase     :							bool			,      _lowercase     :							bool			,      _lowercase     :							bool			,      **_lowercase     :							Dict			,      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= dataset
     SCREAMING_SNAKE_CASE__     				= name
     SCREAMING_SNAKE_CASE__     				= cache_dir
     SCREAMING_SNAKE_CASE__     				= beam_pipeline_options
     SCREAMING_SNAKE_CASE__     				= data_dir
     SCREAMING_SNAKE_CASE__     				= all_configs
     SCREAMING_SNAKE_CASE__     				= save_infos
     SCREAMING_SNAKE_CASE__     				= ignore_verifications
     SCREAMING_SNAKE_CASE__     				= force_redownload
     SCREAMING_SNAKE_CASE__     				= config_kwargs




   def       __a		(       self     :							List[str] ):



     """simple docstring"""






     import apache_beam as beam

     if self._name is not None and self._all_configs:
       print("""Both parameters `name` and `all_configs` can't be used at once.""" )
       exit(1 )
     SCREAMING_SNAKE_CASE__				,	SCREAMING_SNAKE_CASE__     				= self._dataset, self._name
     SCREAMING_SNAKE_CASE__     				= dataset_module_factory(_lowercase )
     SCREAMING_SNAKE_CASE__     				= import_main_class(dataset_module.module_path )
     SCREAMING_SNAKE_CASE__     				= []
     if self._beam_pipeline_options:
       SCREAMING_SNAKE_CASE__     				= beam.options.pipeline_options.PipelineOptions(
           flags=[f"""--{opt.strip()}""" for opt in self._beam_pipeline_options.split(""",""" ) if opt] )
     else:
       SCREAMING_SNAKE_CASE__     				= None
     if self._all_configs and len(builder_cls.BUILDER_CONFIGS ) > 0:
       for builder_config in builder_cls.BUILDER_CONFIGS:
         builders.append(
             builder_cls(
                 config_name=builder_config.name			,      data_dir=self._data_dir			,      hash=dataset_module.hash			,      beam_options=_lowercase			,      cache_dir=self._cache_dir			,      base_path=dataset_module.builder_kwargs.get("""base_path""" )			,      ) )
     else:
       builders.append(
           builder_cls(
               config_name=_lowercase			,      data_dir=self._data_dir			,      beam_options=_lowercase			,      cache_dir=self._cache_dir			,      base_path=dataset_module.builder_kwargs.get("""base_path""" )			,      **self._config_kwargs			,      ) )

     for builder in builders:
       builder.download_and_prepare(
           download_mode=DownloadMode.REUSE_CACHE_IF_EXISTS
           if not self._force_redownload
           else DownloadMode.FORCE_REDOWNLOAD			,      download_config=DownloadConfig(cache_dir=config.DOWNLOADED_DATASETS_PATH )			,      verification_mode=VerificationMode.NO_CHECKS
           if self._ignore_verifications
           else VerificationMode.ALL_CHECKS			,      try_from_hf_gcs=_lowercase			,      )
       if self._save_infos:
         builder._save_infos()

     print("""Apache beam run successful.""" )

     # If save_infos=True, the dataset infos file is created next to the loaded module file.
     # Let's move it to the original directory of the dataset script, to allow the user to
     # upload them on S3 at the same time afterwards.
     if self._save_infos:
       SCREAMING_SNAKE_CASE__     				= os.path.join(builder_cls.get_imported_module_dir()			,      config.DATASETDICT_INFOS_FILENAME )

       SCREAMING_SNAKE_CASE__     				= Path(_lowercase ).name + """.py"""

       SCREAMING_SNAKE_CASE__     				= os.path.join(_lowercase			,      _lowercase )
       if os.path.isfile(_lowercase ):
         SCREAMING_SNAKE_CASE__     				= os.path.dirname(_lowercase )
       elif os.path.isfile(_lowercase ):
         SCREAMING_SNAKE_CASE__     				= path
       else:  # in case of a remote dataset
         print(f"""Dataset Infos file saved at {dataset_infos_path}""" )
         exit(1 )

       # Move datasetinfo back to the user
       SCREAMING_SNAKE_CASE__     				= os.path.join(_lowercase			,      config.DATASETDICT_INFOS_FILENAME )
       copyfile(_lowercase			,      _lowercase )
       print(f"""Dataset Infos file saved at {user_dataset_infos_path}""" )

import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()







@skip_mps
class __snake_case							(		lowerCamelCase_  ,						lowerCamelCase_  ,						lowerCamelCase_  ,						unittest.TestCase						):
   lowerCAmelCase_								=				StableDiffusionModelEditingPipeline
   lowerCAmelCase_								=				TEXT_TO_IMAGE_PARAMS
   lowerCAmelCase_								=				TEXT_TO_IMAGE_BATCH_PARAMS
   lowerCAmelCase_								=				TEXT_TO_IMAGE_IMAGE_PARAMS
   lowerCAmelCase_								=				TEXT_TO_IMAGE_IMAGE_PARAMS




   def       __a		(       self     :							List[str] ):



     """simple docstring"""






     torch.manual_seed(0 )
     SCREAMING_SNAKE_CASE__     				= UNetaDConditionModel(
         block_out_channels=(32, 64)			,      layers_per_block=2			,      sample_size=32			,      in_channels=4			,      out_channels=4			,      down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")			,      up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")			,      cross_attention_dim=32			,      )
     SCREAMING_SNAKE_CASE__     				= DDIMScheduler()
     torch.manual_seed(0 )
     SCREAMING_SNAKE_CASE__     				= AutoencoderKL(
         block_out_channels=[32, 64]			,      in_channels=3			,      out_channels=3			,      down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]			,      up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]			,      latent_channels=4			,      )
     torch.manual_seed(0 )
     SCREAMING_SNAKE_CASE__     				= CLIPTextConfig(
         bos_token_id=0			,      eos_token_id=2			,      hidden_size=32			,      intermediate_size=37			,      layer_norm_eps=1E-05			,      num_attention_heads=4			,      num_hidden_layers=5			,      pad_token_id=1			,      vocab_size=10_00			,      )
     SCREAMING_SNAKE_CASE__     				= CLIPTextModel(_lowercase )
     SCREAMING_SNAKE_CASE__     				= CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""" )

     SCREAMING_SNAKE_CASE__     				= {
         """unet""": unet,
         """scheduler""": scheduler,
         """vae""": vae,
         """text_encoder""": text_encoder,
         """tokenizer""": tokenizer,
         """safety_checker""": None,
         """feature_extractor""": None,
     }
     return components




   def       __a		(       self     :							Optional[int]			,      _lowercase     :							Optional[int]			,      _lowercase     :							Union[str, Any]=0 ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= torch.manual_seed(_lowercase )
     SCREAMING_SNAKE_CASE__     				= {
         """prompt""": """A field of roses""",
         """generator""": generator,
         # Setting height and width to None to prevent OOMs on CPU.
         """height""": None,
         """width""": None,
         """num_inference_steps""": 2,
         """guidance_scale""": 6.0,
         """output_type""": """numpy""",
     }
     return inputs




   def       __a		(       self     :							Dict ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= """cpu"""  # ensure determinism for the device-dependent torch.Generator
     SCREAMING_SNAKE_CASE__     				= self.get_dummy_components()
     SCREAMING_SNAKE_CASE__     				= StableDiffusionModelEditingPipeline(**_lowercase )
     SCREAMING_SNAKE_CASE__     				= sd_pipe.to(_lowercase )
     sd_pipe.set_progress_bar_config(disable=_lowercase )

     SCREAMING_SNAKE_CASE__     				= self.get_dummy_inputs(_lowercase )
     SCREAMING_SNAKE_CASE__     				= sd_pipe(**_lowercase ).images
     SCREAMING_SNAKE_CASE__     				= image[0, -3:, -3:, -1]
     assert image.shape == (1, 64, 64, 3)

     SCREAMING_SNAKE_CASE__     				= np.array([0.47_55, 0.51_32, 0.49_76, 0.39_04, 0.35_54, 0.47_65, 0.51_39, 0.51_58, 0.48_89] )

     assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2




   def       __a		(       self     :							List[str] ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= """cpu"""  # ensure determinism for the device-dependent torch.Generator
     SCREAMING_SNAKE_CASE__     				= self.get_dummy_components()
     SCREAMING_SNAKE_CASE__     				= StableDiffusionModelEditingPipeline(**_lowercase )
     SCREAMING_SNAKE_CASE__     				= sd_pipe.to(_lowercase )
     sd_pipe.set_progress_bar_config(disable=_lowercase )

     SCREAMING_SNAKE_CASE__     				= self.get_dummy_inputs(_lowercase )
     SCREAMING_SNAKE_CASE__     				= """french fries"""
     SCREAMING_SNAKE_CASE__     				= sd_pipe(**_lowercase			,      negative_prompt=_lowercase )
     SCREAMING_SNAKE_CASE__     				= output.images
     SCREAMING_SNAKE_CASE__     				= image[0, -3:, -3:, -1]

     assert image.shape == (1, 64, 64, 3)

     SCREAMING_SNAKE_CASE__     				= np.array([0.49_92, 0.51_01, 0.50_04, 0.39_49, 0.36_04, 0.47_35, 0.52_16, 0.52_04, 0.49_13] )

     assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2




   def       __a		(       self     :							Union[str, Any] ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= """cpu"""  # ensure determinism for the device-dependent torch.Generator
     SCREAMING_SNAKE_CASE__     				= self.get_dummy_components()
     SCREAMING_SNAKE_CASE__     				= EulerAncestralDiscreteScheduler(
         beta_start=0.0_00_85			,      beta_end=0.0_12			,      beta_schedule="""scaled_linear""" )
     SCREAMING_SNAKE_CASE__     				= StableDiffusionModelEditingPipeline(**_lowercase )
     SCREAMING_SNAKE_CASE__     				= sd_pipe.to(_lowercase )
     sd_pipe.set_progress_bar_config(disable=_lowercase )

     SCREAMING_SNAKE_CASE__     				= self.get_dummy_inputs(_lowercase )
     SCREAMING_SNAKE_CASE__     				= sd_pipe(**_lowercase ).images
     SCREAMING_SNAKE_CASE__     				= image[0, -3:, -3:, -1]

     assert image.shape == (1, 64, 64, 3)

     SCREAMING_SNAKE_CASE__     				= np.array([0.47_47, 0.53_72, 0.47_79, 0.49_82, 0.55_43, 0.48_16, 0.52_38, 0.49_04, 0.50_27] )

     assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2




   def       __a		(       self     :							List[Any] ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= """cpu"""  # ensure determinism for the device-dependent torch.Generator
     SCREAMING_SNAKE_CASE__     				= self.get_dummy_components()
     SCREAMING_SNAKE_CASE__     				= PNDMScheduler()
     SCREAMING_SNAKE_CASE__     				= StableDiffusionModelEditingPipeline(**_lowercase )
     SCREAMING_SNAKE_CASE__     				= sd_pipe.to(_lowercase )
     sd_pipe.set_progress_bar_config(disable=_lowercase )

     SCREAMING_SNAKE_CASE__     				= self.get_dummy_inputs(_lowercase )
     # the pipeline does not expect pndm so test if it raises error.
     with self.assertRaises(_lowercase ):
       SCREAMING_SNAKE_CASE__     				= sd_pipe(**_lowercase ).images




   def       __a		(       self     :							List[str] ):



     """simple docstring"""






     super().test_inference_batch_single_identical(expected_max_diff=5E-3 )




   def       __a		(       self     :							Any ):



     """simple docstring"""






     super().test_attention_slicing_forward_pass(expected_max_diff=5E-3 )







@slow
@require_torch_gpu
class __snake_case							(		unittest.TestCase						):




   def       __a		(       self     :							str ):



     """simple docstring"""






     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()




   def       __a		(       self     :							List[Any]			,      _lowercase     :							Dict=0 ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= torch.manual_seed(_lowercase )
     SCREAMING_SNAKE_CASE__     				= {
         """prompt""": """A field of roses""",
         """generator""": generator,
         """num_inference_steps""": 3,
         """guidance_scale""": 7.5,
         """output_type""": """numpy""",
     }
     return inputs




   def       __a		(       self     :							Tuple ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= """CompVis/stable-diffusion-v1-4"""
     SCREAMING_SNAKE_CASE__     				= StableDiffusionModelEditingPipeline.from_pretrained(_lowercase			,      safety_checker=_lowercase )
     pipe.to(_lowercase )
     pipe.set_progress_bar_config(disable=_lowercase )
     pipe.enable_attention_slicing()

     SCREAMING_SNAKE_CASE__     				= self.get_inputs()
     SCREAMING_SNAKE_CASE__     				= pipe(**_lowercase ).images
     SCREAMING_SNAKE_CASE__     				= image[0, -3:, -3:, -1].flatten()

     assert image.shape == (1, 5_12, 5_12, 3)

     SCREAMING_SNAKE_CASE__     				= np.array(
         [0.6_74_94_96, 0.6_38_64_53, 0.51_44_32_67, 0.66_09_49_05, 0.61_92_12_15, 0.5_49_13_32, 0.5_74_44_17, 0.58_07_51_06, 0.5_17_46_58] )

     assert np.abs(expected_slice - image_slice ).max() < 1E-2

     # make sure image changes after editing
     pipe.edit_model("""A pack of roses"""			,      """A pack of blue roses""" )

     SCREAMING_SNAKE_CASE__     				= pipe(**_lowercase ).images
     SCREAMING_SNAKE_CASE__     				= image[0, -3:, -3:, -1].flatten()

     assert image.shape == (1, 5_12, 5_12, 3)

     assert np.abs(expected_slice - image_slice ).max() > 1E-1




   def       __a		(       self     :							Any ):



     """simple docstring"""






     torch.cuda.empty_cache()
     torch.cuda.reset_max_memory_allocated()
     torch.cuda.reset_peak_memory_stats()

     SCREAMING_SNAKE_CASE__     				= """CompVis/stable-diffusion-v1-4"""
     SCREAMING_SNAKE_CASE__     				= DDIMScheduler.from_pretrained(_lowercase			,      subfolder="""scheduler""" )
     SCREAMING_SNAKE_CASE__     				= StableDiffusionModelEditingPipeline.from_pretrained(
         _lowercase			,      scheduler=_lowercase			,      safety_checker=_lowercase )
     SCREAMING_SNAKE_CASE__     				= pipe.to(_lowercase )
     pipe.set_progress_bar_config(disable=_lowercase )
     pipe.enable_attention_slicing(1 )
     pipe.enable_sequential_cpu_offload()

     SCREAMING_SNAKE_CASE__     				= self.get_inputs()
     SCREAMING_SNAKE_CASE__     				= pipe(**_lowercase )

     SCREAMING_SNAKE_CASE__     				= torch.cuda.max_memory_allocated()
     # make sure that less than 4.4 GB is allocated
     assert mem_bytes < 4.4 * 10**9
