






import time
from dataclasses import dataclass
from multiprocessing import Pool
from unittest import TestCase
from unittest.mock import patch

import multiprocess
import numpy as np
import pytest

from datasets.utils.py_utils import (
    NestedDataStructure,
    asdict,
    iflatmap_unordered,
    map_nested,
    temp_seed,
    temporary_assignment,
    zip_dict,
)

from .utils import require_tf, require_torch







def 		SCREAMING_SNAKE_CASE							( __lowerCAmelCase      )			->     List[Any]:  # picklable for multiprocessing
     return x.sum()







def 		SCREAMING_SNAKE_CASE							( __lowerCAmelCase      )			->     Union[str, Any]:  # picklable for multiprocessing
     return i + 1





@dataclass
class       __a  :
 _lowerCAmelCase	:	int
 _lowerCAmelCase	:	str







class       __a  (   A__	):







 def  __lowercase (		self :       str						):


      '''simple docstring'''






      UpperCamelCase__      :       Optional[int]	    =  {}
      UpperCamelCase__      :       Optional[int]	    =  []
      UpperCamelCase__      :       Union[str, Any]	    =  1
      UpperCamelCase__      :       List[str]	    =  [1, 2]
      UpperCamelCase__      :       Dict	    =  {'a': 1, 'b': 2}
      UpperCamelCase__      :       Optional[Any]	    =  {'a': [1, 2], 'b': [3, 4]}
      UpperCamelCase__      :       str	    =  {'a': {'1': 1}, 'b': 2}
      UpperCamelCase__      :       int	    =  {'a': 1, 'b': 2, 'c': 3, 'd': 4}
      UpperCamelCase__      :       Optional[Any]	    =  {}
      UpperCamelCase__      :       Optional[int]	    =  []
      UpperCamelCase__      :       Optional[int]	    =  2
      UpperCamelCase__      :       Optional[int]	    =  [2, 3]
      UpperCamelCase__      :       Optional[int]	    =  {'a': 2, 'b': 3}
      UpperCamelCase__      :       Any	    =  {'a': [2, 3], 'b': [4, 5]}
      UpperCamelCase__      :       List[str]	    =  {'a': {'1': 2}, 'b': 3}
      UpperCamelCase__      :       Optional[int]	    =  {'a': 2, 'b': 3, 'c': 4, 'd': 5}
      self.assertEqual(map_nested(_A							,       _A						)							,       _A						)
      self.assertEqual(map_nested(_A							,       _A						)							,       _A						)
      self.assertEqual(map_nested(_A							,       _A						)							,       _A						)
      self.assertEqual(map_nested(_A							,       _A						)							,       _A						)
      self.assertEqual(map_nested(_A							,       _A						)							,       _A						)
      self.assertEqual(map_nested(_A							,       _A						)							,       _A						)
      self.assertEqual(map_nested(_A							,       _A						)							,       _A						)
      self.assertEqual(map_nested(_A							,       _A						)							,       _A						)

      UpperCamelCase__      :       Tuple	    =  2
      self.assertEqual(map_nested(_A							,       _A							,       num_proc=_A						)							,       _A						)
      self.assertEqual(map_nested(_A							,       _A							,       num_proc=_A						)							,       _A						)
      self.assertEqual(map_nested(_A							,       _A							,       num_proc=_A						)							,       _A						)
      self.assertEqual(map_nested(_A							,       _A							,       num_proc=_A						)							,       _A						)
      self.assertEqual(map_nested(_A							,       _A							,       num_proc=_A						)							,       _A						)
      self.assertEqual(map_nested(_A							,       _A							,       num_proc=_A						)							,       _A						)
      self.assertEqual(map_nested(_A							,       _A							,       num_proc=_A						)							,       _A						)
      self.assertEqual(map_nested(_A							,       _A							,       num_proc=_A						)							,       _A						)

      UpperCamelCase__      :       Optional[int]	    =  {'a': np.eye(2						), 'b': np.zeros(3						), 'c': np.ones(2						)}
      UpperCamelCase__      :       Dict	    =  {'a': 2, 'b': 0, 'c': 2}
      UpperCamelCase__      :       List[str]	    =  {
          'a': np.eye(2						).astype(_A						),
          'b': np.zeros(3						).astype(_A						),
          'c': np.ones(2						).astype(_A						),
      }
      self.assertEqual(map_nested(_A							,       _A							,       map_numpy=_A						)							,       _A						)
      self.assertEqual(
          {k: v.tolist() for k, v in map_nested(_A							,       _A							,       map_numpy=_A						).items()}							,       {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}							,       )
      self.assertEqual(map_nested(_A							,       _A							,       map_numpy=_A							,       num_proc=_A						)							,       _A						)
      self.assertEqual(
          {k: v.tolist() for k, v in map_nested(_A							,       _A							,       map_numpy=_A							,       num_proc=_A						).items()}							,       {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}							,       )
      with self.assertRaises(_A						):  # can't pickle a local lambda
           map_nested(lambda SCREAMING_SNAKE_CASE						: x + 1							,       _A							,       num_proc=_A						)







 def  __lowercase (		self :       Dict						):


      '''simple docstring'''






      UpperCamelCase__      :       Union[str, Any]	    =  {'a': 1, 'b': 2}
      UpperCamelCase__      :       Tuple	    =  {'a': 3, 'b': 4}
      UpperCamelCase__      :       List[Any]	    =  {'a': 5, 'b': 6}
      UpperCamelCase__      :       Dict	    =  sorted([("a", (1, 3, 5)), ("b", (2, 4, 6))]						)
      self.assertEqual(sorted(zip_dict(_A							,       _A							,       _A						)						)							,       _A						)







 def  __lowercase (		self :       List[Any]						):


      '''simple docstring'''






      class       __a  :
       _lowerCAmelCase	:	Optional[int]				=			"""bar"""







      UpperCamelCase__      :       Optional[Any]	    =  Foo()
      self.assertEqual(foo.my_attr							,       "bar"						)
      with temporary_assignment(_A							,       "my_attr"							,       "BAR"						):
           self.assertEqual(foo.my_attr							,       "BAR"						)
      self.assertEqual(foo.my_attr							,       "bar"						)







@pytest.mark.parametrize(
    "iterable_length, num_proc, expected_num_proc" ,   [
        (1, None, 1),
        (1, 1, 1),
        (2, None, 1),
        (2, 1, 1),
        (2, 2, 1),
        (2, 3, 1),
        (3, 2, 1),
        (16, 16, 16),
        (16, 17, 16),
        (17, 16, 16),
    ] ,   )
def 		SCREAMING_SNAKE_CASE							( __lowerCAmelCase ,   __lowerCAmelCase ,   __lowerCAmelCase      )			->     Dict:





     with patch("datasets.utils.py_utils._single_map_nested"      ) as mock_single_map_nested, patch(
         "datasets.parallel.parallel.Pool"      ) as mock_multiprocessing_pool:
          UpperCamelCase__      :       int	    =  {f'{i}': i for i in range(UpperCAmelCase_      )}
          UpperCamelCase__      :       str	    =  map_nested(lambda __lowerCAmelCase      : x + 10 ,   UpperCAmelCase_ ,   num_proc=UpperCAmelCase_ ,   parallel_min_length=16      )
          if expected_num_proc == 1:
               assert mock_single_map_nested.called
               assert not mock_multiprocessing_pool.called
          else:
               assert not mock_single_map_nested.called
               assert mock_multiprocessing_pool.called
               assert mock_multiprocessing_pool.call_args[0][0] == expected_num_proc





class       __a  (   A__	):







 @require_tf
 def  __lowercase (		self :       Union[str, Any]						):


      '''simple docstring'''






      import tensorflow as tf
      from tensorflow.keras import layers

      UpperCamelCase__      :       Optional[Any]	    =  layers.Dense(2						)

      def gen_random_output():
           UpperCamelCase__      :       int	    =  tf.random.uniform((1, 3)						)
           return model(_A						).numpy()

      with temp_seed(42							,       set_tensorflow=_A						):
           UpperCamelCase__      :       Tuple	    =  gen_random_output()
      with temp_seed(42							,       set_tensorflow=_A						):
           UpperCamelCase__      :       Union[str, Any]	    =  gen_random_output()
      UpperCamelCase__      :       str	    =  gen_random_output()

      np.testing.assert_equal(_A							,       _A						)
      self.assertGreater(np.abs(outa - outa						).sum()							,       0						)







 @require_torch
 def  __lowercase (		self :       List[str]						):


      '''simple docstring'''






      import torch

      def gen_random_output():
           UpperCamelCase__      :       Optional[Any]	    =  torch.nn.Linear(3							,       2						)
           UpperCamelCase__      :       Dict	    =  torch.rand(1							,       3						)
           return model(_A						).detach().numpy()

      with temp_seed(42							,       set_pytorch=_A						):
           UpperCamelCase__      :       Optional[int]	    =  gen_random_output()
      with temp_seed(42							,       set_pytorch=_A						):
           UpperCamelCase__      :       Optional[int]	    =  gen_random_output()
      UpperCamelCase__      :       Dict	    =  gen_random_output()

      np.testing.assert_equal(_A							,       _A						)
      self.assertGreater(np.abs(outa - outa						).sum()							,       0						)







 def  __lowercase (		self :       Dict						):


      '''simple docstring'''






      def gen_random_output():
           return np.random.rand(1							,       3						)

      with temp_seed(42						):
           UpperCamelCase__      :       Dict	    =  gen_random_output()
      with temp_seed(42						):
           UpperCamelCase__      :       str	    =  gen_random_output()
      UpperCamelCase__      :       Optional[Any]	    =  gen_random_output()

      np.testing.assert_equal(_A							,       _A						)
      self.assertGreater(np.abs(outa - outa						).sum()							,       0						)







@pytest.mark.parametrize("input_data" ,   [{}]      )
def 		SCREAMING_SNAKE_CASE							( __lowerCAmelCase      )			->     Optional[Any]:
     UpperCamelCase__      :       Any	    =  NestedDataStructure(UpperCAmelCase_      ).data
     assert output_data == input_data







@pytest.mark.parametrize(
    "data, expected_output" ,   [
        ({}, []),
        ([], []),
        ("foo", ["foo"]),
        (["foo", "bar"], ["foo", "bar"]),
        ([["foo", "bar"]], ["foo", "bar"]),
        ([[["foo"], ["bar"]]], ["foo", "bar"]),
        ([[["foo"], "bar"]], ["foo", "bar"]),
        ({"a": 1, "b": 2}, [1, 2]),
        ({"a": [1, 2], "b": [3, 4]}, [1, 2, 3, 4]),
        ({"a": [[1, 2]], "b": [[3, 4]]}, [1, 2, 3, 4]),
        ({"a": [[1, 2]], "b": [3, 4]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [[[3], [4]]]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [[3, 4]]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [3, 4]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [3, [4]]}, [1, 2, 3, 4]),
        ({"a": {"1": 1}, "b": 2}, [1, 2]),
        ({"a": {"1": [1]}, "b": 2}, [1, 2]),
        ({"a": {"1": [1]}, "b": [2]}, [1, 2]),
    ] ,   )
def 		SCREAMING_SNAKE_CASE							( __lowerCAmelCase ,   __lowerCAmelCase      )			->     Dict:
     UpperCamelCase__      :       List[str]	    =  NestedDataStructure(UpperCAmelCase_      ).flatten()
     assert output == expected_output







def 		SCREAMING_SNAKE_CASE							( )			->     Any:
     UpperCamelCase__      :       List[Any]	    =  A(x=1 ,   y="foobar"      )
     UpperCamelCase__      :       int	    =  {'x': 1, 'y': 'foobar'}
     assert asdict(UpperCAmelCase_      ) == expected_output

     UpperCamelCase__      :       Optional[int]	    =  {'a': {'b': A(x=10 ,   y="foo"      )}, 'c': [A(x=20 ,   y="bar"      )]}
     UpperCamelCase__      :       Union[str, Any]	    =  {'a': {'b': {'x': 10, 'y': 'foo'}}, 'c': [{'x': 20, 'y': 'bar'}]}
     assert asdict(UpperCAmelCase_      ) == expected_output

     with pytest.raises(UpperCAmelCase_      ):
          asdict([1, A(x=10 ,   y="foo"      )]      )







def 		SCREAMING_SNAKE_CASE							( __lowerCAmelCase      )			->     Tuple:
     return text.split()







def 		SCREAMING_SNAKE_CASE							( __lowerCAmelCase      )			->     List[Any]:
     yield (time.time(), content)
     time.sleep(2      )
     yield (time.time(), content)







def 		SCREAMING_SNAKE_CASE							( )			->     List[Any]:





     with Pool(2      ) as pool:
          UpperCamelCase__      :       Any	    =  list(iflatmap_unordered(UpperCAmelCase_ ,   _split_text ,   kwargs_iterable=[{"text": "hello there"}] * 10      )      )
          assert out.count("hello"      ) == 10
          assert out.count("there"      ) == 10
          assert len(UpperCAmelCase_      ) == 20

     # check multiprocess from pathos (uses dill for pickling)
     with multiprocess.Pool(2      ) as pool:
          UpperCamelCase__      :       int	    =  list(iflatmap_unordered(UpperCAmelCase_ ,   _split_text ,   kwargs_iterable=[{"text": "hello there"}] * 10      )      )
          assert out.count("hello"      ) == 10
          assert out.count("there"      ) == 10
          assert len(UpperCAmelCase_      ) == 20

     # check that we get items as fast as possible
     with Pool(2      ) as pool:
          UpperCamelCase__      :       Union[str, Any]	    =  []
          for yield_time, content in iflatmap_unordered(
              UpperCAmelCase_ ,   _aseconds_generator_of_aitems_with_timing ,   kwargs_iterable=[{"content": "a"}, {"content": "b"}]      ):
               assert yield_time < time.time() + 0.1, "we should each item directly after it was yielded"
               out.append(UpperCAmelCase_      )
          assert out.count("a"      ) == 2
          assert out.count("b"      ) == 2
          assert len(UpperCAmelCase_      ) == 4







import gc
import tempfile
import unittest

import numpy as np
import torch

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


lowerCamelCase     :							Optional[Any]    =False





class       __a  (   unittest.TestCase	):
 pass







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







 def  __lowercase (		self :       Any						):


      '''simple docstring'''






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







 def  __lowercase (		self :       Dict						):


      '''simple docstring'''






      UpperCamelCase__      :       int	    =  VersatileDiffusionPipeline.from_pretrained("shi-labs/versatile-diffusion"							,       torch_dtype=torch.floataa						)
      pipe.to(SCREAMING_SNAKE_CASE						)
      pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

      UpperCamelCase__      :       Optional[int]	    =  load_image(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg"						)

      UpperCamelCase__      :       Dict	    =  torch.manual_seed(0						)
      UpperCamelCase__      :       str	    =  pipe.dual_guided(
          prompt="first prompt"							,       image=SCREAMING_SNAKE_CASE							,       text_to_image_strength=0.7_5							,       generator=SCREAMING_SNAKE_CASE							,       guidance_scale=7.5							,       num_inference_steps=2							,       output_type="numpy"							,       ).images

      with tempfile.TemporaryDirectory() as tmpdirname:
           pipe.save_pretrained(SCREAMING_SNAKE_CASE						)
           UpperCamelCase__      :       Tuple	    =  VersatileDiffusionPipeline.from_pretrained(SCREAMING_SNAKE_CASE							,       torch_dtype=torch.floataa						)
      pipe.to(SCREAMING_SNAKE_CASE						)
      pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

      UpperCamelCase__      :       int	    =  generator.manual_seed(0						)
      UpperCamelCase__      :       Dict	    =  pipe.dual_guided(
          prompt="first prompt"							,       image=SCREAMING_SNAKE_CASE							,       text_to_image_strength=0.7_5							,       generator=SCREAMING_SNAKE_CASE							,       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  __lowercase (		self :       Optional[int]						):


      '''simple docstring'''






      UpperCamelCase__      :       Optional[int]	    =  VersatileDiffusionPipeline.from_pretrained("shi-labs/versatile-diffusion"							,       torch_dtype=torch.floataa						)
      pipe.to(SCREAMING_SNAKE_CASE						)
      pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

      UpperCamelCase__      :       Tuple	    =  "cyberpunk 2077"
      UpperCamelCase__      :       Any	    =  load_image(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg"						)
      UpperCamelCase__      :       int	    =  torch.manual_seed(0						)
      UpperCamelCase__      :       Any	    =  pipe.dual_guided(
          prompt=SCREAMING_SNAKE_CASE							,       image=SCREAMING_SNAKE_CASE							,       text_to_image_strength=0.7_5							,       generator=SCREAMING_SNAKE_CASE							,       guidance_scale=7.5							,       num_inference_steps=50							,       output_type="numpy"							,       ).images

      UpperCamelCase__      :       Tuple	    =  image[0, 2_53:2_56, 2_53:2_56, -1]

      assert image.shape == (1, 5_12, 5_12, 3)
      UpperCamelCase__      :       Optional[int]	    =  np.array([0.1_4_4_8, 0.1_6_1_9, 0.1_7_4_1, 0.1_0_8_6, 0.1_1_4_7, 0.1_1_2_8, 0.1_1_9_9, 0.1_1_6_5, 0.1_0_0_1]						)

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

      UpperCamelCase__      :       List[str]	    =  "A painting of a squirrel eating a burger "
      UpperCamelCase__      :       List[Any]	    =  torch.manual_seed(0						)
      UpperCamelCase__      :       Dict	    =  pipe.text_to_image(
          prompt=SCREAMING_SNAKE_CASE							,       generator=SCREAMING_SNAKE_CASE							,       guidance_scale=7.5							,       num_inference_steps=50							,       output_type="numpy"						).images

      UpperCamelCase__      :       Union[str, Any]	    =  image[0, 2_53:2_56, 2_53:2_56, -1]

      assert image.shape == (1, 5_12, 5_12, 3)
      UpperCamelCase__      :       Dict	    =  np.array([0.3_3_6_7, 0.3_1_6_9, 0.2_6_5_6, 0.3_8_7_0, 0.4_7_9_0, 0.3_7_9_6, 0.4_0_0_9, 0.4_8_7_8, 0.4_7_7_8]						)

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

      UpperCamelCase__      :       Any	    =  pipe.image_variation(SCREAMING_SNAKE_CASE							,       generator=SCREAMING_SNAKE_CASE							,       output_type="numpy"						).images

      UpperCamelCase__      :       Optional[Any]	    =  image[0, 2_53:2_56, 2_53:2_56, -1]

      assert image.shape == (1, 5_12, 5_12, 3)
      UpperCamelCase__      :       int	    =  np.array([0.3_0_7_6, 0.3_1_2_3, 0.3_2_8_4, 0.3_7_8_2, 0.3_7_7_0, 0.3_8_9_4, 0.4_2_9_7, 0.4_3_3_1, 0.4_4_5_6]						)

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