






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

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

   @property
   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       Union[str, Any]:
       UpperCamelCase__         =    1
       UpperCamelCase__         =    3
       UpperCamelCase__         =    (32, 32)

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

   @property
   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       str:
       torch.manual_seed(0      )
       UpperCamelCase__         =    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 						SCREAMING_SNAKE_CASE__ (					self      )		->       int:
       torch.manual_seed(0      )
       UpperCamelCase__         =    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 						SCREAMING_SNAKE_CASE__ (					self      )		->       Union[str, Any]:
       torch.manual_seed(0      )
       UpperCamelCase__         =    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 						SCREAMING_SNAKE_CASE__ (					self      )		->       int:
       def extract(*snake_case_ ,				**snake_case_      ):
           class    __lowerCamelCase    :

              def __init__(					self      )		->       Any:
                  UpperCamelCase__         =    torch.ones([0]      )







              def 						SCREAMING_SNAKE_CASE__ (					self ,				snake_case_      )		->       Optional[Any]:
                  self.pixel_values.to(_a      )
                  return self







           return Out()

       return extract

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       Tuple:
       UpperCamelCase__         =    """cpu"""  # ensure determinism for the device-dependent torch.Generator
       UpperCamelCase__         =    self.dummy_cond_unet
       UpperCamelCase__         =    PNDMScheduler(skip_prk_steps=_a      )
       UpperCamelCase__         =    self.dummy_vae
       UpperCamelCase__         =    self.dummy_text_encoder
       UpperCamelCase__         =    XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'      )
       UpperCamelCase__         =    77

       UpperCamelCase__         =    self.dummy_image.to(_a      )
       UpperCamelCase__         =    init_image / 2 + 0.5

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

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

       UpperCamelCase__         =    output.images

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

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

       assert image.shape == (1, 32, 32, 3)
       UpperCamelCase__         =    np.array([0.4_427, 0.3_731, 0.4_249, 0.4_941, 0.4_546, 0.4_148, 0.4_193, 0.4_666, 0.4_499]      )

       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 						SCREAMING_SNAKE_CASE__ (					self      )		->       Optional[Any]:
       UpperCamelCase__         =    self.dummy_cond_unet
       UpperCamelCase__         =    PNDMScheduler(skip_prk_steps=_a      )
       UpperCamelCase__         =    self.dummy_vae
       UpperCamelCase__         =    self.dummy_text_encoder
       UpperCamelCase__         =    XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'      )
       UpperCamelCase__         =    77

       UpperCamelCase__         =    self.dummy_image.to(_a      )

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

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

       UpperCamelCase__         =    """A painting of a squirrel eating a burger"""
       UpperCamelCase__         =    torch.manual_seed(0      )
       UpperCamelCase__         =    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 						SCREAMING_SNAKE_CASE__ (					self      )		->       str:
       UpperCamelCase__         =    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
       UpperCamelCase__         =    init_image.resize((760, 504)      )

       UpperCamelCase__         =    """BAAI/AltDiffusion"""
       UpperCamelCase__         =    AltDiffusionImgaImgPipeline.from_pretrained(
           _a ,				safety_checker=_a ,				)
       pipe.to(_a      )
       pipe.set_progress_bar_config(disable=_a      )
       pipe.enable_attention_slicing()

       UpperCamelCase__         =    """A fantasy landscape, trending on artstation"""

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

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

       assert image.shape == (504, 760, 3)
       UpperCamelCase__         =    np.array([0.9_358, 0.9_397, 0.9_599, 0.9_901, 1.0_000, 1.0_000, 0.9_882, 1.0_000, 1.0_000]      )

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







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

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       Tuple:
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()







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

       UpperCamelCase__         =    """BAAI/AltDiffusion"""
       UpperCamelCase__         =    AltDiffusionImgaImgPipeline.from_pretrained(
           _a ,				safety_checker=_a ,				)
       pipe.to(_a      )
       pipe.set_progress_bar_config(disable=_a      )
       pipe.enable_attention_slicing()

       UpperCamelCase__         =    """A fantasy landscape, trending on artstation"""

       UpperCamelCase__         =    torch.manual_seed(0      )
       UpperCamelCase__         =    pipe(
           prompt=_a ,				image=_a ,				strength=0.75 ,				guidance_scale=7.5 ,				generator=_a ,				output_type='np' ,				)
       UpperCamelCase__         =    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 ...configuration_utils import PretrainedConfig
from ...utils import logging


A__					: List[str]=	logging.get_logger(__name__)
class    __lowerCamelCase    (  _a  ):
   a      :				Optional[int] 						="""timm_backbone"""

   def __init__(					self ,				snake_case_=None ,				snake_case_=3 ,				snake_case_=True ,				snake_case_=True ,				snake_case_=None ,				**snake_case_ ,				)		->       Dict:
       super().__init__(**snake_case_      )
       UpperCamelCase__         =    backbone
       UpperCamelCase__         =    num_channels
       UpperCamelCase__         =    features_only
       UpperCamelCase__         =    use_pretrained_backbone
       UpperCamelCase__         =    True
       UpperCamelCase__         =    out_indices if out_indices is not None else (-1,)




