



'''simple docstring'''





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 						__snake_case(				_lowerCAmelCase		)		-> str:  # picklable for multiprocessing
      return x.sum()





def 						__snake_case(				_lowerCAmelCase		)		-> Tuple:  # picklable for multiprocessing
      return i + 1







@dataclass
class     UpperCAmelCase_       :
    """simple docstring"""







    lowercase					=		42
    lowercase					=		42







class     UpperCAmelCase_       (       _a	):
    """simple docstring"""







    def lowerCamelCase       (    self					:     str						):
          snake_case__  :    int    						=      {}
          snake_case__  :    Tuple    						=      []
          snake_case__  :    Any    						=      1
          snake_case__  :    str    						=      [1, 2]
          snake_case__  :    List[str]    						=      {"""a""": 1, """b""": 2}
          snake_case__  :    List[str]    						=      {"""a""": [1, 2], """b""": [3, 4]}
          snake_case__  :    Dict    						=      {"""a""": {"""1""": 1}, """b""": 2}
          snake_case__  :    List[Any]    						=      {"""a""": 1, """b""": 2, """c""": 3, """d""": 4}
          snake_case__  :    List[str]    						=      {}
          snake_case__  :    List[str]    						=      []
          snake_case__  :    str    						=      2
          snake_case__  :    Dict    						=      [2, 3]
          snake_case__  :    List[Any]    						=      {"""a""": 2, """b""": 3}
          snake_case__  :    Dict    						=      {"""a""": [2, 3], """b""": [4, 5]}
          snake_case__  :    List[str]    						=      {"""a""": {"""1""": 2}, """b""": 3}
          snake_case__  :    List[Any]    						=      {"""a""": 2, """b""": 3, """c""": 4, """d""": 5}
          self.assertEqual(map_nested(snake_case_ ,     snake_case_						) ,     snake_case_						)
          self.assertEqual(map_nested(snake_case_ ,     snake_case_						) ,     snake_case_						)
          self.assertEqual(map_nested(snake_case_ ,     snake_case_						) ,     snake_case_						)
          self.assertEqual(map_nested(snake_case_ ,     snake_case_						) ,     snake_case_						)
          self.assertEqual(map_nested(snake_case_ ,     snake_case_						) ,     snake_case_						)
          self.assertEqual(map_nested(snake_case_ ,     snake_case_						) ,     snake_case_						)
          self.assertEqual(map_nested(snake_case_ ,     snake_case_						) ,     snake_case_						)
          self.assertEqual(map_nested(snake_case_ ,     snake_case_						) ,     snake_case_						)

          snake_case__  :    List[str]    						=      2
          self.assertEqual(map_nested(snake_case_ ,     snake_case_ ,     num_proc=snake_case_						) ,     snake_case_						)
          self.assertEqual(map_nested(snake_case_ ,     snake_case_ ,     num_proc=snake_case_						) ,     snake_case_						)
          self.assertEqual(map_nested(snake_case_ ,     snake_case_ ,     num_proc=snake_case_						) ,     snake_case_						)
          self.assertEqual(map_nested(snake_case_ ,     snake_case_ ,     num_proc=snake_case_						) ,     snake_case_						)
          self.assertEqual(map_nested(snake_case_ ,     snake_case_ ,     num_proc=snake_case_						) ,     snake_case_						)
          self.assertEqual(map_nested(snake_case_ ,     snake_case_ ,     num_proc=snake_case_						) ,     snake_case_						)
          self.assertEqual(map_nested(snake_case_ ,     snake_case_ ,     num_proc=snake_case_						) ,     snake_case_						)
          self.assertEqual(map_nested(snake_case_ ,     snake_case_ ,     num_proc=snake_case_						) ,     snake_case_						)

          snake_case__  :    Tuple    						=      {"""a""": np.eye(2						), """b""": np.zeros(3						), """c""": np.ones(2						)}
          snake_case__  :    Union[str, Any]    						=      {"""a""": 2, """b""": 0, """c""": 2}
          snake_case__  :    Optional[Any]    						=      {
              """a""": np.eye(2						).astype(snake_case_						),
              """b""": np.zeros(3						).astype(snake_case_						),
              """c""": np.ones(2						).astype(snake_case_						),
          }
          self.assertEqual(map_nested(snake_case_ ,     snake_case_ ,     map_numpy=snake_case_						) ,     snake_case_						)
          self.assertEqual(
              {k: v.tolist() for k, v in map_nested(snake_case_ ,     snake_case_ ,     map_numpy=snake_case_						).items()} ,     {k: v.tolist() for k, v in expected_map_nested_sna_int.items()} ,     )
          self.assertEqual(map_nested(snake_case_ ,     snake_case_ ,     map_numpy=snake_case_ ,     num_proc=snake_case_						) ,     snake_case_						)
          self.assertEqual(
              {k: v.tolist() for k, v in map_nested(snake_case_ ,     snake_case_ ,     map_numpy=snake_case_ ,     num_proc=snake_case_						).items()} ,     {k: v.tolist() for k, v in expected_map_nested_sna_int.items()} ,     )
          with self.assertRaises(snake_case_						):  # can't pickle a local lambda
                map_nested(lambda snake_case_						: x + 1 ,     snake_case_ ,     num_proc=snake_case_						)




    def lowerCamelCase       (    self					:     int						):
          snake_case__  :    Tuple    						=      {"""a""": 1, """b""": 2}
          snake_case__  :    Dict    						=      {"""a""": 3, """b""": 4}
          snake_case__  :    List[str]    						=      {"""a""": 5, """b""": 6}
          snake_case__  :    Tuple    						=      sorted([("""a""", (1, 3, 5)), ("""b""", (2, 4, 6))]						)
          self.assertEqual(sorted(zip_dict(snake_case_ ,     snake_case_ ,     snake_case_						)						) ,     snake_case_						)




    def lowerCamelCase       (    self					:     int						):







          class     UpperCAmelCase_       :
              """simple docstring"""







              lowercase					=		"bar"
          snake_case__  :    Tuple    						=      Foo()
          self.assertEqual(foo.my_attr ,     """bar"""						)
          with temporary_assignment(snake_case_ ,     """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 						__snake_case(				_lowerCAmelCase     ,   _lowerCAmelCase     ,   _lowerCAmelCase		)		-> Optional[int]:
      with patch("""datasets.utils.py_utils._single_map_nested"""		) as mock_single_map_nested, patch(
          """datasets.parallel.parallel.Pool"""		) as mock_multiprocessing_pool:
            snake_case__  :    Union[str, Any]    						=      {f"{i}": i for i in range(_lowerCAmelCase		)}
            snake_case__  :    int    						=      map_nested(lambda _lowerCAmelCase		: x + 10     ,   _lowerCAmelCase     ,   num_proc=_lowerCAmelCase     ,   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     UpperCAmelCase_       (       _a	):
    """simple docstring"""







    @require_tf
    def lowerCamelCase       (    self					:     Optional[int]						):
          import tensorflow as tf
          from tensorflow.keras import layers

          snake_case__  :    Tuple    						=      layers.Dense(2						)

          def gen_random_output():
                snake_case__  :    Union[str, Any]    						=      tf.random.uniform((1, 3)						)
                return model(snake_case_						).numpy()

          with temp_seed(42 ,     set_tensorflow=snake_case_						):
                snake_case__  :    List[Any]    						=      gen_random_output()
          with temp_seed(42 ,     set_tensorflow=snake_case_						):
                snake_case__  :    List[str]    						=      gen_random_output()
          snake_case__  :    Any    						=      gen_random_output()

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




    @require_torch
    def lowerCamelCase       (    self					:     Union[str, Any]						):
          import torch

          def gen_random_output():
                snake_case__  :    List[str]    						=      torch.nn.Linear(3 ,     2						)
                snake_case__  :    Tuple    						=      torch.rand(1 ,     3						)
                return model(snake_case_						).detach().numpy()

          with temp_seed(42 ,     set_pytorch=snake_case_						):
                snake_case__  :    List[str]    						=      gen_random_output()
          with temp_seed(42 ,     set_pytorch=snake_case_						):
                snake_case__  :    List[str]    						=      gen_random_output()
          snake_case__  :    List[str]    						=      gen_random_output()

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




    def lowerCamelCase       (    self					:     Optional[Any]						):
          def gen_random_output():
                return np.random.rand(1 ,     3						)

          with temp_seed(42						):
                snake_case__  :    Union[str, Any]    						=      gen_random_output()
          with temp_seed(42						):
                snake_case__  :    List[str]    						=      gen_random_output()
          snake_case__  :    Optional[int]    						=      gen_random_output()

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





@pytest.mark.parametrize("""input_data"""     ,   [{}]		)
def 						__snake_case(				_lowerCAmelCase		)		-> Union[str, Any]:
      snake_case__  :    Any    						=      NestedDataStructure(_lowerCAmelCase		).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 						__snake_case(				_lowerCAmelCase     ,   _lowerCAmelCase		)		-> Union[str, Any]:
      snake_case__  :    Dict    						=      NestedDataStructure(_lowerCAmelCase		).flatten()
      assert output == expected_output





def 						__snake_case(				)		-> Any:
      snake_case__  :    Optional[int]    						=      A(x=1     ,   y="""foobar"""		)
      snake_case__  :    Any    						=      {"""x""": 1, """y""": """foobar"""}
      assert asdict(_lowerCAmelCase		) == expected_output

      snake_case__  :    str    						=      {"""a""": {"""b""": A(x=10     ,   y="""foo"""		)}, """c""": [A(x=20     ,   y="""bar"""		)]}
      snake_case__  :    Optional[int]    						=      {"""a""": {"""b""": {"""x""": 10, """y""": """foo"""}}, """c""": [{"""x""": 20, """y""": """bar"""}]}
      assert asdict(_lowerCAmelCase		) == expected_output

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





def 						__snake_case(				_lowerCAmelCase		)		-> Optional[int]:
      return text.split()





def 						__snake_case(				_lowerCAmelCase		)		-> Optional[int]:
      yield (time.time(), content)
      time.sleep(2		)
      yield (time.time(), content)





def 						__snake_case(				)		-> Optional[Any]:
      with Pool(2		) as pool:
            snake_case__  :    Tuple    						=      list(iflatmap_unordered(_lowerCAmelCase     ,   _split_text     ,   kwargs_iterable=[{"""text""": """hello there"""}] * 10		)		)
            assert out.count("""hello"""		) == 10
            assert out.count("""there"""		) == 10
            assert len(_lowerCAmelCase		) == 20

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

      # check that we get items as fast as possible
      with Pool(2		) as pool:
            snake_case__  :    str    						=      []
            for yield_time, content in iflatmap_unordered(
                _lowerCAmelCase     ,   _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(_lowerCAmelCase		)
            assert out.count("""a"""		) == 2
            assert out.count("""b"""		) == 2
            assert len(_lowerCAmelCase		) == 4






"""simple docstring"""






import unittest

import numpy as np
import torch

from diffusers import DDIMPipeline, DDIMScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow, torch_device

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()






class lowerCAmelCase__	(     lowercase,	unittest.TestCase						):






			'''simple docstring'''



			lowerCamelCase__      =			DDIMPipeline
			lowerCamelCase__      =			UNCONDITIONAL_IMAGE_GENERATION_PARAMS
			lowerCamelCase__      =			PipelineTesterMixin.required_optional_params - {
			    """num_images_per_prompt""",
			    """latents""",
			    """callback""",
			    """callback_steps""",
			}
			lowerCamelCase__      =			UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS
			lowerCamelCase__      =			False
			def 					A_    (  self  ):
							torch.manual_seed(0  )
							_lowerCamelCase     :			List[Any]    =    UNetaDModel(
							    block_out_channels=(32, 64)  ,      layers_per_block=2  ,      sample_size=32  ,      in_channels=3  ,      out_channels=3  ,      down_block_types=('DownBlock2D', 'AttnDownBlock2D')  ,      up_block_types=('AttnUpBlock2D', 'UpBlock2D')  ,      )
							_lowerCamelCase     :			List[str]    =    DDIMScheduler()
							_lowerCamelCase     :			Optional[int]    =    {'unet': unet, 'scheduler': scheduler}
							return components
			def 					A_    (  self  ,      lowercase  ,      lowercase=0  ):
							if str(lowercase  ).startswith('mps'  ):
											_lowerCamelCase     :			Dict    =    torch.manual_seed(lowercase  )
							else:
											_lowerCamelCase     :			List[str]    =    torch.Generator(device=lowercase  ).manual_seed(lowercase  )
							_lowerCamelCase     :			Tuple    =    {
							    'batch_size': 1,
							    'generator': generator,
							    'num_inference_steps': 2,
							    'output_type': 'numpy',
							}
							return inputs
			def 					A_    (  self  ):
							_lowerCamelCase     :			Any    =    'cpu'

							_lowerCamelCase     :			Tuple    =    self.get_dummy_components()
							_lowerCamelCase     :			Optional[Any]    =    self.pipeline_class(**lowercase  )
							pipe.to(lowercase  )
							pipe.set_progress_bar_config(disable=lowercase  )

							_lowerCamelCase     :			str    =    self.get_dummy_inputs(lowercase  )
							_lowerCamelCase     :			int    =    pipe(**lowercase  ).images
							_lowerCamelCase     :			Any    =    image[0, -3:, -3:, -1]

							self.assertEqual(image.shape  ,      (1, 32, 32, 3)  )
							_lowerCamelCase     :			Tuple    =    np.array(
							    [1.000E00, 5.717E-01, 4.717E-01, 1.000E00, 0.000E00, 1.000E00, 3.000E-04, 0.000E00, 9.000E-04]  )
							_lowerCamelCase     :			str    =    np.abs(image_slice.flatten() - expected_slice  ).max()
							self.assertLessEqual(lowercase  ,      1E-3  )
			def 					A_    (  self  ):
							super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3  )
			def 					A_    (  self  ):
							super().test_save_load_local(expected_max_difference=3E-3  )
			def 					A_    (  self  ):
							super().test_save_load_optional_components(expected_max_difference=3E-3  )






			def 					A_    (  self  ):
							super().test_inference_batch_single_identical(expected_max_diff=3E-3  )






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






			'''simple docstring'''



			def 					A_    (  self  ):
							_lowerCamelCase     :			Optional[Any]    =    'google/ddpm-cifar10-32'

							_lowerCamelCase     :			Optional[Any]    =    UNetaDModel.from_pretrained(lowercase  )
							_lowerCamelCase     :			Dict    =    DDIMScheduler()

							_lowerCamelCase     :			Dict    =    DDIMPipeline(unet=lowercase  ,      scheduler=lowercase  )
							ddim.to(lowercase  )
							ddim.set_progress_bar_config(disable=lowercase  )

							_lowerCamelCase     :			List[str]    =    torch.manual_seed(0  )
							_lowerCamelCase     :			str    =    ddim(generator=lowercase  ,      eta=0.0  ,      output_type='numpy'  ).images

							_lowerCamelCase     :			Dict    =    image[0, -3:, -3:, -1]

							assert image.shape == (1, 32, 32, 3)
							_lowerCamelCase     :			List[Any]    =    np.array([0.17_23, 0.16_17, 0.16_00, 0.16_26, 0.14_97, 0.15_13, 0.15_05, 0.14_42, 0.14_53]  )

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






			def 					A_    (  self  ):
							_lowerCamelCase     :			Optional[int]    =    'google/ddpm-ema-bedroom-256'

							_lowerCamelCase     :			str    =    UNetaDModel.from_pretrained(lowercase  )
							_lowerCamelCase     :			str    =    DDIMScheduler.from_pretrained(lowercase  )

							_lowerCamelCase     :			Optional[int]    =    DDIMPipeline(unet=lowercase  ,      scheduler=lowercase  )
							ddpm.to(lowercase  )
							ddpm.set_progress_bar_config(disable=lowercase  )

							_lowerCamelCase     :			Tuple    =    torch.manual_seed(0  )
							_lowerCamelCase     :			int    =    ddpm(generator=lowercase  ,      output_type='numpy'  ).images

							_lowerCamelCase     :			Dict    =    image[0, -3:, -3:, -1]

							assert image.shape == (1, 256, 256, 3)
							_lowerCamelCase     :			str    =    np.array([0.00_60, 0.02_01, 0.03_44, 0.00_24, 0.00_18, 0.00_02, 0.00_22, 0.00_00, 0.00_69]  )

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