






'''simple docstring'''





import random
import unittest

import torch

from diffusers import IFSuperResolutionPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin
@skip_mps
class 					lowercase							(				_lowerCamelCase			,							_lowerCamelCase			,							unittest.TestCase   ):


    '''simple docstring'''

    lowerCAmelCase__   					=				IFSuperResolutionPipeline
    lowerCAmelCase__   					=				TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"width", "height"}
    lowerCAmelCase__   					=				TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
    lowerCAmelCase__   					=				PipelineTesterMixin.required_optional_params - {"latents"}

    def  a__					(       self  :  Optional[Any]  )				->				Tuple:

     '''simple docstring'''




     return self._get_superresolution_dummy_components()

    def  a__					(       self  :  Optional[int]					,  __lowerCamelCase  :  List[Any]					,  __lowerCamelCase  :  int=0  )				->				List[Any]:

     '''simple docstring'''




     if str(__lowerCamelCase  ).startswith("mps"  ):
      lowerCamelCase__	=  torch.manual_seed(__lowerCamelCase  )
     else:
      lowerCamelCase__	=  torch.Generator(device=__lowerCamelCase  ).manual_seed(__lowerCamelCase  )

     lowerCamelCase__	=  floats_tensor((1, 3, 32, 32)					,  rng=random.Random(__lowerCamelCase  )  ).to(__lowerCamelCase  )

     lowerCamelCase__	=  {
         "prompt": "A painting of a squirrel eating a burger",
         "image": image,
         "generator": generator,
         "num_inference_steps": 2,
         "output_type": "numpy",
     }

     return inputs

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

     '''simple docstring'''




     self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3  )

    def  a__					(       self  :  Optional[Any]  )				->				Tuple:

     '''simple docstring'''




     self._test_save_load_optional_components()

    @unittest.skipIf(torch_device != "cuda"					,  reason="float16 requires CUDA"  )
    def  a__					(       self  :  Optional[Any]  )				->				Optional[Any]:

     '''simple docstring'''




     super().test_save_load_floataa(expected_max_diff=1E-1  )

    def  a__					(       self  :  Tuple  )				->				Any:

     '''simple docstring'''




     self._test_attention_slicing_forward_pass(expected_max_diff=1E-2  )

    def  a__					(       self  :  Optional[int]  )				->				List[str]:

     '''simple docstring'''




     self._test_save_load_local()




    def  a__					(       self  :  Dict  )				->				List[Any]:

     '''simple docstring'''




     self._test_inference_batch_single_identical(
         expected_max_diff=1E-2					,  )









'''simple docstring'''





import os
import sys
from pathlib import Path

import pytest

from datasets import Dataset, IterableDataset
from datasets.distributed import split_dataset_by_node

from .utils import execute_subprocess_async, get_torch_dist_unique_port, require_torch






def  lowerCamelCase_  (     ):
 lowerCamelCase__	=  Dataset.from_dict({"i": range(17)})
 lowerCamelCase__	=  len(lowercase__)
 lowerCamelCase__	=  3
 lowerCamelCase__	=  [
     split_dataset_by_node(lowercase__					,  rank=lowercase__					,  world_size=lowercase__) for rank in range(lowercase__)
 ]
 assert sum(len(lowercase__) for ds in datasets_per_rank) == full_size
 assert len({tuple(x.values()) for ds in datasets_per_rank for x in ds}) == full_size






def  lowerCamelCase_  (     ):
 def gen():
  return ({"i": i} for i in range(17))

 lowerCamelCase__	=  3
 lowerCamelCase__	=  IterableDataset.from_generator(lowercase__)
 lowerCamelCase__	=  len(list(lowercase__))
 lowerCamelCase__	=  [
     split_dataset_by_node(lowercase__					,  rank=lowercase__					,  world_size=lowercase__) for rank in range(lowercase__)
 ]
 assert sum(len(list(lowercase__)) for ds in datasets_per_rank) == full_size
 assert len({tuple(x.values()) for ds in datasets_per_rank for x in ds}) == full_size






@pytest.mark.parametrize("shards_per_node"					,  [1, 2, 3])
def  lowerCamelCase_  (     lowercase__):
 def gen(lowercase__):
  for shard in shards:
   yield from ({"i": i, "shard": shard} for i in range(17))

 lowerCamelCase__	=  3
 lowerCamelCase__	=  shards_per_node * world_size
 lowerCamelCase__	=  {"shards": [F'''shard_{shard_idx}.txt''' for shard_idx in range(lowercase__)]}
 lowerCamelCase__	=  IterableDataset.from_generator(lowercase__					,  gen_kwargs=lowercase__)
 lowerCamelCase__	=  len(list(lowercase__))
 assert full_ds.n_shards == world_size * shards_per_node
 lowerCamelCase__	=  [
     split_dataset_by_node(lowercase__					,  rank=lowercase__					,  world_size=lowercase__) for rank in range(lowercase__)
 ]
 assert [ds.n_shards for ds in datasets_per_rank] == [shards_per_node] * world_size
 assert sum(len(list(lowercase__)) for ds in datasets_per_rank) == full_size
 assert len({tuple(x.values()) for ds in datasets_per_rank for x in ds}) == full_size






def  lowerCamelCase_  (     ):
 def gen():
  return ({"i": i} for i in range(17))

 lowerCamelCase__	=  2
 lowerCamelCase__	=  IterableDataset.from_generator(lowercase__)
 lowerCamelCase__	=  len(list(lowercase__))

 lowerCamelCase__	=  split_dataset_by_node(lowercase__					,  rank=0					,  world_size=lowercase__).shuffle(seed=42)
 assert len(list(lowercase__)) == 1 + full_size // world_size
 with pytest.raises(lowercase__):
  split_dataset_by_node(lowercase__					,  rank=0					,  world_size=lowercase__).shuffle()

 lowerCamelCase__	=  split_dataset_by_node(full_ds.shuffle(seed=42)					,  rank=0					,  world_size=lowercase__)
 assert len(list(lowercase__)) == 1 + full_size // world_size
 with pytest.raises(lowercase__):
  split_dataset_by_node(full_ds.shuffle()					,  rank=0					,  world_size=lowercase__)






@pytest.mark.parametrize("streaming"					,  [False, True])
@require_torch
@pytest.mark.skipif(os.name == "nt"					,  reason="execute_subprocess_async doesn't support windows")
@pytest.mark.integration
def  lowerCamelCase_  (     lowercase__):
 lowerCamelCase__	=  2
 lowerCamelCase__	=  get_torch_dist_unique_port()
 lowerCamelCase__	=  Path(lowercase__).resolve().parent / "distributed_scripts" / "launch_torch_distributed.py"
 lowerCamelCase__	=  F'''
        -m torch.distributed.launch
        --nproc_per_node={nproc_per_node}
        --master_port={master_port}
        {test_script}
    '''.split()
 lowerCamelCase__	=  F'''
        --streaming={streaming}
    '''.split()
 lowerCamelCase__	=  [sys.executable] + distributed_args + args
 execute_subprocess_async(lowercase__					,  env=os.environ.copy())






@pytest.mark.parametrize(
    "nproc_per_node, num_workers"					,  [
        (2, 2),  # each node has 2 shards and each worker has 1 shards
        (3, 2),  # each node uses all the shards but skips examples, and each worker has 2 shards
    ]					,  )
@require_torch
@pytest.mark.skipif(os.name == "nt"					,  reason="execute_subprocess_async doesn't support windows")
@pytest.mark.integration
def  lowerCamelCase_  (     lowercase__					,  lowercase__):
 lowerCamelCase__	=  True
 lowerCamelCase__	=  get_torch_dist_unique_port()
 lowerCamelCase__	=  Path(lowercase__).resolve().parent / "distributed_scripts" / "launch_torch_distributed.py"
 lowerCamelCase__	=  F'''
        -m torch.distributed.launch
        --nproc_per_node={nproc_per_node}
        --master_port={master_port}
        {test_script}
    '''.split()
 lowerCamelCase__	=  F'''
        --streaming={streaming}
        --num_workers={num_workers}
    '''.split()
 lowerCamelCase__	=  [sys.executable] + distributed_args + args
 execute_subprocess_async(lowercase__					,  env=os.environ.copy())

