




"""simple docstring"""




import gc
import random
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionInstructPixaPixPipeline,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.utils import floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






class 					__magic_name__							( UpperCAmelCase__		,      UpperCAmelCase__		,      UpperCAmelCase__		,      unittest.TestCase       ):






				'''simple docstring'''
				__UpperCamelCase    						=		StableDiffusionInstructPixaPixPipeline
				__UpperCamelCase    						=		TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width", "cross_attention_kwargs"}
				__UpperCamelCase    						=		TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
				__UpperCamelCase    						=		IMAGE_TO_IMAGE_IMAGE_PARAMS
				__UpperCamelCase    						=		IMAGE_TO_IMAGE_IMAGE_PARAMS





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							torch.manual_seed(0				)
							lowerCamelCase       =      UNetaDConditionModel(
							    block_out_channels=(32, 64)	,				layers_per_block=2	,				sample_size=32	,				in_channels=8	,				out_channels=4	,				down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")	,				up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")	,				cross_attention_dim=32	,				)
							lowerCamelCase       =      PNDMScheduler(skip_prk_steps=_a				)
							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	,				)
							torch.manual_seed(0				)
							lowerCamelCase       =      CLIPTextConfig(
							    bos_token_id=0	,				eos_token_id=2	,				hidden_size=32	,				intermediate_size=37	,				layer_norm_eps=1e-0_5	,				num_attention_heads=4	,				num_hidden_layers=5	,				pad_token_id=1	,				vocab_size=1_000	,				)
							lowerCamelCase       =      CLIPTextModel(_a				)
							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      _lowerCAmelCase       ( self	,				_a	,				_a=0				):





							"""simple docstring"""
							lowerCamelCase       =      floats_tensor((1, 3, 32, 32)	,				rng=random.Random(_a				)				).to(_a				)
							lowerCamelCase       =      image.cpu().permute(0	,				2	,				3	,				1				)[0]
							lowerCamelCase       =      Image.fromarray(np.uinta(_a				)				).convert("""RGB"""				)
							if str(_a				).startswith("""mps"""				):
										lowerCamelCase       =      torch.manual_seed(_a				)
							else:
										lowerCamelCase       =      torch.Generator(device=_a				).manual_seed(_a				)
							lowerCamelCase       =      {
							    """prompt""": """A painting of a squirrel eating a burger""",
							    """image""": image,
							    """generator""": generator,
							    """num_inference_steps""": 2,
							    """guidance_scale""": 6.0,
							    """image_guidance_scale""": 1,
							    """output_type""": """numpy""",
							}
							return inputs





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      """cpu"""  # ensure determinism for the device-dependent torch.Generator
							lowerCamelCase       =      self.get_dummy_components()
							lowerCamelCase       =      StableDiffusionInstructPixaPixPipeline(**_a				)
							lowerCamelCase       =      sd_pipe.to(_a				)
							sd_pipe.set_progress_bar_config(disable=_a				)

							lowerCamelCase       =      self.get_dummy_inputs(_a				)
							lowerCamelCase       =      sd_pipe(**_a				).images
							lowerCamelCase       =      image[0, -3:, -3:, -1]
							assert image.shape == (1, 32, 32, 3)
							lowerCamelCase       =      np.array([0.7_526, 0.3_750, 0.4_547, 0.6_117, 0.5_866, 0.5_016, 0.4_327, 0.5_642, 0.4_815]				)

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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      """cpu"""  # ensure determinism for the device-dependent torch.Generator
							lowerCamelCase       =      self.get_dummy_components()
							lowerCamelCase       =      StableDiffusionInstructPixaPixPipeline(**_a				)
							lowerCamelCase       =      sd_pipe.to(_a				)
							sd_pipe.set_progress_bar_config(disable=_a				)

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

							assert image.shape == (1, 32, 32, 3)
							lowerCamelCase       =      np.array([0.7_511, 0.3_642, 0.4_553, 0.6_236, 0.5_797, 0.5_013, 0.4_343, 0.5_611, 0.4_831]				)

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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      """cpu"""  # ensure determinism for the device-dependent torch.Generator
							lowerCamelCase       =      self.get_dummy_components()
							lowerCamelCase       =      StableDiffusionInstructPixaPixPipeline(**_a				)
							lowerCamelCase       =      sd_pipe.to(_a				)
							sd_pipe.set_progress_bar_config(disable=_a				)

							lowerCamelCase       =      self.get_dummy_inputs(_a				)
							lowerCamelCase       =      [inputs["""prompt"""]] * 2

							lowerCamelCase       =      np.array(inputs["""image"""]				).astype(np.floataa				) / 255.0
							lowerCamelCase       =      torch.from_numpy(_a				).unsqueeze(0				).to(_a				)
							lowerCamelCase       =      image / 2 + 0.5
							lowerCamelCase       =      image.permute(0	,				3	,				1	,				2				)
							lowerCamelCase       =      image.repeat(2	,				1	,				1	,				1				)

							lowerCamelCase       =      sd_pipe(**_a				).images
							lowerCamelCase       =      image[-1, -3:, -3:, -1]

							assert image.shape == (2, 32, 32, 3)
							lowerCamelCase       =      np.array([0.5_812, 0.5_748, 0.5_222, 0.5_908, 0.5_695, 0.7_174, 0.6_804, 0.5_523, 0.5_579]				)

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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      """cpu"""  # ensure determinism for the device-dependent torch.Generator
							lowerCamelCase       =      self.get_dummy_components()
							lowerCamelCase       =      EulerAncestralDiscreteScheduler(
							    beta_start=0.00_085	,				beta_end=0.012	,				beta_schedule="""scaled_linear"""				)
							lowerCamelCase       =      StableDiffusionInstructPixaPixPipeline(**_a				)
							lowerCamelCase       =      sd_pipe.to(_a				)
							sd_pipe.set_progress_bar_config(disable=_a				)

							lowerCamelCase       =      self.get_dummy_inputs(_a				)
							lowerCamelCase       =      sd_pipe(**_a				).images
							lowerCamelCase       =      image[0, -3:, -3:, -1]

							lowerCamelCase       =      [round(_a	,				4				) for x in image_slice.flatten().tolist()]
							print(""",""".join([str(_a				) for x in slice]				)				)

							assert image.shape == (1, 32, 32, 3)
							lowerCamelCase       =      np.array([0.7_417, 0.3_842, 0.4_732, 0.5_776, 0.5_891, 0.5_139, 0.4_052, 0.5_673, 0.4_986]				)

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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							super().test_inference_batch_single_identical(expected_max_diff=3e-3				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self.get_dummy_components()
							lowerCamelCase       =      StableDiffusionInstructPixaPixPipeline(**_a				)
							lowerCamelCase       =      VaeImageProcessor(do_resize=_a	,				do_normalize=_a				)
							lowerCamelCase       =      pipe.to(_a				)
							pipe.set_progress_bar_config(disable=_a				)

							lowerCamelCase       =      pipe(**self.get_dummy_inputs_by_type(_a	,				input_image_type="""pt"""				)				)[0]

							lowerCamelCase       =      components["""vae"""]
							lowerCamelCase       =      self.get_dummy_inputs_by_type(_a	,				input_image_type="""pt"""				)

							for image_param in self.image_latents_params:
										if image_param in inputs.keys():
													lowerCamelCase       =      vae.encode(inputs[image_param]				).latent_dist.mode()

							lowerCamelCase       =      pipe(**_a				)[0]

							lowerCamelCase       =      np.abs(out - out_latents_inputs				).max()
							self.assertLess(_a	,				1e-4	,				"""passing latents as image input generate different result from passing image"""				)






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






				'''simple docstring'''





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()





				def      _lowerCAmelCase       ( self	,				_a=0				):





							"""simple docstring"""
							lowerCamelCase       =      torch.manual_seed(_a				)
							lowerCamelCase       =      load_image(
							    """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_pix2pix/example.jpg"""				)
							lowerCamelCase       =      {
							    """prompt""": """turn him into a cyborg""",
							    """image""": image,
							    """generator""": generator,
							    """num_inference_steps""": 3,
							    """guidance_scale""": 7.5,
							    """image_guidance_scale""": 1.0,
							    """output_type""": """numpy""",
							}
							return inputs





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      StableDiffusionInstructPixaPixPipeline.from_pretrained(
							    """timbrooks/instruct-pix2pix"""	,				safety_checker=_a				)
							pipe.to(_a				)
							pipe.set_progress_bar_config(disable=_a				)
							pipe.enable_attention_slicing()

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

							assert image.shape == (1, 512, 512, 3)
							lowerCamelCase       =      np.array([0.5_902, 0.6_015, 0.6_027, 0.5_983, 0.6_092, 0.6_061, 0.5_765, 0.5_785, 0.5_555]				)

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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      StableDiffusionInstructPixaPixPipeline.from_pretrained(
							    """timbrooks/instruct-pix2pix"""	,				safety_checker=_a				)
							lowerCamelCase       =      LMSDiscreteScheduler.from_config(pipe.scheduler.config				)
							pipe.to(_a				)
							pipe.set_progress_bar_config(disable=_a				)
							pipe.enable_attention_slicing()

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

							assert image.shape == (1, 512, 512, 3)
							lowerCamelCase       =      np.array([0.6_578, 0.6_817, 0.6_972, 0.6_761, 0.6_856, 0.6_916, 0.6_428, 0.6_516, 0.6_301]				)

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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      StableDiffusionInstructPixaPixPipeline.from_pretrained(
							    """timbrooks/instruct-pix2pix"""	,				safety_checker=_a				)
							lowerCamelCase       =      DDIMScheduler.from_config(pipe.scheduler.config				)
							pipe.to(_a				)
							pipe.set_progress_bar_config(disable=_a				)
							pipe.enable_attention_slicing()

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

							assert image.shape == (1, 512, 512, 3)
							lowerCamelCase       =      np.array([0.3_828, 0.3_834, 0.3_818, 0.3_792, 0.3_865, 0.3_752, 0.3_792, 0.3_847, 0.3_753]				)

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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      0

							def callback_fn(_a	,				_a	,				_a				) -> None:
										lowerCamelCase       =      True
										nonlocal number_of_steps
										number_of_steps += 1
										if step == 1:
													lowerCamelCase       =      latents.detach().cpu().numpy()
													assert latents.shape == (1, 4, 64, 64)
													lowerCamelCase       =      latents[0, -3:, -3:, -1]
													lowerCamelCase       =      np.array([-0.2_463, -0.4_644, -0.9_756, 1.5_176, 1.4_414, 0.7_866, 0.9_897, 0.8_521, 0.7_983]				)

													assert np.abs(latents_slice.flatten() - expected_slice				).max() < 5e-2
										elif step == 2:
													lowerCamelCase       =      latents.detach().cpu().numpy()
													assert latents.shape == (1, 4, 64, 64)
													lowerCamelCase       =      latents[0, -3:, -3:, -1]
													lowerCamelCase       =      np.array([-0.2_644, -0.4_626, -0.9_653, 1.5_176, 1.4_551, 0.7_686, 0.9_805, 0.8_452, 0.8_115]				)

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

							lowerCamelCase       =      False

							lowerCamelCase       =      StableDiffusionInstructPixaPixPipeline.from_pretrained(
							    """timbrooks/instruct-pix2pix"""	,				safety_checker=_a	,				torch_dtype=torch.floataa				)
							lowerCamelCase       =      pipe.to(_a				)
							pipe.set_progress_bar_config(disable=_a				)
							pipe.enable_attention_slicing()

							lowerCamelCase       =      self.get_inputs()
							pipe(**_a	,				callback=_a	,				callback_steps=1				)
							assert callback_fn.has_been_called
							assert number_of_steps == 3





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							torch.cuda.empty_cache()
							torch.cuda.reset_max_memory_allocated()
							torch.cuda.reset_peak_memory_stats()

							lowerCamelCase       =      StableDiffusionInstructPixaPixPipeline.from_pretrained(
							    """timbrooks/instruct-pix2pix"""	,				safety_checker=_a	,				torch_dtype=torch.floataa				)
							lowerCamelCase       =      pipe.to(_a				)
							pipe.set_progress_bar_config(disable=_a				)
							pipe.enable_attention_slicing(1				)
							pipe.enable_sequential_cpu_offload()

							lowerCamelCase       =      self.get_inputs()
							lowerCamelCase       =      pipe(**_a				)

							lowerCamelCase       =      torch.cuda.max_memory_allocated()
							# make sure that less than 2.2 GB is allocated
							assert mem_bytes < 2.2 * 10**9





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self.get_inputs()
							# resize to resolution that is divisible by 8 but not 16 or 32
							lowerCamelCase       =      inputs["""image"""].resize((504, 504)				)

							lowerCamelCase       =      """timbrooks/instruct-pix2pix"""
							lowerCamelCase       =      StableDiffusionInstructPixaPixPipeline.from_pretrained(
							    _a	,				safety_checker=_a	,				)
							pipe.to(_a				)
							pipe.set_progress_bar_config(disable=_a				)
							pipe.enable_attention_slicing()

							lowerCamelCase       =      pipe(**_a				)
							lowerCamelCase       =      output.images[0]

							lowerCamelCase       =      image[255:258, 383:386, -1]

							assert image.shape == (504, 504, 3)
							lowerCamelCase       =      np.array([0.2_726, 0.2_529, 0.2_664, 0.2_655, 0.2_641, 0.2_642, 0.2_591, 0.2_649, 0.2_590]				)

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









"""simple docstring"""




import argparse

import torch

from diffusers.pipelines.stable_diffusion.convert_from_ckpt import download_from_original_stable_diffusion_ckpt


if __name__ == "__main__":
					lowerCAmelCase							:    Dict  							=							argparse.ArgumentParser()

					parser.add_argument(
					    """--checkpoint_path""", default=None, type=str, required=True, help="""Path to the checkpoint to convert."""
					)
					# !wget https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml
					parser.add_argument(
					    """--original_config_file""",
					    default=None,
					    type=str,
					    help="""The YAML config file corresponding to the original architecture.""",
					)
					parser.add_argument(
					    """--num_in_channels""",
					    default=None,
					    type=int,
					    help="""The number of input channels. If `None` number of input channels will be automatically inferred.""",
					)
					parser.add_argument(
					    """--scheduler_type""",
					    default="""pndm""",
					    type=str,
					    help="""Type of scheduler to use. Should be one of ['pndm', 'lms', 'ddim', 'euler', 'euler-ancestral', 'dpm']""",
					)
					parser.add_argument(
					    """--pipeline_type""",
					    default=None,
					    type=str,
					    help=(
					        """The pipeline type. One of 'FrozenOpenCLIPEmbedder', 'FrozenCLIPEmbedder', 'PaintByExample'"""
					        """. If `None` pipeline will be automatically inferred."""
					    ),
					)
					parser.add_argument(
					    """--image_size""",
					    default=None,
					    type=int,
					    help=(
					        """The image size that the model was trained on. Use 512 for Stable Diffusion v1.X and Stable Siffusion v2"""
					        """ Base. Use 768 for Stable Diffusion v2."""
					    ),
					)
					parser.add_argument(
					    """--prediction_type""",
					    default=None,
					    type=str,
					    help=(
					        """The prediction type that the model was trained on. Use 'epsilon' for Stable Diffusion v1.X and Stable"""
					        """ Diffusion v2 Base. Use 'v_prediction' for Stable Diffusion v2."""
					    ),
					)
					parser.add_argument(
					    """--extract_ema""",
					    action="""store_true""",
					    help=(
					        """Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights"""
					        """ or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield"""
					        """ higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning."""
					    ),
					)
					parser.add_argument(
					    """--upcast_attention""",
					    action="""store_true""",
					    help=(
					        """Whether the attention computation should always be upcasted. This is necessary when running stable"""
					        """ diffusion 2.1."""
					    ),
					)
					parser.add_argument(
					    """--from_safetensors""",
					    action="""store_true""",
					    help="""If `--checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch.""",
					)
					parser.add_argument(
					    """--to_safetensors""",
					    action="""store_true""",
					    help="""Whether to store pipeline in safetensors format or not.""",
					)
					parser.add_argument("""--dump_path""", default=None, type=str, required=True, help="""Path to the output model.""")
					parser.add_argument("""--device""", type=str, help="""Device to use (e.g. cpu, cuda:0, cuda:1, etc.)""")
					parser.add_argument(
					    """--stable_unclip""",
					    type=str,
					    default=None,
					    required=False,
					    help="""Set if this is a stable unCLIP model. One of 'txt2img' or 'img2img'.""",
					)
					parser.add_argument(
					    """--stable_unclip_prior""",
					    type=str,
					    default=None,
					    required=False,
					    help="""Set if this is a stable unCLIP txt2img model. Selects which prior to use. If `--stable_unclip` is set to `txt2img`, the karlo prior (https://huggingface.co/kakaobrain/karlo-v1-alpha/tree/main/prior) is selected by default.""",
					)
					parser.add_argument(
					    """--clip_stats_path""",
					    type=str,
					    help="""Path to the clip stats file. Only required if the stable unclip model's config specifies `model.params.noise_aug_config.params.clip_stats_path`.""",
					    required=False,
					)
					parser.add_argument(
					    """--controlnet""", action="""store_true""", default=None, help="""Set flag if this is a controlnet checkpoint."""
					)
					parser.add_argument("""--half""", action="""store_true""", help="""Save weights in half precision.""")
					parser.add_argument(
					    """--vae_path""",
					    type=str,
					    default=None,
					    required=False,
					    help="""Set to a path, hub id to an already converted vae to not convert it again.""",
					)
					lowerCAmelCase							:    Dict  							=							parser.parse_args()

					lowerCAmelCase							:    str  							=							download_from_original_stable_diffusion_ckpt(
					    checkpoint_path=args.checkpoint_path,
					    original_config_file=args.original_config_file,
					    image_size=args.image_size,
					    prediction_type=args.prediction_type,
					    model_type=args.pipeline_type,
					    extract_ema=args.extract_ema,
					    scheduler_type=args.scheduler_type,
					    num_in_channels=args.num_in_channels,
					    upcast_attention=args.upcast_attention,
					    from_safetensors=args.from_safetensors,
					    device=args.device,
					    stable_unclip=args.stable_unclip,
					    stable_unclip_prior=args.stable_unclip_prior,
					    clip_stats_path=args.clip_stats_path,
					    controlnet=args.controlnet,
					    vae_path=args.vae_path,
					)

					if args.half:
										pipe.to(torch_dtype=torch.floataa)

					if args.controlnet:
										# only save the controlnet model
										pipe.controlnet.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)
					else:
										pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)



