
from ...utils import deprecate
from ..controlnet.pipeline_flax_controlnet import FlaxStableDiffusionControlNetPipeline  # noqa: F401


deprecate(
    '''stable diffusion controlnet''',
    '''0.22.0''',
    '''Importing `FlaxStableDiffusionControlNetPipeline` from diffusers.pipelines.stable_diffusion.flax_pipeline_stable_diffusion_controlnet is deprecated. Please import `from diffusers import FlaxStableDiffusionControlNetPipeline` instead.''',
    standard_warn=False,
    stacklevel=3,
)





"""simple docstring"""
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 				lowerCamelCase  (			_lowerCAmelCase    ,						_lowerCAmelCase    ,						unittest.TestCase				):






					'''simple docstring'''





					_A		:			str						=     StableDiffusionSAGPipeline
					_A		:			Optional[Any]						=     TEXT_TO_IMAGE_PARAMS
					_A		:			Any						=     TEXT_TO_IMAGE_BATCH_PARAMS
					_A		:			Tuple						=     TEXT_TO_IMAGE_IMAGE_PARAMS
					_A		:			Tuple						=     TEXT_TO_IMAGE_IMAGE_PARAMS
					_A		:			List[str]						=     False






					def    lowerCAmelCase_				(       self:							Optional[Any]						)   ->	str:
											torch.manual_seed(0						)
											snake_case_ :Any						=			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						,     )
											snake_case_ :Any						=			DDIMScheduler(
											    beta_start=0.0_0_0_8_5						,     beta_end=0.0_1_2						,     beta_schedule="""scaled_linear"""						,     clip_sample=snake_case						,     set_alpha_to_one=snake_case						,     )
											torch.manual_seed(0						)
											snake_case_ :Optional[int]						=			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						)
											snake_case_ :Union[str, Any]						=			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						,     )
											snake_case_ :Tuple						=			CLIPTextModel(snake_case						)
											snake_case_ :str						=			CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

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






					def    lowerCAmelCase_				(       self:							List[str]						,     snake_case:							Tuple						,     snake_case:							List[str]=0						)   ->	str:
											if str(snake_case						).startswith("""mps"""						):
																	snake_case_ :Tuple						=			torch.manual_seed(snake_case						)
											else:
																	snake_case_ :Optional[int]						=			torch.Generator(device=snake_case						).manual_seed(snake_case						)
											snake_case_ :Any						=			{
											    """prompt""": """.""",
											    """generator""": generator,
											    """num_inference_steps""": 2,
											    """guidance_scale""": 1.0,
											    """sag_scale""": 1.0,
											    """output_type""": """numpy""",
											}
											return inputs






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

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






					'''simple docstring'''





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






					def    lowerCAmelCase_				(       self:							int						)   ->	List[str]:
											snake_case_ :Any						=			StableDiffusionSAGPipeline.from_pretrained("""CompVis/stable-diffusion-v1-4"""						)
											snake_case_ :int						=			sag_pipe.to(snake_case						)
											sag_pipe.set_progress_bar_config(disable=snake_case						)

											snake_case_ :Union[str, Any]						=			"""."""
											snake_case_ :str						=			torch.manual_seed(0						)
											snake_case_ :str						=			sag_pipe(
											    [prompt]						,     generator=snake_case						,     guidance_scale=7.5						,     sag_scale=1.0						,     num_inference_steps=20						,     output_type="""np"""						)

											snake_case_ :List[Any]						=			output.images

											snake_case_ :Tuple						=			image[0, -3:, -3:, -1]

											assert image.shape == (1, 512, 512, 3)
											snake_case_ :List[Any]						=			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:							Dict						)   ->	str:
											snake_case_ :Tuple						=			StableDiffusionSAGPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base"""						)
											snake_case_ :Optional[int]						=			sag_pipe.to(snake_case						)
											sag_pipe.set_progress_bar_config(disable=snake_case						)

											snake_case_ :Tuple						=			"""."""
											snake_case_ :Union[str, Any]						=			torch.manual_seed(0						)
											snake_case_ :Tuple						=			sag_pipe(
											    [prompt]						,     generator=snake_case						,     guidance_scale=7.5						,     sag_scale=1.0						,     num_inference_steps=20						,     output_type="""np"""						)

											snake_case_ :Optional[int]						=			output.images

											snake_case_ :Optional[Any]						=			image[0, -3:, -3:, -1]

											assert image.shape == (1, 512, 512, 3)
											snake_case_ :Tuple						=			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:							List[str]						)   ->	List[str]:
											snake_case_ :Optional[int]						=			StableDiffusionSAGPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base"""						)
											snake_case_ :int						=			sag_pipe.to(snake_case						)
											sag_pipe.set_progress_bar_config(disable=snake_case						)

											snake_case_ :Tuple						=			"""."""
											snake_case_ :Optional[int]						=			torch.manual_seed(0						)
											snake_case_ :List[str]						=			sag_pipe(
											    [prompt]						,     width=768						,     height=512						,     generator=snake_case						,     guidance_scale=7.5						,     sag_scale=1.0						,     num_inference_steps=20						,     output_type="""np"""						,     )

											snake_case_ :Optional[Any]						=			output.images

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