



import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, XLMRobertaTokenizer

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class       _SCREAMING_SNAKE_CASE	(UpperCamelCase			,    UpperCamelCase			,    UpperCamelCase			,    unittest.TestCase ):
     lowerCAmelCase					=			AltDiffusionPipeline
     lowerCAmelCase					=			TEXT_TO_IMAGE_PARAMS
     lowerCAmelCase					=			TEXT_TO_IMAGE_BATCH_PARAMS
     lowerCAmelCase					=			TEXT_TO_IMAGE_IMAGE_PARAMS
     lowerCAmelCase					=			TEXT_TO_IMAGE_IMAGE_PARAMS





     def 	__snake_case    (       self :							int				)->Any:
      torch.manual_seed(0				)
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	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     ,  )
      __SCREAMING_SNAKE_CASE       :    str       =	DDIMScheduler(
          beta_start=0.0_0_0_8_5     ,  beta_end=0.0_1_2     ,  beta_schedule="scaled_linear"     ,  clip_sample=UpperCamelCase     ,  set_alpha_to_one=UpperCamelCase     ,  )
      torch.manual_seed(0				)
      __SCREAMING_SNAKE_CASE       :    Tuple       =	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     ,  )

      # TODO: address the non-deterministic text encoder (fails for save-load tests)
      # torch.manual_seed(0)
      # text_encoder_config = RobertaSeriesConfig(
      #     hidden_size=32,
      #     project_dim=32,
      #     intermediate_size=37,
      #     layer_norm_eps=1e-05,
      #     num_attention_heads=4,
      #     num_hidden_layers=5,
      #     vocab_size=5002,
      # )
      # text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

      torch.manual_seed(0				)
      __SCREAMING_SNAKE_CASE       :    int       =	CLIPTextConfig(
          bos_token_id=0     ,  eos_token_id=2     ,  hidden_size=3_2     ,  projection_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_2     ,  )
      __SCREAMING_SNAKE_CASE       :    Any       =	CLIPTextModel(UpperCamelCase				)

      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta"				)
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	7_7

      __SCREAMING_SNAKE_CASE       :    int       =	{
          "unet": unet,
          "scheduler": scheduler,
          "vae": vae,
          "text_encoder": text_encoder,
          "tokenizer": tokenizer,
          "safety_checker": None,
          "feature_extractor": None,
      }
      return components





     def 	__snake_case    (       self :							List[Any]     ,  UpperCamelCase :							Dict     ,  UpperCamelCase :							Tuple=0				)->Any:
      if str(UpperCamelCase				).startswith("mps"				):
       __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	torch.manual_seed(UpperCamelCase				)
      else:
       __SCREAMING_SNAKE_CASE       :    Dict       =	torch.Generator(device=UpperCamelCase				).manual_seed(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	{
          "prompt": "A painting of a squirrel eating a burger",
          "generator": generator,
          "num_inference_steps": 2,
          "guidance_scale": 6.0,
          "output_type": "numpy",
      }
      return inputs





     def 	__snake_case    (       self :							Dict				)->List[str]:
      super().test_attention_slicing_forward_pass(expected_max_diff=3E-3				)





     def 	__snake_case    (       self :							Optional[int]				)->List[str]:
      super().test_inference_batch_single_identical(expected_max_diff=3E-3				)





     def 	__snake_case    (       self :							int				)->Union[str, Any]:
      __SCREAMING_SNAKE_CASE       :    Tuple       =	"cpu"  # ensure determinism for the device-dependent torch.Generator

      __SCREAMING_SNAKE_CASE       :    List[Any]       =	self.get_dummy_components()
      torch.manual_seed(0				)
      __SCREAMING_SNAKE_CASE       :    str       =	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     ,  vocab_size=5_0_0_2     ,  )
      # TODO: remove after fixing the non-deterministic text encoder
      __SCREAMING_SNAKE_CASE       :    Tuple       =	RobertaSeriesModelWithTransformation(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	text_encoder

      __SCREAMING_SNAKE_CASE       :    Tuple       =	AltDiffusionPipeline(**UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	alt_pipe.to(UpperCamelCase				)
      alt_pipe.set_progress_bar_config(disable=UpperCamelCase				)

      __SCREAMING_SNAKE_CASE       :    Any       =	self.get_dummy_inputs(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	"A photo of an astronaut"
      __SCREAMING_SNAKE_CASE       :    str       =	alt_pipe(**UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    List[str]       =	output.images
      __SCREAMING_SNAKE_CASE       :    int       =	image[0, -3:, -3:, -1]

      assert image.shape == (1, 6_4, 6_4, 3)
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	np.array(
          [0.5_7_4_8_1_6_2, 0.6_0_4_4_7_1_4_5, 0.4_8_8_2_1_2_1_7, 0.5_0_1_0_0_6_3_6, 0.5_4_3_1_1_8_5, 0.4_5_7_6_3_6_8_3, 0.4_9_6_5_7_6_9_6, 0.4_8_1_3_2_7_3_3, 0.4_7_5_7_3_0_9_3]				)

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






     def 	__snake_case    (       self :							Dict				)->Any:
      __SCREAMING_SNAKE_CASE       :    str       =	"cpu"  # ensure determinism for the device-dependent torch.Generator

      __SCREAMING_SNAKE_CASE       :    List[str]       =	self.get_dummy_components()
      __SCREAMING_SNAKE_CASE       :    Dict       =	PNDMScheduler(skip_prk_steps=UpperCamelCase				)
      torch.manual_seed(0				)
      __SCREAMING_SNAKE_CASE       :    List[str]       =	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     ,  vocab_size=5_0_0_2     ,  )
      # TODO: remove after fixing the non-deterministic text encoder
      __SCREAMING_SNAKE_CASE       :    Tuple       =	RobertaSeriesModelWithTransformation(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    List[str]       =	text_encoder
      __SCREAMING_SNAKE_CASE       :    Any       =	AltDiffusionPipeline(**UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    List[str]       =	alt_pipe.to(UpperCamelCase				)
      alt_pipe.set_progress_bar_config(disable=UpperCamelCase				)

      __SCREAMING_SNAKE_CASE       :    Tuple       =	self.get_dummy_inputs(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    List[str]       =	alt_pipe(**UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    str       =	output.images
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	image[0, -3:, -3:, -1]

      assert image.shape == (1, 6_4, 6_4, 3)
      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	np.array(
          [0.5_1_6_0_5_0_9_3, 0.5_7_0_7_2_4_1, 0.4_7_3_6_5_5_0_7, 0.5_0_5_7_8_8_8_6, 0.5_6_3_3_8_7_7, 0.4_6_4_2_5_0_3, 0.5_1_8_2_0_8_1, 0.4_8_7_6_3_4_8_4, 0.4_9_0_8_4_2_3_7]				)

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






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





     def 	__snake_case    (       self :							Union[str, Any]				)->List[Any]:
      # clean up the VRAM after each test
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()





     def 	__snake_case    (       self :							Optional[Any]				)->Optional[int]:
      # make sure here that pndm scheduler skips prk
      __SCREAMING_SNAKE_CASE       :    List[str]       =	AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion"     ,  safety_checker=UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	alt_pipe.to(UpperCamelCase				)
      alt_pipe.set_progress_bar_config(disable=UpperCamelCase				)

      __SCREAMING_SNAKE_CASE       :    int       =	"A painting of a squirrel eating a burger"
      __SCREAMING_SNAKE_CASE       :    Any       =	torch.manual_seed(0				)
      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	alt_pipe([prompt]     ,  generator=UpperCamelCase     ,  guidance_scale=6.0     ,  num_inference_steps=2_0     ,  output_type="np"				)

      __SCREAMING_SNAKE_CASE       :    str       =	output.images

      __SCREAMING_SNAKE_CASE       :    Dict       =	image[0, -3:, -3:, -1]

      assert image.shape == (1, 5_1_2, 5_1_2, 3)
      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	np.array([0.1_0_1_0, 0.0_8_0_0, 0.0_7_9_4, 0.0_8_8_5, 0.0_8_4_3, 0.0_7_6_2, 0.0_7_6_9, 0.0_7_2_9, 0.0_5_8_6]				)

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






     def 	__snake_case    (       self :							Tuple				)->Any:
      __SCREAMING_SNAKE_CASE       :    List[str]       =	DDIMScheduler.from_pretrained("BAAI/AltDiffusion"     ,  subfolder="scheduler"				)

      __SCREAMING_SNAKE_CASE       :    List[Any]       =	AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion"     ,  scheduler=UpperCamelCase     ,  safety_checker=UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Tuple       =	alt_pipe.to(UpperCamelCase				)
      alt_pipe.set_progress_bar_config(disable=UpperCamelCase				)

      __SCREAMING_SNAKE_CASE       :    Dict       =	"A painting of a squirrel eating a burger"
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	torch.manual_seed(0				)

      __SCREAMING_SNAKE_CASE       :    List[str]       =	alt_pipe([prompt]     ,  generator=UpperCamelCase     ,  num_inference_steps=2     ,  output_type="numpy"				)
      __SCREAMING_SNAKE_CASE       :    str       =	output.images

      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	image[0, -3:, -3:, -1]

      assert image.shape == (1, 5_1_2, 5_1_2, 3)
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	np.array([0.4_0_1_9, 0.4_0_5_2, 0.3_8_1_0, 0.4_1_1_9, 0.3_9_1_6, 0.3_9_8_2, 0.4_6_5_1, 0.4_1_9_5, 0.5_3_2_3]				)

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











from abc import ABC, abstractmethod
from typing import Optional, Union

from .. import Dataset, DatasetDict, Features, IterableDataset, IterableDatasetDict, NamedSplit
from ..utils.typing import NestedDataStructureLike, PathLike

class       _SCREAMING_SNAKE_CASE	(UpperCamelCase ):





     def __init__(       self :							List[str]     ,  UpperCamelCase :							Optional[NestedDataStructureLike[PathLike]] = None     ,  UpperCamelCase :							Optional[NamedSplit] = None     ,  UpperCamelCase :							Optional[Features] = None     ,  UpperCamelCase :							str = None     ,  UpperCamelCase :							bool = False     ,  UpperCamelCase :							bool = False     ,  UpperCamelCase :							Optional[int] = None     ,  **UpperCamelCase :							Dict     ,  )->int:
      __SCREAMING_SNAKE_CASE       :    List[str]       =	path_or_paths
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	split if split or isinstance(UpperCamelCase     ,  UpperCamelCase				) else "train"
      __SCREAMING_SNAKE_CASE       :    Dict       =	features
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	cache_dir
      __SCREAMING_SNAKE_CASE       :    Tuple       =	keep_in_memory
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	streaming
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	num_proc
      __SCREAMING_SNAKE_CASE       :    int       =	kwargs






     @abstractmethod
     def 	__snake_case    (       self :							Optional[int]				)->Union[Dataset, DatasetDict, IterableDataset, IterableDatasetDict]:
      pass






class       _SCREAMING_SNAKE_CASE	(UpperCamelCase ):





     def __init__(       self :							Optional[Any]     ,  UpperCamelCase :							Optional[Features] = None     ,  UpperCamelCase :							str = None     ,  UpperCamelCase :							bool = False     ,  UpperCamelCase :							bool = False     ,  UpperCamelCase :							Optional[int] = None     ,  **UpperCamelCase :							Dict     ,  )->Optional[int]:
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	features
      __SCREAMING_SNAKE_CASE       :    str       =	cache_dir
      __SCREAMING_SNAKE_CASE       :    Any       =	keep_in_memory
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	streaming
      __SCREAMING_SNAKE_CASE       :    int       =	num_proc
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	kwargs






     @abstractmethod
     def 	__snake_case    (       self :							str				)->Union[Dataset, IterableDataset]:
      pass






