




import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    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       a__	( UpperCamelCase__       , UpperCamelCase__       , unittest.TestCase				):
     a		:      Optional[Any]		 =     StableDiffusionSAGPipeline
     a		:      List[str]		 =     TEXT_TO_IMAGE_PARAMS
     a		:      List[str]		 =     TEXT_TO_IMAGE_BATCH_PARAMS
     a		:      str		 =     TEXT_TO_IMAGE_IMAGE_PARAMS
     a		:      List[str]		 =     TEXT_TO_IMAGE_IMAGE_PARAMS
     a		:      int		 =     False







     def       lowerCAmelCase_    (     self			)  ->   List[str]:



          '''simple docstring'''



          torch.manual_seed(0			)
          a     =				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  ,		)
          a     =				DDIMScheduler(
              beta_start=0.0_0_0_8_5  ,		beta_end=0.0_1_2  ,		beta_schedule="scaled_linear"  ,		clip_sample=A  ,		set_alpha_to_one=A  ,		)
          torch.manual_seed(0			)
          a     =				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  ,		)
          torch.manual_seed(0			)
          a     =				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=1000  ,		)
          a     =				CLIPTextModel(A			)
          a     =				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"			)

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







     def       lowerCAmelCase_    (     self  ,		A  ,		A=0			)  ->   str:



          '''simple docstring'''



          if str(A			).startswith("mps"			):
               a     =				torch.manual_seed(A			)
          else:
               a     =				torch.Generator(device=A			).manual_seed(A			)
          a     =				{
              "prompt": ".",
              "generator": generator,
              "num_inference_steps": 2,
              "guidance_scale": 1.0,
              "sag_scale": 1.0,
              "output_type": "numpy",
          }
          return inputs







     def       lowerCAmelCase_    (     self			)  ->   List[str]:



          '''simple docstring'''



          super().test_inference_batch_single_identical(expected_max_diff=3e-3			)







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







     def       lowerCAmelCase_    (     self			)  ->   List[Any]:



          '''simple docstring'''



          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()







     def       lowerCAmelCase_    (     self			)  ->   Any:



          '''simple docstring'''



          a     =				StableDiffusionSAGPipeline.from_pretrained("CompVis/stable-diffusion-v1-4"			)
          a     =				sag_pipe.to(A			)
          sag_pipe.set_progress_bar_config(disable=A			)

          a     =				"."
          a     =				torch.manual_seed(0			)
          a     =				sag_pipe(
              [prompt]  ,		generator=A  ,		guidance_scale=7.5  ,		sag_scale=1.0  ,		num_inference_steps=20  ,		output_type="np"			)

          a     =				output.images

          a     =				image[0, -3:, -3:, -1]

          assert image.shape == (1, 512, 512, 3)
          a     =				np.array([0.1_5_6_8, 0.1_7_3_8, 0.1_6_9_5, 0.1_6_9_3, 0.1_5_0_7, 0.1_7_0_5, 0.1_5_4_7, 0.1_7_5_1, 0.1_9_4_9]			)

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







     def       lowerCAmelCase_    (     self			)  ->   int:



          '''simple docstring'''



          a     =				StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"			)
          a     =				sag_pipe.to(A			)
          sag_pipe.set_progress_bar_config(disable=A			)

          a     =				"."
          a     =				torch.manual_seed(0			)
          a     =				sag_pipe(
              [prompt]  ,		generator=A  ,		guidance_scale=7.5  ,		sag_scale=1.0  ,		num_inference_steps=20  ,		output_type="np"			)

          a     =				output.images

          a     =				image[0, -3:, -3:, -1]

          assert image.shape == (1, 512, 512, 3)
          a     =				np.array([0.3_4_5_9, 0.2_8_7_6, 0.2_5_3_7, 0.3_0_0_2, 0.2_6_7_1, 0.2_1_6_0, 0.3_0_2_6, 0.2_2_6_2, 0.2_3_7_1]			)

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







     def       lowerCAmelCase_    (     self			)  ->   Optional[Any]:



          '''simple docstring'''



          a     =				StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"			)
          a     =				sag_pipe.to(A			)
          sag_pipe.set_progress_bar_config(disable=A			)

          a     =				"."
          a     =				torch.manual_seed(0			)
          a     =				sag_pipe(
              [prompt]  ,		width=768  ,		height=512  ,		generator=A  ,		guidance_scale=7.5  ,		sag_scale=1.0  ,		num_inference_steps=20  ,		output_type="np"  ,		)

          a     =				output.images

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







lowercase__				:      Optional[int]	   = [
    "DownloadConfig",
    "DownloadManager",
    "DownloadMode",
    "StreamingDownloadManager",
]

from .download_config import DownloadConfig
from .download_manager import DownloadManager, DownloadMode
from .streaming_download_manager import StreamingDownloadManager

