


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  snake_case_     (unittest.TestCase   ):
					def 				lowerCamelCase__(   self     :Optional[int]   )      ->					Tuple:
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()
					@property
					def 				lowerCamelCase__(   self     :Union[str, Any]   )      ->					int:
												a__    =  1
												a__    =  3
												a__    =  (32, 32)

												a__    =  floats_tensor((batch_size, num_channels) + sizes    ,rng=random.Random(0   )   ).to(lowerCAmelCase_   )
												return image
					@property
					def 				lowerCamelCase__(   self     :Union[str, Any]   )      ->					str:
												torch.manual_seed(0   )
												a__    =  UNetaDConditionModel(
												    block_out_channels=(32, 64)    ,layers_per_block=2    ,sample_size=32    ,in_channels=4    ,out_channels=4    ,down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')    ,up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')    ,cross_attention_dim=32    ,)
												return model
					@property
					def 				lowerCamelCase__(   self     :Optional[int]   )      ->					str:
												torch.manual_seed(0   )
												a__    =  AutoencoderKL(
												    block_out_channels=[32, 64]    ,in_channels=3    ,out_channels=3    ,down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']    ,up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']    ,latent_channels=4    ,)
												return model
					@property
					def 				lowerCamelCase__(   self     :Union[str, Any]   )      ->					int:
												torch.manual_seed(0   )
												a__    =  RobertaSeriesConfig(
												    hidden_size=32    ,project_dim=32    ,intermediate_size=37    ,layer_norm_eps=1E-05    ,num_attention_heads=4    ,num_hidden_layers=5    ,pad_token_id=1    ,vocab_size=50_06    ,)
												return RobertaSeriesModelWithTransformation(lowerCAmelCase_   )
					@property
					def 				lowerCamelCase__(   self     :Optional[Any]   )      ->					Dict:







												def extract(*__snake_case     :int    ,**__snake_case     :int   ):







																			class  snake_case_     :
																								def __init__(   self     :int   )      ->					Optional[int]:
																															a__    =  torch.ones([0]   )



																								def 				lowerCamelCase__(   self     :List[Any]    ,__snake_case     :Optional[int]   )      ->					Union[str, Any]:
																															self.pixel_values.to(lowerCAmelCase_   )
																															return self





																			return Out()

												return extract
					def 				lowerCamelCase__(   self     :List[Any]   )      ->					List[Any]:
												a__    =  'cpu'  # ensure determinism for the device-dependent torch.Generator
												a__    =  self.dummy_cond_unet
												a__    =  PNDMScheduler(skip_prk_steps=lowerCAmelCase_   )
												a__    =  self.dummy_vae
												a__    =  self.dummy_text_encoder
												a__    =  XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'   )
												a__    =  77

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

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

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

												a__    =  output.images

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

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

												assert image.shape == (1, 32, 32, 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 				lowerCamelCase__(   self     :Optional[int]   )      ->					Tuple:
												a__    =  self.dummy_cond_unet
												a__    =  PNDMScheduler(skip_prk_steps=lowerCAmelCase_   )
												a__    =  self.dummy_vae
												a__    =  self.dummy_text_encoder
												a__    =  XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'   )
												a__    =  77

												a__    =  self.dummy_image.to(lowerCAmelCase_   )

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

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

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

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



					@unittest.skipIf(torch_device != 'cuda'    ,'This test requires a GPU'   )
					def 				lowerCamelCase__(   self     :Union[str, Any]   )      ->					List[Any]:
												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_60, 5_04)   )

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

												a__    =  'A fantasy landscape, trending on artstation'

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

												a__    =  image[2_55:2_58, 3_83:3_86, -1]

												assert image.shape == (5_04, 7_60, 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  snake_case_     (unittest.TestCase   ):
					def 				lowerCamelCase__(   self     :Tuple   )      ->					Dict:
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()



					def 				lowerCamelCase__(   self     :str   )      ->					Tuple:
												a__    =  load_image(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
												    '/img2img/sketch-mountains-input.jpg'   )
												a__    =  init_image.resize((7_68, 5_12)   )
												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(
												    lowerCAmelCase_    ,safety_checker=lowerCAmelCase_    ,)
												pipe.to(lowerCAmelCase_   )
												pipe.set_progress_bar_config(disable=lowerCAmelCase_   )
												pipe.enable_attention_slicing()

												a__    =  'A fantasy landscape, trending on artstation'

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

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








from typing import List, Optional, Union

import numpy as np
import PIL.Image

from ...image_processing_utils import BaseImageProcessor, BatchFeature
from ...image_transforms import rescale, resize, to_channel_dimension_format
from ...image_utils import (
    ChannelDimension,
    PILImageResampling,
    get_image_size,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging


snake_case							: Dict      =   logging.get_logger(__name__)







class  snake_case_     (lowerCamelCase_   ):
					UpperCAmelCase__      :					Dict					   =    ['''pixel_values''']
					def __init__(   self     :Optional[Any]    ,__snake_case     :bool = True    ,__snake_case     :int = 32    ,__snake_case     :Union[str, Any]=PILImageResampling.BILINEAR    ,__snake_case     :bool = True    ,**__snake_case     :Tuple    ,)      ->					None:
												a__    =  do_resize
												a__    =  do_rescale
												a__    =  size_divisor
												a__    =  resample
												super().__init__(**__snake_case   )
					def 				lowerCamelCase__(   self     :Union[str, Any]    ,__snake_case     :np.ndarray    ,__snake_case     :int    ,__snake_case     :Tuple    ,__snake_case     :Optional[ChannelDimension] = None    ,**__snake_case     :List[Any]   )      ->					np.ndarray:
												a__   ,      a__    =  get_image_size(__snake_case   )
												# Rounds the height and width down to the closest multiple of size_divisor
												a__    =  height // size_divisor * size_divisor
												a__    =  width // size_divisor * size_divisor
												a__    =  resize(__snake_case    ,(new_h, new_w)    ,resample=__snake_case    ,data_format=__snake_case    ,**__snake_case   )
												return image
					def 				lowerCamelCase__(   self     :List[str]    ,__snake_case     :np.ndarray    ,__snake_case     :float    ,__snake_case     :Optional[ChannelDimension] = None    ,**__snake_case     :str   )      ->					np.ndarray:
												return rescale(image=__snake_case    ,scale=__snake_case    ,data_format=__snake_case    ,**__snake_case   )



					def 				lowerCamelCase__(   self     :Tuple    ,__snake_case     :Union["PIL.Image.Image", TensorType, List["PIL.Image.Image"], List[TensorType]]    ,__snake_case     :Optional[bool] = None    ,__snake_case     :Optional[int] = None    ,__snake_case     :Union[str, Any]=None    ,__snake_case     :Optional[bool] = None    ,__snake_case     :Optional[Union[TensorType, str]] = None    ,__snake_case     :ChannelDimension = ChannelDimension.FIRST    ,**__snake_case     :List[Any]    ,)      ->					BatchFeature:
												a__    =  do_resize if do_resize is not None else self.do_resize
												a__    =  do_rescale if do_rescale is not None else self.do_rescale
												a__    =  size_divisor if size_divisor is not None else self.size_divisor
												a__    =  resample if resample is not None else self.resample

												if do_resize and size_divisor is None:
																			raise ValueError('size_divisor is required for resizing'   )

												a__    =  make_list_of_images(__snake_case   )

												if not valid_images(__snake_case   ):
																			raise ValueError('Invalid image(s)'   )

												# All transformations expect numpy arrays.
												a__    =  [to_numpy_array(__snake_case   ) for img in images]

												if do_resize:
																			a__    =  [self.resize(__snake_case    ,size_divisor=__snake_case    ,resample=__snake_case   ) for image in images]

												if do_rescale:
																			a__    =  [self.rescale(__snake_case    ,scale=1 / 2_55   ) for image in images]

												a__    =  [to_channel_dimension_format(__snake_case    ,__snake_case   ) for image in images]

												a__    =  {'pixel_values': images}
												return BatchFeature(data=__snake_case    ,tensor_type=__snake_case   )




