





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





       '''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 :    Optional[Any]		):
         __snake_case:						Optional[int]   =						1
         __snake_case:						int   =						3
         __snake_case:						Dict   =						(32, 32)

         __snake_case:						List[str]   =						floats_tensor((batch_size, num_channels) + sizes							,				rng=random.Random(0		)		).to(A		)
         return image




       @property
       def 							UpperCAmelCase__ (  self :    int		):
         torch.manual_seed(0		)
         __snake_case:						List[str]   =						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 :    Dict		):
         torch.manual_seed(0		)
         __snake_case:						Tuple   =						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 :    Any		):
         torch.manual_seed(0		)
         __snake_case:						str   =						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=5_006							,				)
         return RobertaSeriesModelWithTransformation(A		)




       @property
       def 							UpperCAmelCase__ (  self :    Tuple		):
         def extract(*A :    Any							,				**A :    Dict		):







           class     __snake_case       :





                  '''simple docstring'''







                  def __init__(  self :    Union[str, Any]		):
                    __snake_case:						Optional[Any]   =						torch.ones([0]		)






                  def 							UpperCAmelCase__ (  self :    int							,				A :    str		):
                    self.pixel_values.to(A		)
                    return self






           return Out()

         return extract




       def 							UpperCAmelCase__ (  self :    Optional[int]		):
         __snake_case:						str   =						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
         __snake_case:						Tuple   =						self.dummy_cond_unet
         __snake_case:						List[str]   =						PNDMScheduler(skip_prk_steps=A		)
         __snake_case:						Dict   =						self.dummy_vae
         __snake_case:						Optional[int]   =						self.dummy_text_encoder
         __snake_case:						str   =						XLMRobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-xlm-roberta"""		)
         __snake_case:						Tuple   =						77

         __snake_case:						Optional[int]   =						self.dummy_image.to(A		)
         __snake_case:						Dict   =						init_image / 2 + 0.5

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

         __snake_case:						Tuple   =						"""A painting of a squirrel eating a burger"""
         __snake_case:						Optional[Any]   =						torch.Generator(device=A		).manual_seed(0		)
         __snake_case:						int   =						alt_pipe(
             [prompt]							,				generator=A							,				guidance_scale=6.0							,				num_inference_steps=2							,				output_type="""np"""							,				image=A							,				)

         __snake_case:						Dict   =						output.images

         __snake_case:						Union[str, Any]   =						torch.Generator(device=A		).manual_seed(0		)
         __snake_case:						str   =						alt_pipe(
             [prompt]							,				generator=A							,				guidance_scale=6.0							,				num_inference_steps=2							,				output_type="""np"""							,				image=A							,				return_dict=A							,				)[0]

         __snake_case:						List[Any]   =						image[0, -3:, -3:, -1]
         __snake_case:						int   =						image_from_tuple[0, -3:, -3:, -1]

         assert image.shape == (1, 32, 32, 3)
         __snake_case:						str   =						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 :    List[Any]		):
         __snake_case:						int   =						self.dummy_cond_unet
         __snake_case:						Optional[int]   =						PNDMScheduler(skip_prk_steps=A		)
         __snake_case:						Any   =						self.dummy_vae
         __snake_case:						Tuple   =						self.dummy_text_encoder
         __snake_case:						str   =						XLMRobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-xlm-roberta"""		)
         __snake_case:						Any   =						77

         __snake_case:						Tuple   =						self.dummy_image.to(A		)

         # put models in fp16
         __snake_case:						List[Any]   =						unet.half()
         __snake_case:						Union[str, Any]   =						vae.half()
         __snake_case:						Optional[Any]   =						bert.half()

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

         __snake_case:						Union[str, Any]   =						"""A painting of a squirrel eating a burger"""
         __snake_case:						Dict   =						torch.manual_seed(0		)
         __snake_case:						Union[str, Any]   =						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 :    int		):
         __snake_case:						int   =						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
         __snake_case:						Union[str, Any]   =						init_image.resize((760, 504)		)

         __snake_case:						List[Any]   =						"""BAAI/AltDiffusion"""
         __snake_case:						List[Any]   =						AltDiffusionImgaImgPipeline.from_pretrained(
             A							,				safety_checker=A							,				)
         pipe.to(A		)
         pipe.set_progress_bar_config(disable=A		)
         pipe.enable_attention_slicing()

         __snake_case:						Any   =						"""A fantasy landscape, trending on artstation"""

         __snake_case:						Dict   =						torch.manual_seed(0		)
         __snake_case:						Tuple   =						pipe(
             prompt=A							,				image=A							,				strength=0.75							,				guidance_scale=7.5							,				generator=A							,				output_type="""np"""							,				)
         __snake_case:						str   =						output.images[0]

         __snake_case:						str   =						image[255:258, 383:386, -1]

         assert image.shape == (504, 760, 3)
         __snake_case:						Optional[int]   =						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     __snake_case       (							unittest.TestCase			):





       '''simple docstring'''







       def 							UpperCAmelCase__ (  self :    Dict		):
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()






       def 							UpperCAmelCase__ (  self :    Any		):
         __snake_case:						Dict   =						load_image(
             """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
             """/img2img/sketch-mountains-input.jpg"""		)
         __snake_case:						List[str]   =						init_image.resize((768, 512)		)
         __snake_case:						Optional[Any]   =						load_numpy(
             """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy"""		)

         __snake_case:						Union[str, Any]   =						"""BAAI/AltDiffusion"""
         __snake_case:						Tuple   =						AltDiffusionImgaImgPipeline.from_pretrained(
             A							,				safety_checker=A							,				)
         pipe.to(A		)
         pipe.set_progress_bar_config(disable=A		)
         pipe.enable_attention_slicing()

         __snake_case:						Optional[int]   =						"""A fantasy landscape, trending on artstation"""

         __snake_case:						Optional[Any]   =						torch.manual_seed(0		)
         __snake_case:						Optional[Any]   =						pipe(
             prompt=A							,				image=A							,				strength=0.75							,				guidance_scale=7.5							,				generator=A							,				output_type="""np"""							,				)
         __snake_case:						List[str]   =						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













__UpperCAmelCase       :   int				       =							[
    "Audio",
    "Array2D",
    "Array3D",
    "Array4D",
    "Array5D",
    "ClassLabel",
    "Features",
    "Sequence",
    "Value",
    "Image",
    "Translation",
    "TranslationVariableLanguages",
]
from .audio import Audio
from .features import ArrayaD, ArrayaD, ArrayaD, ArrayaD, ClassLabel, Features, Sequence, Value
from .image import Image
from .translation import Translation, TranslationVariableLanguages






