




import unittest

import numpy as np
import torch

from diffusers import PNDMPipeline, PNDMScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()





class  __UpperCAmelCase					(	unittest.TestCase						):



   @property
   def 	__magic_name__   (							self							:					str						):
     torch.manual_seed(0						)
     UpperCAmelCase	:							List[Any]						=			UNetaDModel(
         block_out_channels=(3_2, 6_4),		layers_per_block=2,		sample_size=3_2,		in_channels=3,		out_channels=3,		down_block_types=('''DownBlock2D''', '''AttnDownBlock2D'''),		up_block_types=('''AttnUpBlock2D''', '''UpBlock2D'''),		)
     return model





   def 	__magic_name__   (							self							:					Tuple						):
     UpperCAmelCase	:							Any						=			self.dummy_uncond_unet
     UpperCAmelCase	:							Any						=			PNDMScheduler()

     UpperCAmelCase	:							List[str]						=			PNDMPipeline(unet=__A,		scheduler=__A						)
     pndm.to(__A						)
     pndm.set_progress_bar_config(disable=__A						)

     UpperCAmelCase	:							str						=			torch.manual_seed(0						)
     UpperCAmelCase	:							Optional[Any]						=			pndm(generator=__A,		num_inference_steps=2_0,		output_type='''numpy'''						).images

     UpperCAmelCase	:							Tuple						=			torch.manual_seed(0						)
     UpperCAmelCase	:							Optional[Any]						=			pndm(generator=__A,		num_inference_steps=2_0,		output_type='''numpy''',		return_dict=__A						)[0]

     UpperCAmelCase	:							Optional[int]						=			image[0, -3:, -3:, -1]
     UpperCAmelCase	:							Dict						=			image_from_tuple[0, -3:, -3:, -1]

     assert image.shape == (1, 3_2, 3_2, 3)
     UpperCAmelCase	:							Dict						=			np.array([1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0]						)

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






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



   def 	__magic_name__   (							self							:					str						):
     UpperCAmelCase	:							Optional[int]						=			'''google/ddpm-cifar10-32'''

     UpperCAmelCase	:							Optional[int]						=			UNetaDModel.from_pretrained(__A						)
     UpperCAmelCase	:							Union[str, Any]						=			PNDMScheduler()

     UpperCAmelCase	:							List[Any]						=			PNDMPipeline(unet=__A,		scheduler=__A						)
     pndm.to(__A						)
     pndm.set_progress_bar_config(disable=__A						)
     UpperCAmelCase	:							Dict						=			torch.manual_seed(0						)
     UpperCAmelCase	:							Dict						=			pndm(generator=__A,		output_type='''numpy'''						).images

     UpperCAmelCase	:							int						=			image[0, -3:, -3:, -1]

     assert image.shape == (1, 3_2, 3_2, 3)
     UpperCAmelCase	:							int						=			np.array([0.1_5_6_4, 0.1_4_6_4_5, 0.1_4_0_6, 0.1_4_7_1_5, 0.1_2_4_2_5, 0.1_4_0_4_5, 0.1_3_1_1_5, 0.1_2_1_7_5, 0.1_2_5]						)

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








import logging
import os

from .state import PartialState





class  __UpperCAmelCase					(	logging.LoggerAdapter						):



   @staticmethod
   def 	__magic_name__   (							__A							:					str						):
     UpperCAmelCase	:							Dict						=			PartialState()
     return not main_process_only or (main_process_only and state.is_main_process)





   def 	__magic_name__   (							self							:					Union[str, Any],		__A							:					Union[str, Any],		__A							:					Union[str, Any],		*__A							:					Optional[int],		**__A							:					Tuple						):
     if PartialState._shared_state == {}:
       raise RuntimeError(
           '''You must initialize the accelerate state by calling either `PartialState()` or `Accelerator()` before using the logging utility.'''						)
     UpperCAmelCase	:							List[str]						=			kwargs.pop('''main_process_only''',		__A						)
     UpperCAmelCase	:							int						=			kwargs.pop('''in_order''',		__A						)

     if self.isEnabledFor(__A						):
       if self._should_log(__A						):
         UpperCAmelCase			,						UpperCAmelCase	:							Dict						=			self.process(__A,		__A						)
         self.logger.log(__A,		__A,		*__A,		**__A						)

       elif in_order:
         UpperCAmelCase	:							Dict						=			PartialState()
         for i in range(state.num_processes						):
           if i == state.process_index:
             UpperCAmelCase			,						UpperCAmelCase	:							Optional[int]						=			self.process(__A,		__A						)
             self.logger.log(__A,		__A,		*__A,		**__A						)
           state.wait_for_everyone()






def       a__					(		UpperCAmelCase							:      str    ,			UpperCAmelCase							:      str = None      )  ->						Dict:
  if log_level is None:
    UpperCAmelCase	:							Union[str, Any]						=			os.environ.get('''ACCELERATE_LOG_LEVEL'''    ,			UpperCAmelCase      )
  UpperCAmelCase	:							Tuple						=			logging.getLogger(UpperCAmelCase      )
  if log_level is not None:
    logger.setLevel(log_level.upper()      )
    logger.root.setLevel(log_level.upper()      )
  return MultiProcessAdapter(UpperCAmelCase    ,			{}      )





