


import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModel,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableUnCLIPImgaImgPipeline, UNetaDConditionModel
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    floats_tensor,
    load_image,
    load_numpy,
    require_torch_gpu,
    skip_mps,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()






class    UpperCAmelCase__	( A_       ,  A_       ,  A_       ,  unittest.TestCase		):



  """simple docstring"""






  UpperCAmelCase__						:		Tuple              =  StableUnCLIPImgaImgPipeline
  UpperCAmelCase__						:		List[Any]              =  TEXT_GUIDED_IMAGE_VARIATION_PARAMS
  UpperCAmelCase__						:		Optional[int]              =  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
  UpperCAmelCase__						:		Tuple              =  frozenset(
      []		)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
  UpperCAmelCase__						:		Union[str, Any]              =  frozenset([]		)






  def _a (      self   )  ->	List[str]:
       __UpperCamelCase   							=32
       __UpperCamelCase   							=embedder_hidden_size

       # image encoding components

       __UpperCamelCase   							=CLIPImageProcessor(crop_size=32 ,       size=32   )

       torch.manual_seed(0   )
       __UpperCamelCase   							=CLIPVisionModelWithProjection(
           CLIPVisionConfig(
               hidden_size=lowercase_ ,       projection_dim=lowercase_ ,       num_hidden_layers=5 ,       num_attention_heads=4 ,       image_size=32 ,       intermediate_size=37 ,       patch_size=1 ,       )   )

       # regular denoising components

       torch.manual_seed(0   )
       __UpperCamelCase   							=StableUnCLIPImageNormalizer(embedding_dim=lowercase_   )
       __UpperCamelCase   							=DDPMScheduler(beta_schedule='squaredcos_cap_v2'   )

       torch.manual_seed(0   )
       __UpperCamelCase   							=CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'   )

       torch.manual_seed(0   )
       __UpperCamelCase   							=CLIPTextModel(
           CLIPTextConfig(
               bos_token_id=0 ,       eos_token_id=2 ,       hidden_size=lowercase_ ,       projection_dim=32 ,       intermediate_size=37 ,       layer_norm_eps=1E-05 ,       num_attention_heads=4 ,       num_hidden_layers=5 ,       pad_token_id=1 ,       vocab_size=1000 ,       )   )

       torch.manual_seed(0   )
       __UpperCamelCase   							=UNetaDConditionModel(
           sample_size=32 ,       in_channels=4 ,       out_channels=4 ,       down_block_types=('CrossAttnDownBlock2D', 'DownBlock2D') ,       up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D') ,       block_out_channels=(32, 64) ,       attention_head_dim=(2, 4) ,       class_embed_type='projection' ,       projection_class_embeddings_input_dim=embedder_projection_dim * 2 ,       cross_attention_dim=lowercase_ ,       layers_per_block=1 ,       upcast_attention=lowercase_ ,       use_linear_projection=lowercase_ ,       )

       torch.manual_seed(0   )
       __UpperCamelCase   							=DDIMScheduler(
           beta_schedule='scaled_linear' ,       beta_start=0.0_0085 ,       beta_end=0.012 ,       prediction_type='v_prediction' ,       set_alpha_to_one=lowercase_ ,       steps_offset=1 ,       )

       torch.manual_seed(0   )
       __UpperCamelCase   							=AutoencoderKL()

       __UpperCamelCase   							={
           # image encoding components
           'feature_extractor': feature_extractor,
           'image_encoder': image_encoder.eval(),
           # image noising components
           'image_normalizer': image_normalizer.eval(),
           'image_noising_scheduler': image_noising_scheduler,
           # regular denoising components
           'tokenizer': tokenizer,
           'text_encoder': text_encoder.eval(),
           'unet': unet.eval(),
           'scheduler': scheduler,
           'vae': vae.eval(),
       }

       return components






  def _a (      self ,       A_ ,       A_=0 ,       A_=True   )  ->	Tuple:







       if str(lowercase_   ).startswith('mps'   ):
            __UpperCamelCase   							=torch.manual_seed(lowercase_   )
       else:
            __UpperCamelCase   							=torch.Generator(device=lowercase_   ).manual_seed(lowercase_   )

       __UpperCamelCase   							=floats_tensor((1, 3, 32, 32) ,       rng=random.Random(lowercase_   )   ).to(lowercase_   )

       if pil_image:
            __UpperCamelCase   							=input_image * 0.5 + 0.5
            __UpperCamelCase   							=input_image.clamp(0 ,       1   )
            __UpperCamelCase   							=input_image.cpu().permute(0 ,       2 ,       3 ,       1   ).float().numpy()
            __UpperCamelCase   							=DiffusionPipeline.numpy_to_pil(lowercase_   )[0]

       return {
           "prompt": "An anime racoon running a marathon",
           "image": input_image,
           "generator": generator,
           "num_inference_steps": 2,
           "output_type": "np",
       }






  @skip_mps
  def _a (      self   )  ->	Dict:
       __UpperCamelCase   							='cpu'  # ensure determinism for the device-dependent torch.Generator
       __UpperCamelCase   							=self.get_dummy_components()
       __UpperCamelCase   							=StableUnCLIPImgaImgPipeline(**lowercase_   )
       __UpperCamelCase   							=sd_pipe.to(lowercase_   )
       sd_pipe.set_progress_bar_config(disable=lowercase_   )

       __UpperCamelCase   							=self.get_dummy_inputs(lowercase_   )
       inputs.update({'image_embeds': None}   )
       __UpperCamelCase   							=sd_pipe(**lowercase_   ).images
       __UpperCamelCase   							=image[0, -3:, -3:, -1]

       assert image.shape == (1, 32, 32, 3)
       __UpperCamelCase   							=np.array([0.3872, 0.7224, 0.5601, 0.4741, 0.6872, 0.5814, 0.4636, 0.3867, 0.5078]   )

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






  def _a (      self   )  ->	Union[str, Any]:
       __UpperCamelCase   							=torch_device in ['cpu', 'mps']

       self._test_attention_slicing_forward_pass(test_max_difference=lowercase_   )






  def _a (      self   )  ->	Optional[Any]:
       __UpperCamelCase   							=torch_device in ['cpu', 'mps']

       self._test_inference_batch_single_identical(test_max_difference=lowercase_   )






  @unittest.skipIf(
      torch_device != 'cuda' or not is_xformers_available() ,       reason='XFormers attention is only available with CUDA and `xformers` installed' ,       )
  def _a (      self   )  ->	Tuple:
       self._test_xformers_attention_forwardGenerator_pass(test_max_difference=lowercase_   )






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



  """simple docstring"""






  def _a (      self   )  ->	Any:
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()






  def _a (      self   )  ->	str:
       __UpperCamelCase   							=load_image(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'   )

       __UpperCamelCase   							=load_numpy(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_img2img_anime_turtle_fp16.npy'   )

       __UpperCamelCase   							=StableUnCLIPImgaImgPipeline.from_pretrained(
           'fusing/stable-unclip-2-1-l-img2img' ,       torch_dtype=torch.floataa   )
       pipe.to(lowercase_   )
       pipe.set_progress_bar_config(disable=lowercase_   )
       # stable unclip will oom when integration tests are run on a V100,
       # so turn on memory savings
       pipe.enable_attention_slicing()
       pipe.enable_sequential_cpu_offload()

       __UpperCamelCase   							=torch.Generator(device='cpu'   ).manual_seed(0   )
       __UpperCamelCase   							=pipe(lowercase_ ,       'anime turle' ,       generator=lowercase_ ,       output_type='np'   )

       __UpperCamelCase   							=output.images[0]

       assert image.shape == (768, 768, 3)

       assert_mean_pixel_difference(lowercase_ ,       lowercase_   )






  def _a (      self   )  ->	Any:
       __UpperCamelCase   							=load_image(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'   )

       __UpperCamelCase   							=load_numpy(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_h_img2img_anime_turtle_fp16.npy'   )

       __UpperCamelCase   							=StableUnCLIPImgaImgPipeline.from_pretrained(
           'fusing/stable-unclip-2-1-h-img2img' ,       torch_dtype=torch.floataa   )
       pipe.to(lowercase_   )
       pipe.set_progress_bar_config(disable=lowercase_   )
       # stable unclip will oom when integration tests are run on a V100,
       # so turn on memory savings
       pipe.enable_attention_slicing()
       pipe.enable_sequential_cpu_offload()

       __UpperCamelCase   							=torch.Generator(device='cpu'   ).manual_seed(0   )
       __UpperCamelCase   							=pipe(lowercase_ ,       'anime turle' ,       generator=lowercase_ ,       output_type='np'   )

       __UpperCamelCase   							=output.images[0]

       assert image.shape == (768, 768, 3)

       assert_mean_pixel_difference(lowercase_ ,       lowercase_   )






  def _a (      self   )  ->	Optional[int]:
       __UpperCamelCase   							=load_image(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'   )

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

       __UpperCamelCase   							=StableUnCLIPImgaImgPipeline.from_pretrained(
           'fusing/stable-unclip-2-1-h-img2img' ,       torch_dtype=torch.floataa   )
       __UpperCamelCase   							=pipe.to(lowercase_   )
       pipe.set_progress_bar_config(disable=lowercase_   )
       pipe.enable_attention_slicing()
       pipe.enable_sequential_cpu_offload()

       __UpperCamelCase   							=pipe(
           lowercase_ ,       'anime turtle' ,       num_inference_steps=2 ,       output_type='np' ,       )

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









import os
from typing import BinaryIO, Optional, Union

import numpy as np
import pyarrow.parquet as pq

from .. import Audio, Dataset, Features, Image, NamedSplit, Value, config
from ..features.features import FeatureType, _visit
from ..formatting import query_table
from ..packaged_modules import _PACKAGED_DATASETS_MODULES
from ..packaged_modules.parquet.parquet import Parquet
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Features    ):
     __UpperCamelCase   							=np.inf

     def set_batch_size(SCREAMING_SNAKE_CASE__      : FeatureType    ) -> None:
          nonlocal batch_size
          if isinstance(SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__    ):
               __UpperCamelCase   							=min(SCREAMING_SNAKE_CASE__      ,					config.PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS    )
          elif isinstance(SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__    ):
               __UpperCamelCase   							=min(SCREAMING_SNAKE_CASE__      ,					config.PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS    )
          elif isinstance(SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__    ) and feature.dtype == "binary":
               __UpperCamelCase   							=min(SCREAMING_SNAKE_CASE__      ,					config.PARQUET_ROW_GROUP_SIZE_FOR_BINARY_DATASETS    )

     _visit(SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__    )

     return None if batch_size is np.inf else batch_size






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






    def __init__(      self ,       A_ ,       A_ = None ,       A_ = None ,       A_ = None ,       A_ = False ,       A_ = False ,       A_ = None ,       **A_ ,       )  ->	Dict:
         super().__init__(
             A_ ,       split=A_ ,       features=A_ ,       cache_dir=A_ ,       keep_in_memory=A_ ,       streaming=A_ ,       num_proc=A_ ,       **A_ ,       )
         __UpperCamelCase   							=path_or_paths if isinstance(A_ ,       A_   ) else {self.split: path_or_paths}
         __UpperCamelCase   							=_PACKAGED_DATASETS_MODULES['parquet'][1]
         __UpperCamelCase   							=Parquet(
             cache_dir=A_ ,       data_files=A_ ,       features=A_ ,       hash=A_ ,       **A_ ,       )






    def _a (      self   )  ->	List[Any]:
         # Build iterable dataset
         if self.streaming:
              __UpperCamelCase   							=self.builder.as_streaming_dataset(split=self.split   )
         # Build regular (map-style) dataset
         else:
              __UpperCamelCase   							=None
              __UpperCamelCase   							=None
              __UpperCamelCase   							=None
              __UpperCamelCase   							=None

              self.builder.download_and_prepare(
                  download_config=A_ ,       download_mode=A_ ,       verification_mode=A_ ,       base_path=A_ ,       num_proc=self.num_proc ,       )
              __UpperCamelCase   							=self.builder.as_dataset(
                  split=self.split ,       verification_mode=A_ ,       in_memory=self.keep_in_memory   )
         return dataset






class    UpperCAmelCase__	:



    """simple docstring"""






    def __init__(      self ,       A_ ,       A_ ,       A_ = None ,       **A_ ,       )  ->	List[Any]:
         __UpperCamelCase   							=dataset
         __UpperCamelCase   							=path_or_buf
         __UpperCamelCase   							=batch_size or get_writer_batch_size(dataset.features   )
         __UpperCamelCase   							=parquet_writer_kwargs






    def _a (      self   )  ->	int:
         __UpperCamelCase   							=self.batch_size if self.batch_size else config.DEFAULT_MAX_BATCH_SIZE

         if isinstance(self.path_or_buf ,       (str, bytes, os.PathLike)   ):
              with open(self.path_or_buf ,       'wb+'   ) as buffer:
                   __UpperCamelCase   							=self._write(file_obj=A_ ,       batch_size=A_ ,       **self.parquet_writer_kwargs   )
         else:
              __UpperCamelCase   							=self._write(file_obj=self.path_or_buf ,       batch_size=A_ ,       **self.parquet_writer_kwargs   )
         return written






    def _a (      self ,       A_ ,       A_ ,       **A_   )  ->	int:
         __UpperCamelCase   							=0
         __UpperCamelCase   							=parquet_writer_kwargs.pop('path_or_buf' ,       A_   )
         __UpperCamelCase   							=self.dataset.features.arrow_schema

         __UpperCamelCase   							=pq.ParquetWriter(A_ ,       schema=A_ ,       **A_   )

         for offset in logging.tqdm(
             range(0 ,       len(self.dataset   ) ,       A_   ) ,       unit='ba' ,       disable=not logging.is_progress_bar_enabled() ,       desc='Creating parquet from Arrow format' ,       ):
              __UpperCamelCase   							=query_table(
                  table=self.dataset._data ,       key=slice(A_ ,       offset + batch_size   ) ,       indices=self.dataset._indices if self.dataset._indices is not None else None ,       )
              writer.write_table(A_   )
              written += batch.nbytes
         writer.close()
         return written





