






"""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       ):
      # clean up the VRAM after each test
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()







    @property
    def 			UpperCamelCase_      (			self       ):
      lowercase							    =       1
      lowercase							    =       3
      lowercase							    =       (3_2, 3_2)

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







    @property
    def 			UpperCamelCase_      (			self       ):
      torch.manual_seed(0       )
      lowercase							    =       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       ):
      torch.manual_seed(0       )
      lowercase							    =       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       ):
      torch.manual_seed(0       )
      lowercase							    =       RobertaSeriesConfig(
          hidden_size=3_2							,		project_dim=3_2							,		intermediate_size=3_7							,		layer_norm_eps=1e-05							,		num_attention_heads=4							,		num_hidden_layers=5							,		pad_token_id=1							,		vocab_size=5_0_0_6							,		)
      return RobertaSeriesModelWithTransformation(_lowerCamelCase       )







    @property
    def 			UpperCamelCase_      (			self       ):
      def extract(*_lowerCamelCase							,		**_lowerCamelCase       ):

        class        a     :







            def __init__(			self       ):
              lowercase							    =       torch.ones([0]       )







            def 			UpperCamelCase_      (			self							,		_lowerCamelCase       ):
              self.pixel_values.to(_lowerCamelCase       )
              return self


        return Out()

      return extract







    def 			UpperCamelCase_      (			self       ):
      lowercase							    =       'cpu'  # ensure determinism for the device-dependent torch.Generator
      lowercase							    =       self.dummy_cond_unet
      lowercase							    =       PNDMScheduler(skip_prk_steps=_lowerCamelCase       )
      lowercase							    =       self.dummy_vae
      lowercase							    =       self.dummy_text_encoder
      lowercase							    =       XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'       )
      lowercase							    =       7_7

      lowercase							    =       self.dummy_image.to(_lowerCamelCase       )
      lowercase							    =       init_image / 2 + 0.5

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

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

      lowercase							    =       output.images

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

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

      assert image.shape == (1, 3_2, 3_2, 3)
      lowercase							    =       np.array([0.4_4_2_7, 0.3_7_3_1, 0.4_2_4_9, 0.4_9_4_1, 0.4_5_4_6, 0.4_1_4_8, 0.4_1_9_3, 0.4_6_6_6, 0.4_4_9_9]       )

      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       ):
      lowercase							    =       self.dummy_cond_unet
      lowercase							    =       PNDMScheduler(skip_prk_steps=_lowerCamelCase       )
      lowercase							    =       self.dummy_vae
      lowercase							    =       self.dummy_text_encoder
      lowercase							    =       XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'       )
      lowercase							    =       7_7

      lowercase							    =       self.dummy_image.to(_lowerCamelCase       )

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

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

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

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







    @unittest.skipIf(torch_device != 'cuda'							,		'This test requires a GPU'       )
    def 			UpperCamelCase_      (			self       ):
      lowercase							    =       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
      lowercase							    =       init_image.resize((7_6_0, 5_0_4)       )

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

      lowercase							    =       'A fantasy landscape, trending on artstation'

      lowercase							    =       torch.manual_seed(0       )
      lowercase							    =       pipe(
          prompt=_lowerCamelCase							,		image=_lowerCamelCase							,		strength=0.7_5							,		guidance_scale=7.5							,		generator=_lowerCamelCase							,		output_type='np'							,		)
      lowercase							    =       output.images[0]

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

      assert image.shape == (5_0_4, 7_6_0, 3)
      lowercase							    =       np.array([0.9_3_5_8, 0.9_3_9_7, 0.9_5_9_9, 0.9_9_0_1, 1.0_0_0_0, 1.0_0_0_0, 0.9_8_8_2, 1.0_0_0_0, 1.0_0_0_0]       )

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


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







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







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

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

      lowercase							    =       'A fantasy landscape, trending on artstation'

      lowercase							    =       torch.manual_seed(0       )
      lowercase							    =       pipe(
          prompt=_lowerCamelCase							,		image=_lowerCamelCase							,		strength=0.7_5							,		guidance_scale=7.5							,		generator=_lowerCamelCase							,		output_type='np'							,		)
      lowercase							    =       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










"""simple docstring"""



import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
    import jax
    import jax.numpy as jnp
    from flax.jax_utils import replicate
    from flax.training.common_utils import shard

@slow
@require_flax
class        a     (   unittest.TestCase	):







    def 			UpperCamelCase_      (			self       ):
      # clean up the VRAM after each test
      super().tearDown()
      gc.collect()







    def 			UpperCamelCase_      (			self       ):
      lowercase    ,     lowercase							    =       FlaxStableDiffusionPipeline.from_pretrained(
          'stabilityai/stable-diffusion-2'							,		revision='bf16'							,		dtype=jnp.bfloataa							,		)

      lowercase							    =       'A painting of a squirrel eating a burger'
      lowercase							    =       jax.device_count()
      lowercase							    =       num_samples * [prompt]
      lowercase							    =       sd_pipe.prepare_inputs(_lowerCamelCase       )

      lowercase							    =       replicate(_lowerCamelCase       )
      lowercase							    =       shard(_lowerCamelCase       )

      lowercase							    =       jax.random.PRNGKey(0       )
      lowercase							    =       jax.random.split(_lowerCamelCase							,		jax.device_count()       )

      lowercase							    =       sd_pipe(_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		num_inference_steps=2_5							,		jit=_lowerCamelCase       )[0]
      assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

      lowercase							    =       images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]       )
      lowercase							    =       images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

      lowercase							    =       jnp.asarray(jax.device_get(image_slice.flatten()       )       )
      lowercase							    =       jnp.array([0.4_2_3_8, 0.4_4_1_4, 0.4_3_9_5, 0.4_4_5_3, 0.4_6_2_9, 0.4_5_9_0, 0.4_5_3_1, 0.4_5_5_0_8, 0.4_5_1_2]       )
      print(F'output_slice: {output_slice}'       )
      assert jnp.abs(output_slice - expected_slice       ).max() < 1e-2







    def 			UpperCamelCase_      (			self       ):
      lowercase							    =       'stabilityai/stable-diffusion-2'
      lowercase    ,     lowercase							    =       FlaxDPMSolverMultistepScheduler.from_pretrained(_lowerCamelCase							,		subfolder='scheduler'       )
      lowercase    ,     lowercase							    =       FlaxStableDiffusionPipeline.from_pretrained(
          _lowerCamelCase							,		scheduler=_lowerCamelCase							,		revision='bf16'							,		dtype=jnp.bfloataa							,		)
      lowercase							    =       scheduler_params

      lowercase							    =       'A painting of a squirrel eating a burger'
      lowercase							    =       jax.device_count()
      lowercase							    =       num_samples * [prompt]
      lowercase							    =       sd_pipe.prepare_inputs(_lowerCamelCase       )

      lowercase							    =       replicate(_lowerCamelCase       )
      lowercase							    =       shard(_lowerCamelCase       )

      lowercase							    =       jax.random.PRNGKey(0       )
      lowercase							    =       jax.random.split(_lowerCamelCase							,		jax.device_count()       )

      lowercase							    =       sd_pipe(_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		num_inference_steps=2_5							,		jit=_lowerCamelCase       )[0]
      assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

      lowercase							    =       images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]       )
      lowercase							    =       images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

      lowercase							    =       jnp.asarray(jax.device_get(image_slice.flatten()       )       )
      lowercase							    =       jnp.array([0.4_3_3_6, 0.4_2_9_6_9, 0.4_4_5_3, 0.4_1_9_9, 0.4_2_9_7, 0.4_5_3_1, 0.4_4_3_4, 0.4_4_3_4, 0.4_2_9_7]       )
      print(F'output_slice: {output_slice}'       )
      assert jnp.abs(output_slice - expected_slice       ).max() < 1e-2


