



import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizer

from diffusers import (
    AltDiffusionImgaImgPipeline,
    AutoencoderKL,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()
class    __UpperCamelCase			(       unittest.TestCase				):

 """simple docstring"""



 def 						_UpperCAmelCase       (	self      ) ->			List[str]:
  # clean up the VRAM after each test
  super().tearDown()
  gc.collect()
  torch.cuda.empty_cache()

 @property
 def 						_UpperCAmelCase       (	self      ) ->			Dict:
  a__												= 1
  a__												= 3
  a__												= (3_2, 3_2)

  a__												= floats_tensor((batch_size, num_channels) + sizes		,				rng=random.Random(0      )      ).to(SCREAMING_SNAKE_CASE      )
  return image

 @property
 def 						_UpperCAmelCase       (	self      ) ->			Union[str, Any]:
  torch.manual_seed(0      )
  a__												= UNetaDConditionModel(
      block_out_channels=(3_2, 6_4)		,				layers_per_block=2		,				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		,				)
  return model

 @property
 def 						_UpperCAmelCase       (	self      ) ->			Dict:
  torch.manual_seed(0      )
  a__												= 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		,				)
  return model

 @property
 def 						_UpperCAmelCase       (	self      ) ->			int:
  torch.manual_seed(0      )
  a__												= RobertaSeriesConfig(
      hidden_size=3_2		,				project_dim=3_2		,				intermediate_size=3_7		,				layer_norm_eps=1e-0_5		,				num_attention_heads=4		,				num_hidden_layers=5		,				pad_token_id=1		,				vocab_size=5_0_0_6		,				)
  return RobertaSeriesModelWithTransformation(SCREAMING_SNAKE_CASE      )

 @property
 def 						_UpperCAmelCase       (	self      ) ->			int:
  def extract(*SCREAMING_SNAKE_CASE		,				**SCREAMING_SNAKE_CASE      ):
   class    __UpperCamelCase			:

    """simple docstring"""



    def __init__(	self      ) ->			Tuple:
     a__												= torch.ones([0]      )



    def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE      ) ->			Tuple:
     self.pixel_values.to(SCREAMING_SNAKE_CASE      )
     return self





   return Out()

  return extract

 def 						_UpperCAmelCase       (	self      ) ->			Union[str, Any]:
  a__												= '''cpu'''  # ensure determinism for the device-dependent torch.Generator
  a__												= self.dummy_cond_unet
  a__												= PNDMScheduler(skip_prk_steps=SCREAMING_SNAKE_CASE      )
  a__												= self.dummy_vae
  a__												= self.dummy_text_encoder
  a__												= XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''      )
  a__												= 7_7

  a__												= self.dummy_image.to(SCREAMING_SNAKE_CASE      )
  a__												= init_image / 2 + 0.5

  # make sure here that pndm scheduler skips prk
  a__												= AltDiffusionImgaImgPipeline(
      unet=SCREAMING_SNAKE_CASE		,				scheduler=SCREAMING_SNAKE_CASE		,				vae=SCREAMING_SNAKE_CASE		,				text_encoder=SCREAMING_SNAKE_CASE		,				tokenizer=SCREAMING_SNAKE_CASE		,				safety_checker=SCREAMING_SNAKE_CASE		,				feature_extractor=self.dummy_extractor		,				)
  a__												= VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor		,				do_normalize=SCREAMING_SNAKE_CASE      )
  a__												= alt_pipe.to(SCREAMING_SNAKE_CASE      )
  alt_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE      )

  a__												= '''A painting of a squirrel eating a burger'''
  a__												= torch.Generator(device=SCREAMING_SNAKE_CASE      ).manual_seed(0      )
  a__												= alt_pipe(
      [prompt]		,				generator=SCREAMING_SNAKE_CASE		,				guidance_scale=6.0		,				num_inference_steps=2		,				output_type='''np'''		,				image=SCREAMING_SNAKE_CASE		,				)

  a__												= output.images

  a__												= torch.Generator(device=SCREAMING_SNAKE_CASE      ).manual_seed(0      )
  a__												= alt_pipe(
      [prompt]		,				generator=SCREAMING_SNAKE_CASE		,				guidance_scale=6.0		,				num_inference_steps=2		,				output_type='''np'''		,				image=SCREAMING_SNAKE_CASE		,				return_dict=SCREAMING_SNAKE_CASE		,				)[0]

  a__												= image[0, -3:, -3:, -1]
  a__												= image_from_tuple[0, -3:, -3:, -1]

  assert image.shape == (1, 3_2, 3_2, 3)
  a__												= np.array([0.44_27, 0.37_31, 0.42_49, 0.49_41, 0.45_46, 0.41_48, 0.41_93, 0.46_66, 0.44_99]      )

  assert np.abs(image_slice.flatten() - expected_slice      ).max() < 5e-3
  assert np.abs(image_from_tuple_slice.flatten() - expected_slice      ).max() < 5e-3

 @unittest.skipIf(torch_device != '''cuda'''		,				'''This test requires a GPU'''      )
 def 						_UpperCAmelCase       (	self      ) ->			Union[str, Any]:
  a__												= self.dummy_cond_unet
  a__												= PNDMScheduler(skip_prk_steps=SCREAMING_SNAKE_CASE      )
  a__												= self.dummy_vae
  a__												= self.dummy_text_encoder
  a__												= XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''      )
  a__												= 7_7

  a__												= self.dummy_image.to(SCREAMING_SNAKE_CASE      )

  # put models in fp16
  a__												= unet.half()
  a__												= vae.half()
  a__												= bert.half()

  # make sure here that pndm scheduler skips prk
  a__												= AltDiffusionImgaImgPipeline(
      unet=SCREAMING_SNAKE_CASE		,				scheduler=SCREAMING_SNAKE_CASE		,				vae=SCREAMING_SNAKE_CASE		,				text_encoder=SCREAMING_SNAKE_CASE		,				tokenizer=SCREAMING_SNAKE_CASE		,				safety_checker=SCREAMING_SNAKE_CASE		,				feature_extractor=self.dummy_extractor		,				)
  a__												= VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor		,				do_normalize=SCREAMING_SNAKE_CASE      )
  a__												= alt_pipe.to(SCREAMING_SNAKE_CASE      )
  alt_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE      )

  a__												= '''A painting of a squirrel eating a burger'''
  a__												= torch.manual_seed(0      )
  a__												= alt_pipe(
      [prompt]		,				generator=SCREAMING_SNAKE_CASE		,				num_inference_steps=2		,				output_type='''np'''		,				image=SCREAMING_SNAKE_CASE		,				).images

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



 @unittest.skipIf(torch_device != '''cuda'''		,				'''This test requires a GPU'''      )
 def 						_UpperCAmelCase       (	self      ) ->			Tuple:
  a__												= load_image(
      '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
      '''/img2img/sketch-mountains-input.jpg'''      )
  # resize to resolution that is divisible by 8 but not 16 or 32
  a__												= init_image.resize((7_6_0, 5_0_4)      )

  a__												= '''BAAI/AltDiffusion'''
  a__												= AltDiffusionImgaImgPipeline.from_pretrained(
      SCREAMING_SNAKE_CASE		,				safety_checker=SCREAMING_SNAKE_CASE		,				)
  pipe.to(SCREAMING_SNAKE_CASE      )
  pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE      )
  pipe.enable_attention_slicing()

  a__												= '''A fantasy landscape, trending on artstation'''

  a__												= torch.manual_seed(0      )
  a__												= pipe(
      prompt=SCREAMING_SNAKE_CASE		,				image=SCREAMING_SNAKE_CASE		,				strength=0.75		,				guidance_scale=7.5		,				generator=SCREAMING_SNAKE_CASE		,				output_type='''np'''		,				)
  a__												= output.images[0]

  a__												= image[2_5_5:2_5_8, 3_8_3:3_8_6, -1]

  assert image.shape == (5_0_4, 7_6_0, 3)
  a__												= np.array([0.93_58, 0.93_97, 0.95_99, 0.99_01, 1.00_00, 1.00_00, 0.98_82, 1.00_00, 1.00_00]      )

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





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

 """simple docstring"""



 def 						_UpperCAmelCase       (	self      ) ->			Optional[Any]:
  # clean up the VRAM after each test
  super().tearDown()
  gc.collect()
  torch.cuda.empty_cache()



 def 						_UpperCAmelCase       (	self      ) ->			Optional[int]:
  a__												= load_image(
      '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
      '''/img2img/sketch-mountains-input.jpg'''      )
  a__												= init_image.resize((7_6_8, 5_1_2)      )
  a__												= load_numpy(
      '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy'''      )

  a__												= '''BAAI/AltDiffusion'''
  a__												= AltDiffusionImgaImgPipeline.from_pretrained(
      SCREAMING_SNAKE_CASE		,				safety_checker=SCREAMING_SNAKE_CASE		,				)
  pipe.to(SCREAMING_SNAKE_CASE      )
  pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE      )
  pipe.enable_attention_slicing()

  a__												= '''A fantasy landscape, trending on artstation'''

  a__												= torch.manual_seed(0      )
  a__												= pipe(
      prompt=SCREAMING_SNAKE_CASE		,				image=SCREAMING_SNAKE_CASE		,				strength=0.75		,				guidance_scale=7.5		,				generator=SCREAMING_SNAKE_CASE		,				output_type='''np'''		,				)
  a__												= output.images[0]

  assert image.shape == (5_1_2, 7_6_8, 3)
  # img2img is flaky across GPUs even in fp32, so using MAE here
  assert np.abs(expected_image - image      ).max() < 1e-2










import logging
import os

from .state import PartialState
class    __UpperCamelCase			(       logging.LoggerAdapter				):

 """simple docstring"""



 @staticmethod
 def 						_UpperCAmelCase       (	SCREAMING_SNAKE_CASE      ) ->			Optional[Any]:
  a__												= PartialState()
  return not main_process_only or (main_process_only and state.is_main_process)



 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				*SCREAMING_SNAKE_CASE		,				**SCREAMING_SNAKE_CASE      ) ->			Optional[int]:
  if PartialState._shared_state == {}:
   raise RuntimeError(
       '''You must initialize the accelerate state by calling either `PartialState()` or `Accelerator()` before using the logging utility.'''      )
  a__												= kwargs.pop('''main_process_only'''		,				SCREAMING_SNAKE_CASE      )
  a__												= kwargs.pop('''in_order'''		,				SCREAMING_SNAKE_CASE      )

  if self.isEnabledFor(SCREAMING_SNAKE_CASE      ):
   if self._should_log(SCREAMING_SNAKE_CASE      ):
    a__							,			a__												= self.process(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      )
    self.logger.log(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				*SCREAMING_SNAKE_CASE		,				**SCREAMING_SNAKE_CASE      )

   elif in_order:
    a__												= PartialState()
    for i in range(state.num_processes      ):
     if i == state.process_index:
      a__							,			a__												= self.process(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      )
      self.logger.log(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				*SCREAMING_SNAKE_CASE		,				**SCREAMING_SNAKE_CASE      )
     state.wait_for_everyone()






def     __a					(		__UpperCAmelCase							,  __UpperCAmelCase = None       ):
 if log_level is None:
  a__												= os.environ.get('''ACCELERATE_LOG_LEVEL'''							,  __UpperCAmelCase       )
 a__												= logging.getLogger(__UpperCAmelCase       )
 if log_level is not None:
  logger.setLevel(log_level.upper()       )
  logger.root.setLevel(log_level.upper()       )
 return MultiProcessAdapter(__UpperCAmelCase							,  {}       )





