'''simple docstring'''






import random
import unittest

import numpy as np
import torch

from diffusers import DDIMScheduler, LDMSuperResolutionPipeline, UNetaDModel, VQModel
from diffusers.utils import PIL_INTERPOLATION, floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch


enable_full_determinism()







class lowercase_					(		unittest.TestCase ):







   """simple docstring"""

   @property
   def   __UpperCAmelCase      (			self			: List[str]    )     -> Optional[int]:
        _A								=     1
        _A								=     3
        _A								=     (32, 32)

        _A								=     floats_tensor((batch_size, num_channels) + sizes,   rng=random.Random(0    )    ).to(UpperCamelCase__    )
        return image
   @property
   def   __UpperCAmelCase      (			self			: Any    )     -> Union[str, Any]:
        torch.manual_seed(0    )
        _A								=     UNetaDModel(
            block_out_channels=(32, 64),   layers_per_block=2,   sample_size=32,   in_channels=6,   out_channels=3,   down_block_types=('DownBlock2D', 'AttnDownBlock2D'),   up_block_types=('AttnUpBlock2D', 'UpBlock2D'),   )
        return model
   @property
   def   __UpperCAmelCase      (			self			: Any    )     -> Dict:
        torch.manual_seed(0    )
        _A								=     VQModel(
            block_out_channels=[32, 64],   in_channels=3,   out_channels=3,   down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D'],   up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D'],   latent_channels=3,   )
        return model
   def   __UpperCAmelCase      (			self			: List[Any]    )     -> Optional[Any]:
        _A								=     'cpu'
        _A								=     self.dummy_uncond_unet
        _A								=     DDIMScheduler()
        _A								=     self.dummy_vq_model

        _A								=     LDMSuperResolutionPipeline(unet=UpperCamelCase__,   vqvae=UpperCamelCase__,   scheduler=UpperCamelCase__    )
        ldm.to(UpperCamelCase__    )
        ldm.set_progress_bar_config(disable=UpperCamelCase__    )

        _A								=     self.dummy_image.to(UpperCamelCase__    )

        _A								=     torch.Generator(device=UpperCamelCase__    ).manual_seed(0    )
        _A								=     ldm(image=UpperCamelCase__,   generator=UpperCamelCase__,   num_inference_steps=2,   output_type='numpy'    ).images

        _A								=     image[0, -3:, -3:, -1]

        assert image.shape == (1, 64, 64, 3)
        _A								=     np.array([0.8_678, 0.8_245, 0.6_381, 0.6_830, 0.4_385, 0.5_599, 0.4_641, 0.6_201, 0.5_150]    )

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

   @unittest.skipIf(torch_device != 'cuda',   'This test requires a GPU'    )
   def   __UpperCAmelCase      (			self			: Union[str, Any]    )     -> Union[str, Any]:
        _A								=     self.dummy_uncond_unet
        _A								=     DDIMScheduler()
        _A								=     self.dummy_vq_model

        # put models in fp16
        _A								=     unet.half()
        _A								=     vqvae.half()

        _A								=     LDMSuperResolutionPipeline(unet=UpperCamelCase__,   vqvae=UpperCamelCase__,   scheduler=UpperCamelCase__    )
        ldm.to(UpperCamelCase__    )
        ldm.set_progress_bar_config(disable=UpperCamelCase__    )

        _A								=     self.dummy_image.to(UpperCamelCase__    )

        _A								=     ldm(UpperCamelCase__,   num_inference_steps=2,   output_type='numpy'    ).images

        assert image.shape == (1, 64, 64, 3)







@slow
@require_torch
class lowercase_					(		unittest.TestCase ):







   """simple docstring"""

   def   __UpperCAmelCase      (			self			: str    )     -> Dict:
        _A								=     load_image(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
            '/vq_diffusion/teddy_bear_pool.png'    )
        _A								=     init_image.resize((64, 64),   resample=PIL_INTERPOLATION['lanczos']    )

        _A								=     LDMSuperResolutionPipeline.from_pretrained('duongna/ldm-super-resolution',   device_map='auto'    )
        ldm.set_progress_bar_config(disable=UpperCamelCase__    )

        _A								=     torch.manual_seed(0    )
        _A								=     ldm(image=UpperCamelCase__,   generator=UpperCamelCase__,   num_inference_steps=20,   output_type='numpy'    ).images

        _A								=     image[0, -3:, -3:, -1]

        assert image.shape == (1, 2_56, 2_56, 3)
        _A								=     np.array([0.7_644, 0.7_679, 0.7_642, 0.7_633, 0.7_666, 0.7_560, 0.7_425, 0.7_257, 0.6_907]    )

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




'''simple docstring'''






import torch

from accelerate import Accelerator, DistributedType







class lowercase_					:







   """simple docstring"""

   def __enter__(			self			: Optional[int]    )     -> Union[str, Any]:
        if self.enabled:
             _A								=     self.model.no_sync()
             self.model_sync_obj.__enter__()

        return self
   def __exit__(			self			: Dict,   UpperCamelCase__			: Any,   UpperCamelCase__			: int,   UpperCamelCase__			: Any    )     -> str:
        if self.enabled:
             # Average all models on exit
             self._sync_and_avg_model_params()
             self.model_sync_obj.__exit__(UpperCamelCase__,   UpperCamelCase__,   UpperCamelCase__    )
   def __init__(			self			: Optional[int],   UpperCamelCase__			: Accelerator,   UpperCamelCase__			: torch.nn.Module,   UpperCamelCase__			: int,   UpperCamelCase__			: bool = True    )     -> str:
        if accelerator.distributed_type not in [
            DistributedType.NO,
            DistributedType.MULTI_CPU,
            DistributedType.MULTI_GPU,
        ]:
             raise NotImplementedError('LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)'    )
        _A								=     enabled and accelerator.distributed_type != DistributedType.NO
        _A								=     0
        if self.enabled:
             _A								=     accelerator
             _A								=     model
             _A								=     local_sgd_steps
   def   __UpperCAmelCase      (			self			: Any    )     -> int:
        self.num_steps += 1
        if not self.enabled:
             return

        if self.num_steps % self.local_sgd_steps == 0:
             self._sync_and_avg_model_params()

   def   __UpperCAmelCase      (			self			: Dict    )     -> int:
        self.accelerator.wait_for_everyone()
        with self.accelerator.autocast():
             for param in self.model.parameters():
                  _A								=     self.accelerator.reduce(param.data,   reduction='mean'    )



