





"""simple docstring"""





import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
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 PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class 		lowercase ( _UpperCAmelCase					,  _UpperCAmelCase					,  unittest.TestCase ):
      _SCREAMING_SNAKE_CASE         =				StableDiffusionParadigmsPipeline
      _SCREAMING_SNAKE_CASE         =				TEXT_TO_IMAGE_PARAMS
      _SCREAMING_SNAKE_CASE         =				TEXT_TO_IMAGE_BATCH_PARAMS
      _SCREAMING_SNAKE_CASE         =				TEXT_TO_IMAGE_IMAGE_PARAMS
      _SCREAMING_SNAKE_CASE         =				TEXT_TO_IMAGE_IMAGE_PARAMS





      def    _snake_case      (	self			)		->					Optional[Any]:
            torch.manual_seed(0			)
            lowerCAmelCase  						=	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				,						attention_head_dim=(2, 4)				,						use_linear_projection=lowercase				,						)
            lowerCAmelCase  						=	DDIMParallelScheduler(
                beta_start=0.00_085				,						beta_end=0.012				,						beta_schedule="""scaled_linear"""				,						clip_sample=lowercase				,						set_alpha_to_one=lowercase				,						)
            torch.manual_seed(0			)
            lowerCAmelCase  						=	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				,						sample_size=128				,						)
            torch.manual_seed(0			)
            lowerCAmelCase  						=	CLIPTextConfig(
                bos_token_id=0				,						eos_token_id=2				,						hidden_size=32				,						intermediate_size=37				,						layer_norm_eps=1e-05				,						num_attention_heads=4				,						num_hidden_layers=5				,						pad_token_id=1				,						vocab_size=1_000				,						hidden_act="""gelu"""				,						projection_dim=512				,						)
            lowerCAmelCase  						=	CLIPTextModel(lowercase			)
            lowerCAmelCase  						=	CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""			)

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





      def    _snake_case      (	self				,						lowercase				,						lowercase=0			)		->					Optional[int]:
            if str(lowercase			).startswith("""mps"""			):
                  lowerCAmelCase  						=	torch.manual_seed(lowercase			)
            else:
                  lowerCAmelCase  						=	torch.Generator(device=lowercase			).manual_seed(lowercase			)
            lowerCAmelCase  						=	{
                """prompt""": """a photograph of an astronaut riding a horse""",
                """generator""": generator,
                """num_inference_steps""": 10,
                """guidance_scale""": 6.0,
                """output_type""": """numpy""",
                """parallel""": 3,
                """debug""": True,
            }
            return inputs





      def    _snake_case      (	self			)		->					int:
            lowerCAmelCase  						=	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
            lowerCAmelCase  						=	self.get_dummy_components()
            lowerCAmelCase  						=	StableDiffusionParadigmsPipeline(**lowercase			)
            lowerCAmelCase  						=	sd_pipe.to(lowercase			)
            sd_pipe.set_progress_bar_config(disable=lowercase			)

            lowerCAmelCase  						=	self.get_dummy_inputs(lowercase			)
            lowerCAmelCase  						=	sd_pipe(**lowercase			).images
            lowerCAmelCase  						=	image[0, -3:, -3:, -1]
            assert image.shape == (1, 64, 64, 3)

            lowerCAmelCase  						=	np.array([0.4_773, 0.5_417, 0.4_723, 0.4_925, 0.5_631, 0.4_752, 0.5_240, 0.4_935, 0.5_023]			)

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





      def    _snake_case      (	self			)		->					str:
            lowerCAmelCase  						=	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
            lowerCAmelCase  						=	self.get_dummy_components()
            torch.manual_seed(0			)
            lowerCAmelCase  						=	DDPMParallelScheduler()
            torch.manual_seed(0			)
            lowerCAmelCase  						=	StableDiffusionParadigmsPipeline(**lowercase			)
            lowerCAmelCase  						=	sd_pipe.to(lowercase			)
            sd_pipe.set_progress_bar_config(disable=lowercase			)

            lowerCAmelCase  						=	self.get_dummy_inputs(lowercase			)
            lowerCAmelCase  						=	sd_pipe(**lowercase			).images
            lowerCAmelCase  						=	image[0, -3:, -3:, -1]
            assert image.shape == (1, 64, 64, 3)

            lowerCAmelCase  						=	np.array([0.3_573, 0.4_420, 0.4_960, 0.4_799, 0.3_796, 0.3_879, 0.4_819, 0.4_365, 0.4_468]			)

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





      def    _snake_case      (	self			)		->					Optional[int]:
            super().test_inference_batch_consistent(batch_sizes=[1, 2]			)





      def    _snake_case      (	self			)		->					Union[str, Any]:
            super().test_inference_batch_single_identical(batch_size=2				,						expected_max_diff=3e-3			)





      def    _snake_case      (	self			)		->					int:
            lowerCAmelCase  						=	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
            lowerCAmelCase  						=	self.get_dummy_components()
            lowerCAmelCase  						=	StableDiffusionParadigmsPipeline(**lowercase			)
            lowerCAmelCase  						=	sd_pipe.to(lowercase			)
            sd_pipe.set_progress_bar_config(disable=lowercase			)

            lowerCAmelCase  						=	self.get_dummy_inputs(lowercase			)
            lowerCAmelCase  						=	"""french fries"""
            lowerCAmelCase  						=	sd_pipe(**lowercase				,						negative_prompt=lowercase			)
            lowerCAmelCase  						=	output.images
            lowerCAmelCase  						=	image[0, -3:, -3:, -1]

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

            lowerCAmelCase  						=	np.array([0.4_771, 0.5_420, 0.4_683, 0.4_918, 0.5_636, 0.4_725, 0.5_230, 0.4_923, 0.5_015]			)

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







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





      def    _snake_case      (	self			)		->					Dict:
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()





      def    _snake_case      (	self				,						lowercase=0			)		->					List[Any]:
            lowerCAmelCase  						=	torch.Generator(device=lowercase			).manual_seed(lowercase			)
            lowerCAmelCase  						=	{
                """prompt""": """a photograph of an astronaut riding a horse""",
                """generator""": generator,
                """num_inference_steps""": 10,
                """guidance_scale""": 7.5,
                """output_type""": """numpy""",
                """parallel""": 3,
                """debug""": True,
            }
            return inputs





      def    _snake_case      (	self			)		->					int:
            lowerCAmelCase  						=	"""stabilityai/stable-diffusion-2-base"""
            lowerCAmelCase  						=	DDIMParallelScheduler.from_pretrained(lowercase				,						subfolder="""scheduler"""			)
            lowerCAmelCase  						=	StableDiffusionParadigmsPipeline.from_pretrained(lowercase				,						scheduler=lowercase				,						safety_checker=lowercase			)
            pipe.to(lowercase			)
            pipe.set_progress_bar_config(disable=lowercase			)
            pipe.enable_attention_slicing()

            lowerCAmelCase  						=	self.get_inputs()
            lowerCAmelCase  						=	pipe(**lowercase			).images
            lowerCAmelCase  						=	image[0, -3:, -3:, -1].flatten()

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

            lowerCAmelCase  						=	np.array([0.9_622, 0.9_602, 0.9_748, 0.9_591, 0.9_630, 0.9_691, 0.9_661, 0.9_631, 0.9_741]			)

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













"""simple docstring"""





import unittest
from dataclasses import dataclass
from typing import Optional

from transformers.utils import ModelOutput

@dataclass
class 		lowercase ( _UpperCAmelCase ):
      _SCREAMING_SNAKE_CASE         =				42
      _SCREAMING_SNAKE_CASE         =				None
      _SCREAMING_SNAKE_CASE         =				None







class 		lowercase ( unittest.TestCase ):





      def    _snake_case      (	self			)		->					Optional[Any]:
            lowerCAmelCase  						=	ModelOutputTest(a=30			)
            self.assertEqual(x.a				,						30			)
            self.assertIsNone(x.b			)
            self.assertIsNone(x.c			)
            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	x.d





      def    _snake_case      (	self			)		->					Optional[Any]:
            lowerCAmelCase  						=	ModelOutputTest(a=30				,						b=10			)
            self.assertEqual(x[0]				,						30			)
            self.assertEqual(x[1]				,						10			)
            self.assertEqual(x[:2]				,						(30, 10)			)
            self.assertEqual(x[:]				,						(30, 10)			)

            lowerCAmelCase  						=	ModelOutputTest(a=30				,						c=10			)
            self.assertEqual(x[0]				,						30			)
            self.assertEqual(x[1]				,						10			)
            self.assertEqual(x[:2]				,						(30, 10)			)
            self.assertEqual(x[:]				,						(30, 10)			)





      def    _snake_case      (	self			)		->					Tuple:
            lowerCAmelCase  						=	ModelOutputTest(a=30				,						b=10			)
            self.assertEqual(x["""a"""]				,						30			)
            self.assertEqual(x["""b"""]				,						10			)
            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	x["""c"""]

            lowerCAmelCase  						=	ModelOutputTest(a=30				,						c=10			)
            self.assertEqual(x["""a"""]				,						30			)
            self.assertEqual(x["""c"""]				,						10			)
            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	x["""b"""]





      def    _snake_case      (	self			)		->					Union[str, Any]:
            lowerCAmelCase  						=	ModelOutputTest(a=30			)
            self.assertEqual(list(x.keys()			)				,						["""a"""]			)
            self.assertEqual(list(x.values()			)				,						[30]			)
            self.assertEqual(list(x.items()			)				,						[("""a""", 30)]			)
            self.assertEqual(list(lowercase			)				,						["""a"""]			)

            lowerCAmelCase  						=	ModelOutputTest(a=30				,						b=10			)
            self.assertEqual(list(x.keys()			)				,						["""a""", """b"""]			)
            self.assertEqual(list(x.values()			)				,						[30, 10]			)
            self.assertEqual(list(x.items()			)				,						[("""a""", 30), ("""b""", 10)]			)
            self.assertEqual(list(lowercase			)				,						["""a""", """b"""]			)

            lowerCAmelCase  						=	ModelOutputTest(a=30				,						c=10			)
            self.assertEqual(list(x.keys()			)				,						["""a""", """c"""]			)
            self.assertEqual(list(x.values()			)				,						[30, 10]			)
            self.assertEqual(list(x.items()			)				,						[("""a""", 30), ("""c""", 10)]			)
            self.assertEqual(list(lowercase			)				,						["""a""", """c"""]			)

            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	x.update({"""d""": 20}			)
            with self.assertRaises(lowercase			):
                  del x["a"]
            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	x.pop("""a"""			)
            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	x.setdefault("""d"""				,						32			)





      def    _snake_case      (	self			)		->					Dict:
            lowerCAmelCase  						=	ModelOutputTest(a=30			)
            lowerCAmelCase  						=	10
            self.assertEqual(x.a				,						10			)
            self.assertEqual(x["""a"""]				,						10			)





      def    _snake_case      (	self			)		->					List[str]:
            lowerCAmelCase  						=	ModelOutputTest(a=30			)
            lowerCAmelCase  						=	10
            self.assertEqual(x.a				,						10			)
            self.assertEqual(x["""a"""]				,						10			)





      def    _snake_case      (	self			)		->					str:
            lowerCAmelCase  						=	ModelOutputTest({"""a""": 30, """b""": 10}			)
            self.assertEqual(list(x.keys()			)				,						["""a""", """b"""]			)
            self.assertEqual(x.a				,						30			)
            self.assertEqual(x.b				,						10			)





      def    _snake_case      (	self			)		->					Dict:
            lowerCAmelCase  						=	ModelOutputTest([("""a""", 30), ("""b""", 10)]			)
            self.assertEqual(list(x.keys()			)				,						["""a""", """b"""]			)
            self.assertEqual(x.a				,						30			)
            self.assertEqual(x.b				,						10			)

            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	ModelOutputTest([("""a""", 30), (10, 10)]			)

            lowerCAmelCase  						=	ModelOutputTest(a=(30, 30)			)
            self.assertEqual(list(x.keys()			)				,						["""a"""]			)
            self.assertEqual(x.a				,						(30, 30)			)






