




from tempfile import TemporaryDirectory
from unittest import TestCase
from unittest.mock import MagicMock, patch

from transformers import AutoModel, TFAutoModel
from transformers.onnx import FeaturesManager
from transformers.testing_utils import SMALL_MODEL_IDENTIFIER, require_tf, require_torch






@require_torch
@require_tf
class  _lowerCAmelCase     (  _lowercase	):

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Union[str, Any]		       =     SMALL_MODEL_IDENTIFIER
        lowerCAmelCase__				:  List[str]		       =     '''pt'''
        lowerCAmelCase__				:  int		       =     '''tf'''

      def 	__magic_name__(   self      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  Optional[Any]		       =     AutoModel.from_pretrained(self.test_model		)
        model_pt.save_pretrained(__UpperCAmelCase		)

      def 	__magic_name__(   self      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  str		       =     TFAutoModel.from_pretrained(self.test_model      ,				from_pt=__UpperCAmelCase		)
        model_tf.save_pretrained(__UpperCAmelCase		)

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  int		       =     '''mock_framework'''

        # Framework provided - return whatever the user provides
        lowerCAmelCase__				:  Tuple		       =     FeaturesManager.determine_framework(self.test_model      ,				__UpperCAmelCase		)
        self.assertEqual(__UpperCAmelCase      ,				__UpperCAmelCase		)

        # Local checkpoint and framework provided - return provided framework
        # PyTorch checkpoint
        with TemporaryDirectory() as local_pt_ckpt:
          self._setup_pt_ckpt(__UpperCAmelCase		)
          lowerCAmelCase__				:  int		       =     FeaturesManager.determine_framework(__UpperCAmelCase      ,				__UpperCAmelCase		)
          self.assertEqual(__UpperCAmelCase      ,				__UpperCAmelCase		)

        # TensorFlow checkpoint
        with TemporaryDirectory() as local_tf_ckpt:
          self._setup_tf_ckpt(__UpperCAmelCase		)
          lowerCAmelCase__				:  List[str]		       =     FeaturesManager.determine_framework(__UpperCAmelCase      ,				__UpperCAmelCase		)
          self.assertEqual(__UpperCAmelCase      ,				__UpperCAmelCase		)

      def 	__magic_name__(   self		):
        # PyTorch checkpoint
        with TemporaryDirectory() as local_pt_ckpt:
          self._setup_pt_ckpt(__UpperCAmelCase		)
          lowerCAmelCase__				:  Tuple		       =     FeaturesManager.determine_framework(__UpperCAmelCase		)
          self.assertEqual(__UpperCAmelCase      ,				self.framework_pt		)

        # TensorFlow checkpoint
        with TemporaryDirectory() as local_tf_ckpt:
          self._setup_tf_ckpt(__UpperCAmelCase		)
          lowerCAmelCase__				:  int		       =     FeaturesManager.determine_framework(__UpperCAmelCase		)
          self.assertEqual(__UpperCAmelCase      ,				self.framework_tf		)

        # Invalid local checkpoint
        with TemporaryDirectory() as local_invalid_ckpt:
          with self.assertRaises(__UpperCAmelCase		):
            lowerCAmelCase__				:  Tuple		       =     FeaturesManager.determine_framework(__UpperCAmelCase		)







      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  List[Any]		       =     MagicMock(return_value=__UpperCAmelCase		)
        with patch('''transformers.onnx.features.is_tf_available'''      ,				__UpperCAmelCase		):
          lowerCAmelCase__				:  int		       =     FeaturesManager.determine_framework(self.test_model		)
          self.assertEqual(__UpperCAmelCase      ,				self.framework_pt		)

        # PyTorch not in environment -> use TensorFlow
        lowerCAmelCase__				:  Union[str, Any]		       =     MagicMock(return_value=__UpperCAmelCase		)
        with patch('''transformers.onnx.features.is_torch_available'''      ,				__UpperCAmelCase		):
          lowerCAmelCase__				:  Optional[int]		       =     FeaturesManager.determine_framework(self.test_model		)
          self.assertEqual(__UpperCAmelCase      ,				self.framework_tf		)

        # Both in environment -> use PyTorch
        lowerCAmelCase__				:  Union[str, Any]		       =     MagicMock(return_value=__UpperCAmelCase		)
        lowerCAmelCase__				:  List[Any]		       =     MagicMock(return_value=__UpperCAmelCase		)
        with patch('''transformers.onnx.features.is_tf_available'''      ,				__UpperCAmelCase		), patch(
            '''transformers.onnx.features.is_torch_available'''      ,				__UpperCAmelCase		):
          lowerCAmelCase__				:  List[Any]		       =     FeaturesManager.determine_framework(self.test_model		)
          self.assertEqual(__UpperCAmelCase      ,				self.framework_pt		)

        # Both not in environment -> raise error
        lowerCAmelCase__				:  Optional[int]		       =     MagicMock(return_value=__UpperCAmelCase		)
        lowerCAmelCase__				:  Dict		       =     MagicMock(return_value=__UpperCAmelCase		)
        with patch('''transformers.onnx.features.is_tf_available'''      ,				__UpperCAmelCase		), patch(
            '''transformers.onnx.features.is_torch_available'''      ,				__UpperCAmelCase		):
          with self.assertRaises(__UpperCAmelCase		):
            lowerCAmelCase__				:  Tuple		       =     FeaturesManager.determine_framework(self.test_model		)







import unittest

import numpy as np

from diffusers import LMSDiscreteScheduler, OnnxStableDiffusionInpaintPipeline
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  _lowerCAmelCase     (  _lowercase    ,			unittest.TestCase	):
      # FIXME: add fast tests
      pass






@nightly
@require_onnxruntime
@require_torch_gpu
class  _lowerCAmelCase     (  unittest.TestCase	):

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

      @property
      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  str		       =     ort.SessionOptions()
        lowerCAmelCase__				:  List[Any]		       =     False
        return options

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Tuple		       =     load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/in_paint/overture-creations-5sI6fQgYIuo.png'''		)
        lowerCAmelCase__				:  Optional[int]		       =     load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/in_paint/overture-creations-5sI6fQgYIuo_mask.png'''		)
        lowerCAmelCase__				:  List[Any]		       =     OnnxStableDiffusionInpaintPipeline.from_pretrained(
            '''runwayml/stable-diffusion-inpainting'''      ,				revision='''onnx'''      ,				safety_checker=__UpperCAmelCase      ,				feature_extractor=__UpperCAmelCase      ,				provider=self.gpu_provider      ,				sess_options=self.gpu_options      ,				)
        pipe.set_progress_bar_config(disable=__UpperCAmelCase		)

        lowerCAmelCase__				:  Optional[int]		       =     '''A red cat sitting on a park bench'''

        lowerCAmelCase__				:  List[Any]		       =     np.random.RandomState(0		)
        lowerCAmelCase__				:  int		       =     pipe(
            prompt=__UpperCAmelCase      ,				image=__UpperCAmelCase      ,				mask_image=__UpperCAmelCase      ,				guidance_scale=7.5      ,				num_inference_steps=10      ,				generator=__UpperCAmelCase      ,				output_type='''np'''      ,				)
        lowerCAmelCase__				:  Any		       =     output.images
        lowerCAmelCase__				:  Any		       =     images[0, 255:258, 255:258, -1]

        assert images.shape == (1, 512, 512, 3)
        lowerCAmelCase__				:  Any		       =     np.array([0.2514, 0.3007, 0.3517, 0.1790, 0.2382, 0.3167, 0.1944, 0.2273, 0.2464]		)

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







      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  List[str]		       =     load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/in_paint/overture-creations-5sI6fQgYIuo.png'''		)
        lowerCAmelCase__				:  str		       =     load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/in_paint/overture-creations-5sI6fQgYIuo_mask.png'''		)
        lowerCAmelCase__				:  Union[str, Any]		       =     LMSDiscreteScheduler.from_pretrained(
            '''runwayml/stable-diffusion-inpainting'''      ,				subfolder='''scheduler'''      ,				revision='''onnx'''		)
        lowerCAmelCase__				:  Union[str, Any]		       =     OnnxStableDiffusionInpaintPipeline.from_pretrained(
            '''runwayml/stable-diffusion-inpainting'''      ,				revision='''onnx'''      ,				scheduler=__UpperCAmelCase      ,				safety_checker=__UpperCAmelCase      ,				feature_extractor=__UpperCAmelCase      ,				provider=self.gpu_provider      ,				sess_options=self.gpu_options      ,				)
        pipe.set_progress_bar_config(disable=__UpperCAmelCase		)

        lowerCAmelCase__				:  str		       =     '''A red cat sitting on a park bench'''

        lowerCAmelCase__				:  Union[str, Any]		       =     np.random.RandomState(0		)
        lowerCAmelCase__				:  Optional[Any]		       =     pipe(
            prompt=__UpperCAmelCase      ,				image=__UpperCAmelCase      ,				mask_image=__UpperCAmelCase      ,				guidance_scale=7.5      ,				num_inference_steps=20      ,				generator=__UpperCAmelCase      ,				output_type='''np'''      ,				)
        lowerCAmelCase__				:  Tuple		       =     output.images
        lowerCAmelCase__				:  Optional[Any]		       =     images[0, 255:258, 255:258, -1]

        assert images.shape == (1, 512, 512, 3)
        lowerCAmelCase__				:  Optional[int]		       =     np.array([0.0086, 0.0077, 0.0083, 0.0093, 0.0107, 0.0139, 0.0094, 0.0097, 0.0125]		)

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

