


import random
import unittest

import numpy as np

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionImgaImgPipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
	import onnxruntime as ort



class 							A    (       lowerCamelCase__ ,  unittest.TestCase							):
					'''simple docstring'''






					A__			= 'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'

					def   lowerCamelCase__     (self    :						str      ,					_UpperCAmelCase    :						Tuple=0		)						->   Union[str, Any]:




							"""simple docstring"""


							lowercase__     			=		floats_tensor((1, 3, 128, 128)      ,					rng=random.Random(lowercase_		)		)
							lowercase__     			=		np.random.RandomState(lowercase_		)
							lowercase__     			=		{
							    """prompt""": """A painting of a squirrel eating a burger""",
							    """image""": image,
							    """generator""": generator,
							    """num_inference_steps""": 3,
							    """strength""": 0.75,
							    """guidance_scale""": 7.5,
							    """output_type""": """numpy""",
							}
							return inputs

					def   lowerCamelCase__     (self    :						Optional[int]		)						->   Optional[int]:




							"""simple docstring"""


							lowercase__     			=		OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint      ,					provider="""CPUExecutionProvider"""		)
							pipe.set_progress_bar_config(disable=lowercase_		)

							lowercase__     			=		self.get_dummy_inputs()
							lowercase__     			=		pipe(**lowercase_		).images
							lowercase__     			=		image[0, -3:, -3:, -1].flatten()

							assert image.shape == (1, 128, 128, 3)
							lowercase__     			=		np.array([0.69_643, 0.58_484, 0.50_314, 0.58_760, 0.55_368, 0.59_643, 0.51_529, 0.41_217, 0.49_087]		)
							assert np.abs(image_slice - expected_slice		).max() < 1E-1

					def   lowerCamelCase__     (self    :						Any		)						->   Any:




							"""simple docstring"""


							lowercase__     			=		OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint      ,					provider="""CPUExecutionProvider"""		)
							lowercase__     			=		PNDMScheduler.from_config(pipe.scheduler.config      ,					skip_prk_steps=lowercase_		)
							pipe.set_progress_bar_config(disable=lowercase_		)

							lowercase__     			=		self.get_dummy_inputs()
							lowercase__     			=		pipe(**lowercase_		).images
							lowercase__     			=		image[0, -3:, -3:, -1]

							assert image.shape == (1, 128, 128, 3)
							lowercase__     			=		np.array([0.61_737, 0.54_642, 0.53_183, 0.54_465, 0.52_742, 0.60_525, 0.49_969, 0.40_655, 0.48_154]		)

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

					def   lowerCamelCase__     (self    :						List[str]		)						->   Optional[Any]:




							"""simple docstring"""


							lowercase__     			=		OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint      ,					provider="""CPUExecutionProvider"""		)
							lowercase__     			=		LMSDiscreteScheduler.from_config(pipe.scheduler.config		)
							pipe.set_progress_bar_config(disable=lowercase_		)

							# warmup pass to apply optimizations
							lowercase__     			=		pipe(**self.get_dummy_inputs()		)

							lowercase__     			=		self.get_dummy_inputs()
							lowercase__     			=		pipe(**lowercase_		).images
							lowercase__     			=		image[0, -3:, -3:, -1]

							assert image.shape == (1, 128, 128, 3)
							lowercase__     			=		np.array([0.52_761, 0.59_977, 0.49_033, 0.49_619, 0.54_282, 0.50_311, 0.47_600, 0.40_918, 0.45_203]		)

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

					def   lowerCamelCase__     (self    :						Dict		)						->   Union[str, Any]:




							"""simple docstring"""


							lowercase__     			=		OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint      ,					provider="""CPUExecutionProvider"""		)
							lowercase__     			=		EulerDiscreteScheduler.from_config(pipe.scheduler.config		)
							pipe.set_progress_bar_config(disable=lowercase_		)

							lowercase__     			=		self.get_dummy_inputs()
							lowercase__     			=		pipe(**lowercase_		).images
							lowercase__     			=		image[0, -3:, -3:, -1]

							assert image.shape == (1, 128, 128, 3)
							lowercase__     			=		np.array([0.52_911, 0.60_004, 0.49_229, 0.49_805, 0.54_502, 0.50_680, 0.47_777, 0.41_028, 0.45_304]		)

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

					def   lowerCamelCase__     (self    :						str		)						->   Optional[Any]:




							"""simple docstring"""


							lowercase__     			=		OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint      ,					provider="""CPUExecutionProvider"""		)
							lowercase__     			=		EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config		)
							pipe.set_progress_bar_config(disable=lowercase_		)

							lowercase__     			=		self.get_dummy_inputs()
							lowercase__     			=		pipe(**lowercase_		).images
							lowercase__     			=		image[0, -3:, -3:, -1]

							assert image.shape == (1, 128, 128, 3)
							lowercase__     			=		np.array([0.52_911, 0.60_004, 0.49_229, 0.49_805, 0.54_502, 0.50_680, 0.47_777, 0.41_028, 0.45_304]		)

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


					def   lowerCamelCase__     (self    :						Any		)						->   str:




							"""simple docstring"""


							lowercase__     			=		OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint      ,					provider="""CPUExecutionProvider"""		)
							lowercase__     			=		DPMSolverMultistepScheduler.from_config(pipe.scheduler.config		)
							pipe.set_progress_bar_config(disable=lowercase_		)

							lowercase__     			=		self.get_dummy_inputs()
							lowercase__     			=		pipe(**lowercase_		).images
							lowercase__     			=		image[0, -3:, -3:, -1]

							assert image.shape == (1, 128, 128, 3)
							lowercase__     			=		np.array([0.65_331, 0.58_277, 0.48_204, 0.56_059, 0.53_665, 0.56_235, 0.50_969, 0.40_009, 0.46_552]		)

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






@nightly
@require_onnxruntime
@require_torch_gpu
class 							A    (       unittest.TestCase							):
					'''simple docstring'''






					@property
					def   lowerCamelCase__     (self    :						Tuple		)						->   Any:




							"""simple docstring"""


							return (
							    "CUDAExecutionProvider",
							    {
							        "gpu_mem_limit": "15000000000",  # 15GB
							        "arena_extend_strategy": "kSameAsRequested",
							    },
							)

					@property
					def   lowerCamelCase__     (self    :						Union[str, Any]		)						->   Tuple:




							"""simple docstring"""


							lowercase__     			=		ort.SessionOptions()
							lowercase__     			=		False
							return options

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




							"""simple docstring"""


							lowercase__     			=		load_image(
							    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
							    """/img2img/sketch-mountains-input.jpg"""		)
							lowercase__     			=		init_image.resize((768, 512)		)
							# using the PNDM scheduler by default
							lowercase__     			=		OnnxStableDiffusionImgaImgPipeline.from_pretrained(
							    """CompVis/stable-diffusion-v1-4"""      ,					revision="""onnx"""      ,					safety_checker=lowercase_      ,					feature_extractor=lowercase_      ,					provider=self.gpu_provider      ,					sess_options=self.gpu_options      ,					)
							pipe.set_progress_bar_config(disable=lowercase_		)

							lowercase__     			=		"""A fantasy landscape, trending on artstation"""

							lowercase__     			=		np.random.RandomState(0		)
							lowercase__     			=		pipe(
							    prompt=lowercase_      ,					image=lowercase_      ,					strength=0.75      ,					guidance_scale=7.5      ,					num_inference_steps=10      ,					generator=lowercase_      ,					output_type="""np"""      ,					)
							lowercase__     			=		output.images
							lowercase__     			=		images[0, 255:258, 383:386, -1]

							assert images.shape == (1, 512, 768, 3)
							lowercase__     			=		np.array([0.4_909, 0.5_059, 0.5_372, 0.4_623, 0.4_876, 0.5_049, 0.4_820, 0.4_956, 0.5_019]		)
							# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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


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




							"""simple docstring"""


							lowercase__     			=		load_image(
							    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
							    """/img2img/sketch-mountains-input.jpg"""		)
							lowercase__     			=		init_image.resize((768, 512)		)
							lowercase__     			=		LMSDiscreteScheduler.from_pretrained(
							    """runwayml/stable-diffusion-v1-5"""      ,					subfolder="""scheduler"""      ,					revision="""onnx"""		)
							lowercase__     			=		OnnxStableDiffusionImgaImgPipeline.from_pretrained(
							    """runwayml/stable-diffusion-v1-5"""      ,					revision="""onnx"""      ,					scheduler=lowercase_      ,					safety_checker=lowercase_      ,					feature_extractor=lowercase_      ,					provider=self.gpu_provider      ,					sess_options=self.gpu_options      ,					)
							pipe.set_progress_bar_config(disable=lowercase_		)

							lowercase__     			=		"""A fantasy landscape, trending on artstation"""

							lowercase__     			=		np.random.RandomState(0		)
							lowercase__     			=		pipe(
							    prompt=lowercase_      ,					image=lowercase_      ,					strength=0.75      ,					guidance_scale=7.5      ,					num_inference_steps=20      ,					generator=lowercase_      ,					output_type="""np"""      ,					)
							lowercase__     			=		output.images
							lowercase__     			=		images[0, 255:258, 383:386, -1]

							assert images.shape == (1, 512, 768, 3)
							lowercase__     			=		np.array([0.8_043, 0.926, 0.9_581, 0.8_119, 0.8_954, 0.913, 0.7_209, 0.7_463, 0.7_431]		)
							# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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




from collections import defaultdict
from typing import Optional

from ..image_utils import load_image
from ..utils import (
    add_end_docstrings,
    is_torch_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, ChunkPipeline


if is_torch_available():
       import torch

       from ..models.auto.modeling_auto import MODEL_FOR_MASK_GENERATION_MAPPING

lowerCamelCase           =					logging.get_logger(__name__)







@add_end_docstrings(lowerCamelCase__							)
class 						__magic_name__			(							lowerCamelCase__							):

   '''simple docstring'''




   def __init__(				self,    **lowercase_       )							->    Optional[int]:






     """simple docstring"""






     super().__init__(**lowercase_       )
     requires_backends(self,    '''vision'''       )
     requires_backends(self,    '''torch'''       )

     if self.framework != "pt":
       raise ValueError(F"""The {self.__class__} is only available in PyTorch."""       )

     self.check_model_type(lowercase_       )



   def  _UpperCAmelCase (				self,    **lowercase_       )							->    Optional[Any]:






     """simple docstring"""






     a__  ={}
     a__  ={}
     a__  ={}
     # preprocess args
     if "points_per_batch" in kwargs:
       a__  =kwargs['''points_per_batch''']
     if "points_per_crop" in kwargs:
       a__  =kwargs['''points_per_crop''']
     if "crops_n_layers" in kwargs:
       a__  =kwargs['''crops_n_layers''']
     if "crop_overlap_ratio" in kwargs:
       a__  =kwargs['''crop_overlap_ratio''']
     if "crop_n_points_downscale_factor" in kwargs:
       a__  =kwargs['''crop_n_points_downscale_factor''']
     # postprocess args
     if "pred_iou_thresh" in kwargs:
       a__  =kwargs['''pred_iou_thresh''']
     if "stability_score_offset" in kwargs:
       a__  =kwargs['''stability_score_offset''']
     if "mask_threshold" in kwargs:
       a__  =kwargs['''mask_threshold''']
     if "stability_score_thresh" in kwargs:
       a__  =kwargs['''stability_score_thresh''']
     if "crops_nms_thresh" in kwargs:
       a__  =kwargs['''crops_nms_thresh''']
     if "output_rle_mask" in kwargs:
       a__  =kwargs['''output_rle_mask''']
     if "output_bboxes_mask" in kwargs:
       a__  =kwargs['''output_bboxes_mask''']
     return preprocess_kwargs, forward_params, postprocess_kwargs



   def __call__(				self,    lowercase_,    *lowercase_,    lowercase_=None,    lowercase_=None,    **lowercase_       )							->    List[Any]:






     """simple docstring"""






     return super().__call__(lowercase_,    *lowercase_,    num_workers=lowercase_,    batch_size=lowercase_,    **lowercase_       )



   def  _UpperCAmelCase (				self,    lowercase_,    lowercase_=64,    lowercase_ = 0,    lowercase_ = 512 / 1500,    lowercase_ = 32,    lowercase_ = 1,    )							->    Any:






     """simple docstring"""






     a__  =load_image(lowercase_       )
     a__  =self.image_processor.size['''longest_edge''']
     a__,      a__,      a__,      a__  =self.image_processor.generate_crop_boxes(
         lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_       )
     a__  =self.image_processor(images=lowercase_,    return_tensors='''pt'''       )

     with self.device_placement():
       if self.framework == "pt":
         a__  =self.get_inference_context()
         with inference_context():
           a__  =self._ensure_tensor_on_device(lowercase_,    device=self.device       )
           a__  =self.model.get_image_embeddings(model_inputs.pop('''pixel_values'''       )       )
           a__  =image_embeddings

     a__  =grid_points.shape[1]
     a__  =points_per_batch if points_per_batch is not None else n_points

     if points_per_batch <= 0:
       raise ValueError(
           '''Cannot have points_per_batch<=0. Must be >=1 to returned batched outputs. '''
           '''To return all points at once, set points_per_batch to None'''       )

     for i in range(0,    lowercase_,    lowercase_       ):
       a__  =grid_points[:, i : i + points_per_batch, :, :]
       a__  =input_labels[:, i : i + points_per_batch]
       a__  =i == n_points - points_per_batch
       yield {
           "input_points": batched_points,
           "input_labels": labels,
           "input_boxes": crop_boxes,
           "is_last": is_last,
           **model_inputs,
       }



   def  _UpperCAmelCase (				self,    lowercase_,    lowercase_=0.88,    lowercase_=0.95,    lowercase_=0,    lowercase_=1,    )							->    int:






     """simple docstring"""






     a__  =model_inputs.pop('''input_boxes'''       )
     a__  =model_inputs.pop('''is_last'''       )
     a__  =model_inputs.pop('''original_sizes'''       ).tolist()
     a__  =model_inputs.pop('''reshaped_input_sizes'''       ).tolist()

     a__  =self.model(**lowercase_       )

     # post processing happens here in order to avoid CPU GPU copies of ALL the masks
     a__  =model_outputs['''pred_masks''']
     a__  =self.image_processor.post_process_masks(
         lowercase_,    lowercase_,    lowercase_,    lowercase_,    binarize=lowercase_       )
     a__  =model_outputs['''iou_scores''']
     a__,      a__,      a__  =self.image_processor.filter_masks(
         masks[0],    iou_scores[0],    original_sizes[0],    input_boxes[0],    lowercase_,    lowercase_,    lowercase_,    lowercase_,    )
     return {
         "masks": masks,
         "is_last": is_last,
         "boxes": boxes,
         "iou_scores": iou_scores,
     }



   def  _UpperCAmelCase (				self,    lowercase_,    lowercase_=False,    lowercase_=False,    lowercase_=0.7,    )							->    Any:






     """simple docstring"""






     a__  =[]
     a__  =[]
     a__  =[]
     for model_output in model_outputs:
       all_scores.append(model_output.pop('''iou_scores'''       )       )
       all_masks.extend(model_output.pop('''masks'''       )       )
       all_boxes.append(model_output.pop('''boxes'''       )       )

     a__  =torch.cat(lowercase_       )
     a__  =torch.cat(lowercase_       )
     a__,      a__,      a__,      a__  =self.image_processor.post_process_for_mask_generation(
         lowercase_,    lowercase_,    lowercase_,    lowercase_       )

     a__  =defaultdict(lowercase_       )
     for output in model_outputs:
       for k, v in output.items():
         extra[k].append(lowercase_       )

     a__  ={}
     if output_rle_mask:
       a__  =rle_mask

     if output_bboxes_mask:
       a__  =bounding_boxes

     return {"masks": output_masks, "scores": iou_scores, **optional, **extra}
