

import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionDualGuidedPipeline
from diffusers.utils.testing_utils import load_image, nightly, require_torch_gpu, torch_device


_a			=       False






@nightly
@require_torch_gpu
class 							A_			(		unittest.TestCase ):


     def   UpperCAmelCase  (							self      :		Any )							->     List[Any]:
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()


     def   UpperCAmelCase  (							self      :		str )							->     List[str]:
      __lowerCAmelCase:       List[Any]            =				VersatileDiffusionDualGuidedPipeline.from_pretrained('shi-labs/versatile-diffusion' )
      # remove text_unet
      pipe.remove_unused_weights()
      pipe.to(UpperCAmelCase )
      pipe.set_progress_bar_config(disable=UpperCAmelCase )

      __lowerCAmelCase:       int            =				load_image(
          'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg' )

      __lowerCAmelCase:       Optional[int]            =				torch.manual_seed(0 )
      __lowerCAmelCase:       Optional[Any]            =				pipe(
          prompt='first prompt'      ,  image=UpperCAmelCase      ,  text_to_image_strength=0.75      ,  generator=UpperCAmelCase      ,  guidance_scale=7.5      ,  num_inference_steps=2      ,  output_type='numpy'      ,  ).images

      with tempfile.TemporaryDirectory() as tmpdirname:
       pipe.save_pretrained(UpperCAmelCase )
       __lowerCAmelCase:       Any            =				VersatileDiffusionDualGuidedPipeline.from_pretrained(UpperCAmelCase )

      pipe.to(UpperCAmelCase )
      pipe.set_progress_bar_config(disable=UpperCAmelCase )

      __lowerCAmelCase:       Any            =				generator.manual_seed(0 )
      __lowerCAmelCase:       Tuple            =				pipe(
          prompt='first prompt'      ,  image=UpperCAmelCase      ,  text_to_image_strength=0.75      ,  generator=UpperCAmelCase      ,  guidance_scale=7.5      ,  num_inference_steps=2      ,  output_type='numpy'      ,  ).images

      assert np.abs(image - new_image ).sum() < 1E-5, "Models don't have the same forward pass"


     def   UpperCAmelCase  (							self      :		int )							->     Tuple:
      __lowerCAmelCase:       List[str]            =				VersatileDiffusionDualGuidedPipeline.from_pretrained('shi-labs/versatile-diffusion' )
      pipe.remove_unused_weights()
      pipe.to(UpperCAmelCase )
      pipe.set_progress_bar_config(disable=UpperCAmelCase )

      __lowerCAmelCase:       str            =				"cyberpunk 2077"
      __lowerCAmelCase:       Any            =				load_image(
          'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg' )
      __lowerCAmelCase:       Optional[int]            =				torch.manual_seed(0 )
      __lowerCAmelCase:       str            =				pipe(
          prompt=UpperCAmelCase      ,  image=UpperCAmelCase      ,  text_to_image_strength=0.75      ,  generator=UpperCAmelCase      ,  guidance_scale=7.5      ,  num_inference_steps=5_0      ,  output_type='numpy'      ,  ).images

      __lowerCAmelCase:       int            =				image[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

      assert image.shape == (1, 5_1_2, 5_1_2, 3)
      __lowerCAmelCase:       Optional[Any]            =				np.array([0.0787, 0.0849, 0.0826, 0.0812, 0.0807, 0.0795, 0.0818, 0.0798, 0.0779] )

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








import json
import os
from dataclasses import dataclass

import numpy as np
import pyarrow as pa

import datasets
from utils import get_duration


_a			=       1_0_0_0_0_0_0_0_0_0_0_0
_a			=       1_0_0_0_0

_a	,     _a			=       os.path.split(__file__)
_a			=       os.path.join(RESULTS_BASEPATH, '''results''', RESULTS_FILENAME.replace('''.py''', '''.json'''))







def _a			(	SCREAMING_SNAKE_CASE		:      int   , SCREAMING_SNAKE_CASE		:      int		)   ->   datasets.Dataset:

 """simple docstring"""






 __lowerCAmelCase:       List[Any]            =				pa.Table.from_pydict({'col': [0] * chunk_size}		)
 __lowerCAmelCase:       Dict            =				pa.concat_tables([table] * (num_examples // chunk_size)		)
 return datasets.Dataset(SCREAMING_SNAKE_CASE   , fingerprint='table_100B'		)






@dataclass
class 							A_			:
     _lowercase      :					int
     _lowercase      :					int
     _lowercase      :					int
     _lowercase      :					int


     def   UpperCAmelCase  (							self      :		Optional[Any] )							->     Dict:
      __lowerCAmelCase:       Dict            =				np.random.default_rng(self.seed )
      __lowerCAmelCase:       List[Any]            =				rng.integers(low=self.low      ,  high=self.high      ,  size=self.size ).tolist()


     def __iter__(							self      :		Any )							->     Optional[int]:
      return iter(self._sampled_values )


     def __len__(							self      :		Union[str, Any] )							->     Dict:
      return self.size







@get_duration
def _a			(	SCREAMING_SNAKE_CASE		:      datasets.Dataset		)   ->   List[str]:

 """simple docstring"""






 __lowerCAmelCase:       List[Any]            =				dataset[0]







@get_duration
def _a			(	SCREAMING_SNAKE_CASE		:      datasets.Dataset		)   ->   Dict:

 """simple docstring"""






 __lowerCAmelCase:       Dict            =				dataset[-1]







@get_duration
def _a			(	SCREAMING_SNAKE_CASE		:      datasets.Dataset		)   ->   Any:

 """simple docstring"""






 __lowerCAmelCase:       Union[str, Any]            =				dataset[range(len(SCREAMING_SNAKE_CASE		) // 2   , len(SCREAMING_SNAKE_CASE		) // 2 + 10_24		)]







@get_duration
def _a			(	SCREAMING_SNAKE_CASE		:      datasets.Dataset		)   ->   Optional[Any]:

 """simple docstring"""






 __lowerCAmelCase:       int            =				dataset[RandIter(0   , len(SCREAMING_SNAKE_CASE		)   , 10_24   , seed=42		)]







def _a			(	)   ->   Tuple:

 """simple docstring"""






 __lowerCAmelCase:       str            =				{'num examples': SPEED_TEST_N_EXAMPLES}
 __lowerCAmelCase:       Tuple            =				(get_first_row, get_last_row, get_batch_of_aaaa_rows, get_batch_of_aaaa_random_rows)
 print('generating dataset'		)
 __lowerCAmelCase:       str            =				generate_100B_dataset(num_examples=SCREAMING_SNAKE_CASE   , chunk_size=SCREAMING_SNAKE_CASE		)
 print('Functions'		)
 for func in functions:
  print(func.__name__		)
  __lowerCAmelCase:       Dict            =				func(SCREAMING_SNAKE_CASE		)

 with open(SCREAMING_SNAKE_CASE   , 'wb'		) as f:
  f.write(json.dumps(SCREAMING_SNAKE_CASE		).encode('utf-8'		)		)


if __name__ == "__main__":  # useful to run the profiler
  benchmark_table_100B()





