
"""simple docstring"""
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 			UpperCamelCase      (							_UpperCAmelCase      ,       unittest.TestCase  ):
     A__      =			"""hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"""
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__=0				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Any							    = floats_tensor((1, 3, 128, 128)  ,     rng=random.Random(lowercase__				)				)
       _SCREAMING_SNAKE_CASE   :				int							    = np.random.RandomState(lowercase__				)
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = {
           "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 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				List[Any]							    = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint  ,     provider="CPUExecutionProvider"				)
       pipe.set_progress_bar_config(disable=lowercase__				)

       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = self.get_dummy_inputs()
       _SCREAMING_SNAKE_CASE   :				str							    = pipe(**lowercase__				).images
       _SCREAMING_SNAKE_CASE   :				Dict							    = image[0, -3:, -3:, -1].flatten()

       assert image.shape == (1, 128, 128, 3)
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = 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 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				List[str]							    = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint  ,     provider="CPUExecutionProvider"				)
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = PNDMScheduler.from_config(pipe.scheduler.config  ,     skip_prk_steps=lowercase__				)
       pipe.set_progress_bar_config(disable=lowercase__				)

       _SCREAMING_SNAKE_CASE   :				Tuple							    = self.get_dummy_inputs()
       _SCREAMING_SNAKE_CASE   :				Any							    = pipe(**lowercase__				).images
       _SCREAMING_SNAKE_CASE   :				int							    = image[0, -3:, -3:, -1]

       assert image.shape == (1, 128, 128, 3)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = 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 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint  ,     provider="CPUExecutionProvider"				)
       _SCREAMING_SNAKE_CASE   :				Any							    = LMSDiscreteScheduler.from_config(pipe.scheduler.config				)
       pipe.set_progress_bar_config(disable=lowercase__				)

       # warmup pass to apply optimizations
       _SCREAMING_SNAKE_CASE   :				List[str]							    = pipe(**self.get_dummy_inputs()				)

       _SCREAMING_SNAKE_CASE   :				Any							    = self.get_dummy_inputs()
       _SCREAMING_SNAKE_CASE   :				Tuple							    = pipe(**lowercase__				).images
       _SCREAMING_SNAKE_CASE   :				str							    = image[0, -3:, -3:, -1]

       assert image.shape == (1, 128, 128, 3)
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = 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 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint  ,     provider="CPUExecutionProvider"				)
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = EulerDiscreteScheduler.from_config(pipe.scheduler.config				)
       pipe.set_progress_bar_config(disable=lowercase__				)

       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = self.get_dummy_inputs()
       _SCREAMING_SNAKE_CASE   :				Any							    = pipe(**lowercase__				).images
       _SCREAMING_SNAKE_CASE   :				Dict							    = image[0, -3:, -3:, -1]

       assert image.shape == (1, 128, 128, 3)
       _SCREAMING_SNAKE_CASE   :				Tuple							    = 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 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


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

       _SCREAMING_SNAKE_CASE   :				List[str]							    = self.get_dummy_inputs()
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = pipe(**lowercase__				).images
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = image[0, -3:, -3:, -1]

       assert image.shape == (1, 128, 128, 3)
       _SCREAMING_SNAKE_CASE   :				int							    = 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 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				str							    = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint  ,     provider="CPUExecutionProvider"				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config				)
       pipe.set_progress_bar_config(disable=lowercase__				)

       _SCREAMING_SNAKE_CASE   :				Tuple							    = self.get_dummy_inputs()
       _SCREAMING_SNAKE_CASE   :				Any							    = pipe(**lowercase__				).images
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = image[0, -3:, -3:, -1]

       assert image.shape == (1, 128, 128, 3)
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = 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 			UpperCamelCase      (							unittest.TestCase  ):
     @property
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       return (
           "CUDAExecutionProvider",
           {
               "gpu_mem_limit": "15000000000",  # 15GB
               "arena_extend_strategy": "kSameAsRequested",
           },
       )
     @property
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = ort.SessionOptions()
       _SCREAMING_SNAKE_CASE   :				Tuple							    = False
       return options
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Dict							    = load_image(
           "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
           "/img2img/sketch-mountains-input.jpg"				)
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = init_image.resize((768, 512)				)
       # using the PNDM scheduler by default
       _SCREAMING_SNAKE_CASE   :				Dict							    = 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__				)

       _SCREAMING_SNAKE_CASE   :				int							    = "A fantasy landscape, trending on artstation"

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

       assert images.shape == (1, 512, 768, 3)
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = 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 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Dict							    = load_image(
           "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
           "/img2img/sketch-mountains-input.jpg"				)
       _SCREAMING_SNAKE_CASE   :				Any							    = init_image.resize((768, 512)				)
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = LMSDiscreteScheduler.from_pretrained(
           "runwayml/stable-diffusion-v1-5"  ,     subfolder="scheduler"  ,     revision="onnx"				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = 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__				)

       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = "A fantasy landscape, trending on artstation"

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

       assert images.shape == (1, 512, 768, 3)
       _SCREAMING_SNAKE_CASE   :				str							    = 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





"""simple docstring"""
import re
import warnings
from contextlib import contextmanager

from ...processing_utils import ProcessorMixin
class 			UpperCamelCase      (							__SCREAMING_SNAKE_CASE  ):
     A__      =			["""image_processor""", """tokenizer"""]
     A__      =			"""AutoImageProcessor"""
     A__      =			"""AutoTokenizer"""
     def __init__(     self  ,     snake_case__=None  ,     snake_case__=None  ,     **snake_case__				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				int							    = None
       if "feature_extractor" in kwargs:
         warnings.warn(
             "The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"
             " instead."  ,     snake_case__  ,     )
         _SCREAMING_SNAKE_CASE   :				Any							    = kwargs.pop("feature_extractor"				)

       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = image_processor if image_processor is not None else feature_extractor
       if image_processor is None:
         raise ValueError("You need to specify an `image_processor`."				)
       if tokenizer is None:
         raise ValueError("You need to specify a `tokenizer`."				)

       super().__init__(snake_case__  ,     snake_case__				)
       _SCREAMING_SNAKE_CASE   :				Any							    = self.image_processor
       _SCREAMING_SNAKE_CASE   :				Any							    = False
     def __call__(     self  ,     *snake_case__  ,     **snake_case__				):




       """simple docstring"""


       if self._in_target_context_manager:
         return self.current_processor(*snake_case__  ,     **snake_case__				)

       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = kwargs.pop("images"  ,     snake_case__				)
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = kwargs.pop("text"  ,     snake_case__				)
       if len(snake_case__				) > 0:
         _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = args[0]
         _SCREAMING_SNAKE_CASE   :				List[str]							    = args[1:]

       if images is None and text is None:
         raise ValueError("You need to specify either an `images` or `text` input to process."				)

       if images is not None:
         _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = self.image_processor(snake_case__  ,     *snake_case__  ,     **snake_case__				)
       if text is not None:
         _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = self.tokenizer(snake_case__  ,     **snake_case__				)

       if text is None:
         return inputs
       elif images is None:
         return encodings
       else:
         _SCREAMING_SNAKE_CASE   :				List[Any]							    = encodings["input_ids"]
         return inputs
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     *snake_case__  ,     **snake_case__				):




       """simple docstring"""


       return self.tokenizer.batch_decode(*snake_case__  ,     **snake_case__				)
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     *snake_case__  ,     **snake_case__				):




       """simple docstring"""


       return self.tokenizer.decode(*snake_case__  ,     **snake_case__				)
     @contextmanager
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       warnings.warn(
           "`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your "
           "labels by using the argument `text` of the regular `__call__` method (either in the same call as "
           "your images inputs, or in a separate call."				)
       _SCREAMING_SNAKE_CASE   :				Dict							    = True
       _SCREAMING_SNAKE_CASE   :				str							    = self.tokenizer
       yield
       _SCREAMING_SNAKE_CASE   :				str							    = self.image_processor
       _SCREAMING_SNAKE_CASE   :				int							    = False
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__=False  ,     snake_case__=None				):




       """simple docstring"""


       if added_vocab is None:
         _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = self.tokenizer.get_added_vocab()

       _SCREAMING_SNAKE_CASE   :				int							    = {}

       while tokens:
         _SCREAMING_SNAKE_CASE   :				List[Any]							    = re.search(r"<s_(.*?)>"  ,     snake_case__  ,     re.IGNORECASE				)
         if start_token is None:
           break
         _SCREAMING_SNAKE_CASE   :				List[Any]							    = start_token.group(1				)
         _SCREAMING_SNAKE_CASE   :				Any							    = re.search(rF'''</s_{key}>'''  ,     snake_case__  ,     re.IGNORECASE				)
         _SCREAMING_SNAKE_CASE   :				Tuple							    = start_token.group()
         if end_token is None:
           _SCREAMING_SNAKE_CASE   :				Optional[int]							    = tokens.replace(snake_case__  ,     ""				)
         else:
           _SCREAMING_SNAKE_CASE   :				Any							    = end_token.group()
           _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = re.escape(snake_case__				)
           _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = re.escape(snake_case__				)
           _SCREAMING_SNAKE_CASE   :				str							    = re.search(F'''{start_token_escaped}(.*?){end_token_escaped}'''  ,     snake_case__  ,     re.IGNORECASE				)
           if content is not None:
             _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = content.group(1				).strip()
             if r"<s_" in content and r"</s_" in content:  # non-leaf node
               _SCREAMING_SNAKE_CASE   :				Dict							    = self.tokenajson(snake_case__  ,     is_inner_value=snake_case__  ,     added_vocab=snake_case__				)
               if value:
                 if len(snake_case__				) == 1:
                   _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = value[0]
                 _SCREAMING_SNAKE_CASE   :				Dict							    = value
             else:  # leaf nodes
               _SCREAMING_SNAKE_CASE   :				Dict							    = []
               for leaf in content.split(r"<sep/>"				):
                 _SCREAMING_SNAKE_CASE   :				int							    = leaf.strip()
                 if leaf in added_vocab and leaf[0] == "<" and leaf[-2:] == "/>":
                   _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = leaf[1:-2]  # for categorical special tokens
                 output[key].append(snake_case__				)
               if len(output[key]				) == 1:
                 _SCREAMING_SNAKE_CASE   :				Dict							    = output[key][0]

           _SCREAMING_SNAKE_CASE   :				List[Any]							    = tokens[tokens.find(snake_case__				) + len(snake_case__				) :].strip()
           if tokens[:6] == r"<sep/>":  # non-leaf nodes
             return [output] + self.tokenajson(tokens[6:]  ,     is_inner_value=snake_case__  ,     added_vocab=snake_case__				)

       if len(snake_case__				):
         return [output] if is_inner_value else output
       else:
         return [] if is_inner_value else {"text_sequence": tokens}
     @property
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       warnings.warn(
           "`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."  ,     snake_case__  ,     )
       return self.image_processor_class




     @property
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       warnings.warn(
           "`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."  ,     snake_case__  ,     )
       return self.image_processor



