




"""simple docstring"""






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   a__			(       unittest.TestCase  ):





						def       __UpperCamelCase     (							self      :	int)						->       Any:




											"""simple docstring"""







											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()





						@property
						def       __UpperCamelCase     (							self      :	List[Any])						->       str:




											"""simple docstring"""







											_lowerCAmelCase:Optional[int]									=  1
											_lowerCAmelCase:Union[str, Any]									=  3
											_lowerCAmelCase:Union[str, Any]									=  (32, 32)

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





						@property
						def       __UpperCamelCase     (							self      :	List[Any])						->       str:




											"""simple docstring"""







											torch.manual_seed(0)
											_lowerCAmelCase:int									=  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       __UpperCamelCase     (							self      :	Optional[int])						->       List[str]:




											"""simple docstring"""







											torch.manual_seed(0)
											_lowerCAmelCase:Optional[Any]									=  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       __UpperCamelCase     (							self      :	str)						->       Tuple:




											"""simple docstring"""







											torch.manual_seed(0)
											_lowerCAmelCase:Optional[Any]									=  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=5006     ,)
											return RobertaSeriesModelWithTransformation(a__)





						@property
						def       __UpperCamelCase     (							self      :	Optional[Any])						->       List[Any]:




											"""simple docstring"""







											def extract(*a__      :	List[str]     ,**a__      :	List[Any]):




																class   a__			:





																						def __init__(							self      :	List[Any])						->       Optional[Any]:




																											"""simple docstring"""







																											_lowerCAmelCase:Any									=  torch.ones([0])







																						def       __UpperCamelCase     (							self      :	List[Any]     ,a__      :	str)						->       List[Any]:




																											"""simple docstring"""







																											self.pixel_values.to(a__)
																											return self



																return Out()

											return extract





						def       __UpperCamelCase     (							self      :	Tuple)						->       str:




											"""simple docstring"""







											_lowerCAmelCase:Union[str, Any]									=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
											_lowerCAmelCase:Union[str, Any]									=  self.dummy_cond_unet
											_lowerCAmelCase:Any									=  PNDMScheduler(skip_prk_steps=a__)
											_lowerCAmelCase:Optional[int]									=  self.dummy_vae
											_lowerCAmelCase:List[str]									=  self.dummy_text_encoder
											_lowerCAmelCase:List[Any]									=  XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta''')
											_lowerCAmelCase:Tuple									=  77

											_lowerCAmelCase:List[Any]									=  self.dummy_image.to(a__)
											_lowerCAmelCase:Any									=  init_image / 2 + 0.5

											# make sure here that pndm scheduler skips prk
											_lowerCAmelCase:Tuple									=  AltDiffusionImgaImgPipeline(
											    unet=a__     ,scheduler=a__     ,vae=a__     ,text_encoder=a__     ,tokenizer=a__     ,safety_checker=a__     ,feature_extractor=self.dummy_extractor     ,)
											_lowerCAmelCase:str									=  VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor     ,do_normalize=a__)
											_lowerCAmelCase:Union[str, Any]									=  alt_pipe.to(a__)
											alt_pipe.set_progress_bar_config(disable=a__)

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

											_lowerCAmelCase:Dict									=  output.images

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

											_lowerCAmelCase:Any									=  image[0, -3:, -3:, -1]
											_lowerCAmelCase:Tuple									=  image_from_tuple[0, -3:, -3:, -1]

											assert image.shape == (1, 32, 32, 3)
											_lowerCAmelCase:Union[str, Any]									=  np.array([0.4427, 0.3731, 0.4249, 0.4941, 0.4546, 0.4148, 0.4193, 0.4666, 0.4499])

											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      :	Dict)						->       int:




											"""simple docstring"""







											_lowerCAmelCase:Dict									=  self.dummy_cond_unet
											_lowerCAmelCase:List[Any]									=  PNDMScheduler(skip_prk_steps=a__)
											_lowerCAmelCase:Dict									=  self.dummy_vae
											_lowerCAmelCase:Tuple									=  self.dummy_text_encoder
											_lowerCAmelCase:Union[str, Any]									=  XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta''')
											_lowerCAmelCase:List[Any]									=  77

											_lowerCAmelCase:Any									=  self.dummy_image.to(a__)

											# put models in fp16
											_lowerCAmelCase:List[Any]									=  unet.half()
											_lowerCAmelCase:List[str]									=  vae.half()
											_lowerCAmelCase:List[str]									=  bert.half()

											# make sure here that pndm scheduler skips prk
											_lowerCAmelCase:List[str]									=  AltDiffusionImgaImgPipeline(
											    unet=a__     ,scheduler=a__     ,vae=a__     ,text_encoder=a__     ,tokenizer=a__     ,safety_checker=a__     ,feature_extractor=self.dummy_extractor     ,)
											_lowerCAmelCase:Optional[Any]									=  VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor     ,do_normalize=a__)
											_lowerCAmelCase:Any									=  alt_pipe.to(a__)
											alt_pipe.set_progress_bar_config(disable=a__)

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

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







						@unittest.skipIf(torch_device != '''cuda'''     ,'''This test requires a GPU''')
						def       __UpperCamelCase     (							self      :	Optional[Any])						->       str:




											"""simple docstring"""







											_lowerCAmelCase:List[Any]									=  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
											_lowerCAmelCase:Dict									=  init_image.resize((760, 504))

											_lowerCAmelCase:List[Any]									=  '''BAAI/AltDiffusion'''
											_lowerCAmelCase:List[str]									=  AltDiffusionImgaImgPipeline.from_pretrained(
											    a__     ,safety_checker=a__     ,)
											pipe.to(a__)
											pipe.set_progress_bar_config(disable=a__)
											pipe.enable_attention_slicing()

											_lowerCAmelCase:Optional[int]									=  '''A fantasy landscape, trending on artstation'''

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

											_lowerCAmelCase:Union[str, Any]									=  image[255:258, 383:386, -1]

											assert image.shape == (504, 760, 3)
											_lowerCAmelCase:Tuple									=  np.array([0.9358, 0.9397, 0.9599, 0.9901, 1.0000, 1.0000, 0.9882, 1.0000, 1.0000])

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




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





						def       __UpperCamelCase     (							self      :	Any)						->       Dict:




											"""simple docstring"""







											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()







						def       __UpperCamelCase     (							self      :	str)						->       List[Any]:




											"""simple docstring"""







											_lowerCAmelCase:int									=  load_image(
											    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
											    '''/img2img/sketch-mountains-input.jpg''')
											_lowerCAmelCase:Tuple									=  init_image.resize((768, 512))
											_lowerCAmelCase:int									=  load_numpy(
											    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy''')

											_lowerCAmelCase:Optional[int]									=  '''BAAI/AltDiffusion'''
											_lowerCAmelCase:Optional[int]									=  AltDiffusionImgaImgPipeline.from_pretrained(
											    a__     ,safety_checker=a__     ,)
											pipe.to(a__)
											pipe.set_progress_bar_config(disable=a__)
											pipe.enable_attention_slicing()

											_lowerCAmelCase:Tuple									=  '''A fantasy landscape, trending on artstation'''

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

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












"""simple docstring"""






from typing import List, Union

import numpy as np

from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging, requires_backends
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
				from PIL import Image

				from ..image_utils import load_image

if is_torch_available():
				import torch

				from ..models.auto.modeling_auto import MODEL_FOR_DEPTH_ESTIMATION_MAPPING

UpperCamelCase__						=			logging.get_logger(__name__)




@add_end_docstrings(UpperCamelCase_  )
class   a__			(       UpperCamelCase_  ):





						def __init__(							self      :	int     ,*a__      :	Optional[Any]     ,**a__      :	Union[str, Any])						->       Tuple:




											"""simple docstring"""







											super().__init__(*a__     ,**a__)
											requires_backends(self     ,'''vision''')
											self.check_model_type(a__)





						def __call__(							self      :	str     ,a__      :	Union[str, List[str], "Image.Image", List["Image.Image"]]     ,**a__      :	List[str])						->       Optional[int]:




											"""simple docstring"""







											return super().__call__(a__     ,**a__)





						def       __UpperCamelCase     (							self      :	Union[str, Any]     ,**a__      :	List[Any])						->       Any:




											"""simple docstring"""







											return {}, {}, {}





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




											"""simple docstring"""







											_lowerCAmelCase:List[str]									=  load_image(a__)
											_lowerCAmelCase:int									=  image.size
											_lowerCAmelCase:int									=  self.image_processor(images=a__     ,return_tensors=self.framework)
											return model_inputs





						def       __UpperCamelCase     (							self      :	Dict     ,a__      :	List[str])						->       Union[str, Any]:




											"""simple docstring"""







											_lowerCAmelCase:Any									=  self.model(**a__)
											return model_outputs







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




											"""simple docstring"""







											_lowerCAmelCase:Optional[int]									=  model_outputs.predicted_depth
											_lowerCAmelCase:Union[str, Any]									=  torch.nn.functional.interpolate(
											    predicted_depth.unsqueeze(1)     ,size=self.image_size[::-1]     ,mode='''bicubic'''     ,align_corners=a__)
											_lowerCAmelCase:List[str]									=  prediction.squeeze().cpu().numpy()
											_lowerCAmelCase:Any									=  (output * 255 / np.max(a__)).astype('''uint8''')
											_lowerCAmelCase:Dict									=  Image.fromarray(a__)
											_lowerCAmelCase:Tuple									=  {}
											_lowerCAmelCase:Optional[int]									=  predicted_depth
											_lowerCAmelCase:str									=  depth
											return output_dict






