



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			(       _lowercase	,	unittest.TestCase				):

 """simple docstring"""



 _lowercase   :    int      						=					'''hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'''

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE=0      ) ->			Union[str, Any]:
  a__												= floats_tensor((1, 3, 1_2_8, 1_2_8)		,				rng=random.Random(SCREAMING_SNAKE_CASE      )      )
  a__												= np.random.RandomState(SCREAMING_SNAKE_CASE      )
  a__												= {
      '''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 						_UpperCAmelCase       (	self      ) ->			str:
  a__												= OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint		,				provider='''CPUExecutionProvider'''      )
  pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE      )

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

  assert image.shape == (1, 1_2_8, 1_2_8, 3)
  a__												= np.array([0.6_96_43, 0.5_84_84, 0.5_03_14, 0.5_87_60, 0.5_53_68, 0.5_96_43, 0.5_15_29, 0.4_12_17, 0.4_90_87]      )
  assert np.abs(image_slice - expected_slice      ).max() < 1e-1

 def 						_UpperCAmelCase       (	self      ) ->			Dict:
  a__												= OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint		,				provider='''CPUExecutionProvider'''      )
  a__												= PNDMScheduler.from_config(pipe.scheduler.config		,				skip_prk_steps=SCREAMING_SNAKE_CASE      )
  pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE      )

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

  assert image.shape == (1, 1_2_8, 1_2_8, 3)
  a__												= np.array([0.6_17_37, 0.5_46_42, 0.5_31_83, 0.5_44_65, 0.5_27_42, 0.6_05_25, 0.4_99_69, 0.4_06_55, 0.4_81_54]      )

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

 def 						_UpperCAmelCase       (	self      ) ->			int:
  a__												= OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint		,				provider='''CPUExecutionProvider'''      )
  a__												= LMSDiscreteScheduler.from_config(pipe.scheduler.config      )
  pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE      )

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

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

  assert image.shape == (1, 1_2_8, 1_2_8, 3)
  a__												= np.array([0.5_27_61, 0.5_99_77, 0.4_90_33, 0.4_96_19, 0.5_42_82, 0.5_03_11, 0.4_76_00, 0.4_09_18, 0.4_52_03]      )

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

 def 						_UpperCAmelCase       (	self      ) ->			Optional[Any]:
  a__												= OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint		,				provider='''CPUExecutionProvider'''      )
  a__												= EulerDiscreteScheduler.from_config(pipe.scheduler.config      )
  pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE      )

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

  assert image.shape == (1, 1_2_8, 1_2_8, 3)
  a__												= np.array([0.5_29_11, 0.6_00_04, 0.4_92_29, 0.4_98_05, 0.5_45_02, 0.5_06_80, 0.4_77_77, 0.4_10_28, 0.4_53_04]      )

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

 def 						_UpperCAmelCase       (	self      ) ->			Any:
  a__												= OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint		,				provider='''CPUExecutionProvider'''      )
  a__												= EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config      )
  pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE      )

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

  assert image.shape == (1, 1_2_8, 1_2_8, 3)
  a__												= np.array([0.5_29_11, 0.6_00_04, 0.4_92_29, 0.4_98_05, 0.5_45_02, 0.5_06_80, 0.4_77_77, 0.4_10_28, 0.4_53_04]      )

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



 def 						_UpperCAmelCase       (	self      ) ->			str:
  a__												= OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint		,				provider='''CPUExecutionProvider'''      )
  a__												= DPMSolverMultistepScheduler.from_config(pipe.scheduler.config      )
  pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE      )

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

  assert image.shape == (1, 1_2_8, 1_2_8, 3)
  a__												= np.array([0.6_53_31, 0.5_82_77, 0.4_82_04, 0.5_60_59, 0.5_36_65, 0.5_62_35, 0.5_09_69, 0.4_00_09, 0.4_65_52]      )

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





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

 """simple docstring"""



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

 @property
 def 						_UpperCAmelCase       (	self      ) ->			str:
  a__												= ort.SessionOptions()
  a__												= False
  return options

 def 						_UpperCAmelCase       (	self      ) ->			List[Any]:
  a__												= load_image(
      '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
      '''/img2img/sketch-mountains-input.jpg'''      )
  a__												= init_image.resize((7_6_8, 5_1_2)      )
  # using the PNDM scheduler by default
  a__												= OnnxStableDiffusionImgaImgPipeline.from_pretrained(
      '''CompVis/stable-diffusion-v1-4'''		,				revision='''onnx'''		,				safety_checker=SCREAMING_SNAKE_CASE		,				feature_extractor=SCREAMING_SNAKE_CASE		,				provider=self.gpu_provider		,				sess_options=self.gpu_options		,				)
  pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE      )

  a__												= '''A fantasy landscape, trending on artstation'''

  a__												= np.random.RandomState(0      )
  a__												= pipe(
      prompt=SCREAMING_SNAKE_CASE		,				image=SCREAMING_SNAKE_CASE		,				strength=0.75		,				guidance_scale=7.5		,				num_inference_steps=1_0		,				generator=SCREAMING_SNAKE_CASE		,				output_type='''np'''		,				)
  a__												= output.images
  a__												= images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

  assert images.shape == (1, 5_1_2, 7_6_8, 3)
  a__												= np.array([0.49_09, 0.50_59, 0.53_72, 0.46_23, 0.48_76, 0.50_49, 0.48_20, 0.49_56, 0.50_19]      )
  # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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



 def 						_UpperCAmelCase       (	self      ) ->			Optional[int]:
  a__												= load_image(
      '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
      '''/img2img/sketch-mountains-input.jpg'''      )
  a__												= init_image.resize((7_6_8, 5_1_2)      )
  a__												= LMSDiscreteScheduler.from_pretrained(
      '''runwayml/stable-diffusion-v1-5'''		,				subfolder='''scheduler'''		,				revision='''onnx'''      )
  a__												= OnnxStableDiffusionImgaImgPipeline.from_pretrained(
      '''runwayml/stable-diffusion-v1-5'''		,				revision='''onnx'''		,				scheduler=SCREAMING_SNAKE_CASE		,				safety_checker=SCREAMING_SNAKE_CASE		,				feature_extractor=SCREAMING_SNAKE_CASE		,				provider=self.gpu_provider		,				sess_options=self.gpu_options		,				)
  pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE      )

  a__												= '''A fantasy landscape, trending on artstation'''

  a__												= np.random.RandomState(0      )
  a__												= pipe(
      prompt=SCREAMING_SNAKE_CASE		,				image=SCREAMING_SNAKE_CASE		,				strength=0.75		,				guidance_scale=7.5		,				num_inference_steps=2_0		,				generator=SCREAMING_SNAKE_CASE		,				output_type='''np'''		,				)
  a__												= output.images
  a__												= images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

  assert images.shape == (1, 5_1_2, 7_6_8, 3)
  a__												= np.array([0.80_43, 0.9_26, 0.95_81, 0.81_19, 0.89_54, 0.9_13, 0.72_09, 0.74_63, 0.74_31]      )
  # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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










import os
from pathlib import Path
from unittest.mock import patch

import pytest
import zstandard as zstd

from datasets.download.download_config import DownloadConfig
from datasets.utils.file_utils import (
    OfflineModeIsEnabled,
    cached_path,
    fsspec_get,
    fsspec_head,
    ftp_get,
    ftp_head,
    get_from_cache,
    http_get,
    http_head,
)


a_      :      int    =     '\\n    Text data.\n    Second line of data.'

a_      :      Optional[Any]    =     'file'






@pytest.fixture(scope='''session'''       )
def     __a					(		__UpperCAmelCase       ):
 a__												= tmp_path_factory.mktemp('''data'''       ) / (FILE_PATH + '''.zstd''')
 a__												= bytes(__UpperCAmelCase							,  '''utf-8'''       )
 with zstd.open(__UpperCAmelCase							,  '''wb'''       ) as f:
  f.write(__UpperCAmelCase       )
 return path






@pytest.fixture
def     __a					(		__UpperCAmelCase       ):
 with open(os.path.join(tmpfs.local_root_dir							,  __UpperCAmelCase       )							,  '''w'''       ) as f:
  f.write(__UpperCAmelCase       )
 return FILE_PATH






@pytest.mark.parametrize('''compression_format'''							,  ['''gzip''', '''xz''', '''zstd''']       )
def     __a					(		__UpperCAmelCase							,  __UpperCAmelCase							,  __UpperCAmelCase							,  __UpperCAmelCase							,  __UpperCAmelCase							,  __UpperCAmelCase       ):
 a__												= {'''gzip''': gz_file, '''xz''': xz_file, '''zstd''': zstd_path}
 a__												= input_paths[compression_format]
 a__												= tmp_path / '''cache'''
 a__												= DownloadConfig(cache_dir=__UpperCAmelCase							,  extract_compressed_file=__UpperCAmelCase       )
 a__												= cached_path(__UpperCAmelCase							,  download_config=__UpperCAmelCase       )
 with open(__UpperCAmelCase       ) as f:
  a__												= f.read()
 with open(__UpperCAmelCase       ) as f:
  a__												= f.read()
 assert extracted_file_content == expected_file_content






@pytest.mark.parametrize('''default_extracted'''							,  [True, False]       )
@pytest.mark.parametrize('''default_cache_dir'''							,  [True, False]       )
def     __a					(		__UpperCAmelCase							,  __UpperCAmelCase							,  __UpperCAmelCase							,  __UpperCAmelCase							,  __UpperCAmelCase       ):
 a__												= '''custom_cache'''
 a__												= '''custom_extracted_dir'''
 a__												= tmp_path / '''custom_extracted_path'''
 if default_extracted:
  a__												= ('''downloads''' if default_cache_dir else custom_cache_dir, '''extracted''')
 else:
  monkeypatch.setattr('''datasets.config.EXTRACTED_DATASETS_DIR'''							,  __UpperCAmelCase       )
  monkeypatch.setattr('''datasets.config.EXTRACTED_DATASETS_PATH'''							,  str(__UpperCAmelCase       )       )
  a__												= custom_extracted_path.parts[-2:] if default_cache_dir else (custom_cache_dir, custom_extracted_dir)

 a__												= xz_file
 a__												= (
     DownloadConfig(extract_compressed_file=__UpperCAmelCase       )
     if default_cache_dir
     else DownloadConfig(cache_dir=tmp_path / custom_cache_dir							,  extract_compressed_file=__UpperCAmelCase       )
 )
 a__												= cached_path(__UpperCAmelCase							,  download_config=__UpperCAmelCase       )
 assert Path(__UpperCAmelCase       ).parent.parts[-2:] == expected






def     __a					(		__UpperCAmelCase       ):
 # absolute path
 a__												= str(Path(__UpperCAmelCase       ).resolve()       )
 assert cached_path(__UpperCAmelCase       ) == text_file
 # relative path
 a__												= str(Path(__UpperCAmelCase       ).resolve().relative_to(Path(os.getcwd()       )       )       )
 assert cached_path(__UpperCAmelCase       ) == text_file






def     __a					(		__UpperCAmelCase       ):
 # absolute path
 a__												= str(tmp_path.resolve() / '''__missing_file__.txt'''       )
 with pytest.raises(__UpperCAmelCase       ):
  cached_path(__UpperCAmelCase       )
 # relative path
 a__												= '''./__missing_file__.txt'''
 with pytest.raises(__UpperCAmelCase       ):
  cached_path(__UpperCAmelCase       )






def     __a					(		__UpperCAmelCase       ):
 a__												= get_from_cache(f"tmp://{tmpfs_file}"       )
 with open(__UpperCAmelCase       ) as f:
  a__												= f.read()
 assert output_file_content == FILE_CONTENT






@patch('''datasets.config.HF_DATASETS_OFFLINE'''							,  __UpperCAmelCase       )
def     __a					(		):
 with pytest.raises(__UpperCAmelCase       ):
  cached_path('''https://huggingface.co'''       )






@patch('''datasets.config.HF_DATASETS_OFFLINE'''							,  __UpperCAmelCase       )
def     __a					(		__UpperCAmelCase       ):
 a__												= tmp_path_factory.mktemp('''data'''       ) / '''file.html'''
 with pytest.raises(__UpperCAmelCase       ):
  http_get('''https://huggingface.co'''							,  temp_file=__UpperCAmelCase       )
 with pytest.raises(__UpperCAmelCase       ):
  http_head('''https://huggingface.co'''       )






@patch('''datasets.config.HF_DATASETS_OFFLINE'''							,  __UpperCAmelCase       )
def     __a					(		__UpperCAmelCase       ):
 a__												= tmp_path_factory.mktemp('''data'''       ) / '''file.html'''
 with pytest.raises(__UpperCAmelCase       ):
  ftp_get('''ftp://huggingface.co'''							,  temp_file=__UpperCAmelCase       )
 with pytest.raises(__UpperCAmelCase       ):
  ftp_head('''ftp://huggingface.co'''       )






@patch('''datasets.config.HF_DATASETS_OFFLINE'''							,  __UpperCAmelCase       )
def     __a					(		__UpperCAmelCase       ):
 a__												= tmp_path_factory.mktemp('''data'''       ) / '''file.html'''
 with pytest.raises(__UpperCAmelCase       ):
  fsspec_get('''s3://huggingface.co'''							,  temp_file=__UpperCAmelCase       )
 with pytest.raises(__UpperCAmelCase       ):
  fsspec_head('''s3://huggingface.co'''       )





