



"""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  lowerCamelCase_						(UpperCamelCase__		:		List[str]							):  # picklable for multiprocessing
    return x.sum()





def  lowerCamelCase_						(UpperCamelCase__		:		str							):  # picklable for multiprocessing
    return i + 1
@dataclass
class 							_UpperCAmelCase							:





     '''simple docstring'''



     a__									=42
     a__									=42



class 							_UpperCAmelCase							( a      ):





     '''simple docstring'''



     def  __lowerCAmelCase	(					self		)							->					Optional[int]:
         _UpperCAmelCase							:  str											=    {}
         _UpperCAmelCase							:  int											=    []
         _UpperCAmelCase							:  Dict											=    1
         _UpperCAmelCase							:  Tuple											=    [1, 2]
         _UpperCAmelCase							:  Dict											=    {'''a''': 1, '''b''': 2}
         _UpperCAmelCase							:  Optional[int]											=    {'''a''': [1, 2], '''b''': [3, 4]}
         _UpperCAmelCase							:  str											=    {'''a''': {'''1''': 1}, '''b''': 2}
         _UpperCAmelCase							:  Optional[int]											=    {'''a''': 1, '''b''': 2, '''c''': 3, '''d''': 4}
         _UpperCAmelCase							:  List[Any]											=    {}
         _UpperCAmelCase							:  Union[str, Any]											=    []
         _UpperCAmelCase							:  Union[str, Any]											=    2
         _UpperCAmelCase							:  int											=    [2, 3]
         _UpperCAmelCase							:  str											=    {'''a''': 2, '''b''': 3}
         _UpperCAmelCase							:  Optional[Any]											=    {'''a''': [2, 3], '''b''': [4, 5]}
         _UpperCAmelCase							:  Union[str, Any]											=    {'''a''': {'''1''': 2}, '''b''': 3}
         _UpperCAmelCase							:  Dict											=    {'''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							:  List[Any]											=    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							:  Tuple											=    {'''a''': np.eye(2		), '''b''': np.zeros(3		), '''c''': np.ones(2		)}
         _UpperCAmelCase							:  Dict											=    {'''a''': 2, '''b''': 0, '''c''': 2}
         _UpperCAmelCase							:  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 A		: x + 1       ,	A       ,	num_proc=A		)





     def  __lowerCAmelCase	(					self		)							->					int:
         _UpperCAmelCase							:  Optional[Any]											=    {'''a''': 1, '''b''': 2}
         _UpperCAmelCase							:  Optional[Any]											=    {'''a''': 3, '''b''': 4}
         _UpperCAmelCase							:  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  __lowerCAmelCase	(					self		)							->					Union[str, Any]:
         class 							_UpperCAmelCase							:





              '''simple docstring'''



              a__									='''bar'''



         _UpperCAmelCase							:  Tuple											=    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  lowerCamelCase_						(UpperCamelCase__		:		str     ,       UpperCamelCase__		:		str     ,       UpperCamelCase__		:		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:
        _UpperCAmelCase							:  List[Any]											=    {F'{i}': i for i in range(UpperCamelCase__							)}
        _UpperCAmelCase							:  Optional[Any]											=    map_nested(lambda UpperCamelCase__							: 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 							_UpperCAmelCase							( a      ):





     '''simple docstring'''



     @require_tf
     def  __lowerCAmelCase	(					self		)							->					Union[str, Any]:
         import tensorflow as tf
         from tensorflow.keras import layers

         _UpperCAmelCase							:  Optional[int]											=    layers.Dense(2		)

         def gen_random_output():
             _UpperCAmelCase							:  Tuple											=    tf.random.uniform((1, 3)		)
             return model(A		).numpy()

         with temp_seed(4_2       ,	set_tensorflow=A		):
             _UpperCAmelCase							:  List[str]											=    gen_random_output()
         with temp_seed(4_2       ,	set_tensorflow=A		):
             _UpperCAmelCase							:  Optional[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  __lowerCAmelCase	(					self		)							->					List[str]:
         import torch

         def gen_random_output():
             _UpperCAmelCase							:  Tuple											=    torch.nn.Linear(3       ,	2		)
             _UpperCAmelCase							:  Tuple											=    torch.rand(1       ,	3		)
             return model(A		).detach().numpy()

         with temp_seed(4_2       ,	set_pytorch=A		):
             _UpperCAmelCase							:  Optional[Any]											=    gen_random_output()
         with temp_seed(4_2       ,	set_pytorch=A		):
             _UpperCAmelCase							:  Optional[Any]											=    gen_random_output()
         _UpperCAmelCase							:  str											=    gen_random_output()

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





     def  __lowerCAmelCase	(					self		)							->					Any:
         def gen_random_output():
             return np.random.rand(1       ,	3		)

         with temp_seed(4_2		):
             _UpperCAmelCase							:  List[str]											=    gen_random_output()
         with temp_seed(4_2		):
             _UpperCAmelCase							:  int											=    gen_random_output()
         _UpperCAmelCase							:  Dict											=    gen_random_output()

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





@pytest.mark.parametrize('''input_data'''     ,       [{}]							)
def  lowerCamelCase_						(UpperCamelCase__		:		str							):
    _UpperCAmelCase							:  Optional[int]											=    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  lowerCamelCase_						(UpperCamelCase__		:		Union[str, Any]     ,       UpperCamelCase__		:		Tuple							):
    _UpperCAmelCase							:  str											=    NestedDataStructure(UpperCamelCase__							).flatten()
    assert output == expected_output





def  lowerCamelCase_						():
    _UpperCAmelCase							:  Dict											=    A(x=1     ,       y='''foobar'''							)
    _UpperCAmelCase							:  str											=    {'''x''': 1, '''y''': '''foobar'''}
    assert asdict(UpperCamelCase__							) == expected_output

    _UpperCAmelCase							:  int											=    {'''a''': {'''b''': A(x=10     ,       y='''foo'''							)}, '''c''': [A(x=20     ,       y='''bar'''							)]}
    _UpperCAmelCase							:  Tuple											=    {'''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  lowerCamelCase_						(UpperCamelCase__		:		str							):
    return text.split()





def  lowerCamelCase_						(UpperCamelCase__		:		Tuple							):
    yield (time.time(), content)
    time.sleep(2							)
    yield (time.time(), content)





def  lowerCamelCase_						():
    with Pool(2							) as pool:
        _UpperCAmelCase							:  Tuple											=    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							:  str											=    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							:  Dict											=    []
        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







"""simple docstring"""
import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionAttendAndExcitePipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_numpy, skip_mps, slow
from diffusers.utils.testing_utils import require_torch_gpu

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


_lowerCAmelCase    :Any     =					False
@skip_mps
class 							_UpperCAmelCase							( a		,a		,a		,unittest.TestCase      ):





     '''simple docstring'''



     a__									=StableDiffusionAttendAndExcitePipeline
     a__									=False
     a__									=TEXT_TO_IMAGE_PARAMS
     a__									=TEXT_TO_IMAGE_BATCH_PARAMS.union({'''token_indices'''}      )
     a__									=TEXT_TO_IMAGE_IMAGE_PARAMS
     a__									=TEXT_TO_IMAGE_IMAGE_PARAMS





     @classmethod
     def  __lowerCAmelCase	(					cls		)							->					List[str]:
         super().setUpClass()
         torch.use_deterministic_algorithms(A		)





     @classmethod
     def  __lowerCAmelCase	(					cls		)							->					Union[str, Any]:
         super().tearDownClass()
         torch.use_deterministic_algorithms(A		)





     def  __lowerCAmelCase	(					self		)							->					Tuple:
         torch.manual_seed(0		)
         _UpperCAmelCase							:  Optional[int]											=    UNetaDConditionModel(
             block_out_channels=(3_2, 6_4)       ,	layers_per_block=1       ,	sample_size=3_2       ,	in_channels=4       ,	out_channels=4       ,	down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')       ,	up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')       ,	cross_attention_dim=3_2       ,	attention_head_dim=(2, 4)       ,	use_linear_projection=A       ,	)
         _UpperCAmelCase							:  List[Any]											=    DDIMScheduler(
             beta_start=0.00_085       ,	beta_end=0.012       ,	beta_schedule='''scaled_linear'''       ,	clip_sample=A       ,	set_alpha_to_one=A       ,	)
         torch.manual_seed(0		)
         _UpperCAmelCase							:  int											=    AutoencoderKL(
             block_out_channels=[3_2, 6_4]       ,	in_channels=3       ,	out_channels=3       ,	down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']       ,	up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']       ,	latent_channels=4       ,	sample_size=1_2_8       ,	)
         torch.manual_seed(0		)
         _UpperCAmelCase							:  int											=    CLIPTextConfig(
             bos_token_id=0       ,	eos_token_id=2       ,	hidden_size=3_2       ,	intermediate_size=3_7       ,	layer_norm_eps=1E-05       ,	num_attention_heads=4       ,	num_hidden_layers=5       ,	pad_token_id=1       ,	vocab_size=1_0_0_0       ,	hidden_act='''gelu'''       ,	projection_dim=5_1_2       ,	)
         _UpperCAmelCase							:  List[str]											=    CLIPTextModel(A		)
         _UpperCAmelCase							:  str											=    CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)

         _UpperCAmelCase							:  Union[str, Any]											=    {
             '''unet''': unet,
             '''scheduler''': scheduler,
             '''vae''': vae,
             '''text_encoder''': text_encoder,
             '''tokenizer''': tokenizer,
             '''safety_checker''': None,
             '''feature_extractor''': None,
         }

         return components





     def  __lowerCAmelCase	(					self       ,	A       ,	A=0		)							->					List[Any]:
         if str(A		).startswith('''mps'''		):
             _UpperCAmelCase							:  Optional[int]											=    torch.manual_seed(A		)
         else:
             _UpperCAmelCase							:  Union[str, Any]											=    torch.Generator(device=A		).manual_seed(A		)
         _UpperCAmelCase							:  List[str]											=    {
             '''prompt''': '''a cat and a frog''',
             '''token_indices''': [2, 5],
             '''generator''': generator,
             '''num_inference_steps''': 1,
             '''guidance_scale''': 6.0,
             '''output_type''': '''numpy''',
             '''max_iter_to_alter''': 2,
             '''thresholds''': {0: 0.7},
         }
         return inputs





     def  __lowerCAmelCase	(					self		)							->					int:
         _UpperCAmelCase							:  List[str]											=    '''cpu'''

         _UpperCAmelCase							:  Tuple											=    self.get_dummy_components()
         _UpperCAmelCase							:  int											=    self.pipeline_class(**A		)
         pipe.to(A		)
         pipe.set_progress_bar_config(disable=A		)

         _UpperCAmelCase							:  Dict											=    self.get_dummy_inputs(A		)
         _UpperCAmelCase							:  Union[str, Any]											=    pipe(**A		).images
         _UpperCAmelCase							:  Tuple											=    image[0, -3:, -3:, -1]

         self.assertEqual(image.shape       ,	(1, 6_4, 6_4, 3)		)
         _UpperCAmelCase							:  int											=    np.array(
             [0.63_905_364, 0.62_897_307, 0.48_599_017, 0.5_133_624, 0.5_550_048, 0.45_769_516, 0.50_326_973, 0.5_023_139, 0.45_384_496]		)
         _UpperCAmelCase							:  Tuple											=    np.abs(image_slice.flatten() - expected_slice		).max()
         self.assertLessEqual(A       ,	1E-3		)





     def  __lowerCAmelCase	(					self		)							->					Dict:
         super().test_cpu_offload_forward_pass(expected_max_diff=5E-4		)





     def  __lowerCAmelCase	(					self		)							->					List[str]:
         # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
         self._test_inference_batch_consistent(batch_sizes=[1, 2]		)





     def  __lowerCAmelCase	(					self		)							->					Union[str, Any]:
         self._test_inference_batch_single_identical(batch_size=2       ,	expected_max_diff=7E-4		)





     def  __lowerCAmelCase	(					self		)							->					List[str]:
         super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3		)





     def  __lowerCAmelCase	(					self		)							->					List[str]:
         super().test_pt_np_pil_outputs_equivalent(expected_max_diff=5E-4		)





     def  __lowerCAmelCase	(					self		)							->					str:
         super().test_save_load_local(expected_max_difference=5E-4		)





     def  __lowerCAmelCase	(					self		)							->					Optional[int]:
         super().test_save_load_optional_components(expected_max_difference=4E-4		)



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





     '''simple docstring'''



     @classmethod
     def  __lowerCAmelCase	(					cls		)							->					Union[str, Any]:
         super().setUpClass()
         torch.use_deterministic_algorithms(A		)





     @classmethod
     def  __lowerCAmelCase	(					cls		)							->					Optional[int]:
         super().tearDownClass()
         torch.use_deterministic_algorithms(A		)





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





     def  __lowerCAmelCase	(					self		)							->					str:
         _UpperCAmelCase							:  Any											=    torch.manual_seed(5_1		)

         _UpperCAmelCase							:  Optional[Any]											=    StableDiffusionAttendAndExcitePipeline.from_pretrained(
             '''CompVis/stable-diffusion-v1-4'''       ,	safety_checker=A       ,	torch_dtype=torch.floataa		)
         pipe.to('''cuda'''		)

         _UpperCAmelCase							:  Optional[int]											=    '''a painting of an elephant with glasses'''
         _UpperCAmelCase							:  int											=    [5, 7]

         _UpperCAmelCase							:  Dict											=    pipe(
             prompt=A       ,	token_indices=A       ,	guidance_scale=7.5       ,	generator=A       ,	num_inference_steps=5       ,	max_iter_to_alter=5       ,	output_type='''numpy'''       ,	).images[0]

         _UpperCAmelCase							:  List[Any]											=    load_numpy(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/attend-and-excite/elephant_glasses.npy'''		)
         assert np.abs((expected_image - image).max()		) < 5E-1


