




import tempfile
import unittest

import numpy as np

from diffusers import (
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionPipeline,
    PNDMScheduler,
)
from diffusers.utils.testing_utils import is_onnx_available, nightly, require_onnxruntime, require_torch_gpu

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
      import onnxruntime as ort





class     SCREAMING_SNAKE_CASE__		(_a    ,					unittest.TestCase			):
     lowercase_ :      int				= "hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"
     def        A__      (						self	:		str       ,						__lowerCamelCase	:		str=0						):



       """simple docstring"""
       lowerCAmelCase__										=      np.random.RandomState(__lowerCamelCase						)
       lowerCAmelCase__										=      {
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''guidance_scale''': 7.5,
           '''output_type''': '''numpy''',
       }
       return inputs
     def        A__      (						self	:		List[Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,						provider='''CPUExecutionProvider'''						)
       pipe.set_progress_bar_config(disable=__lowerCamelCase						)

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

       assert image.shape == (1, 1_28, 1_28, 3)
       lowerCAmelCase__										=      np.array([0.6_5072, 0.5_8492, 0.4_8219, 0.5_5521, 0.5_3180, 0.5_5939, 0.5_0697, 0.3_9800, 0.4_6455]						)

       assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-2
     def        A__      (						self	:		List[str]						):



       """simple docstring"""
       lowerCAmelCase__										=      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,						provider='''CPUExecutionProvider'''						)
       lowerCAmelCase__										=      PNDMScheduler.from_config(pipe.scheduler.config       ,						skip_prk_steps=__lowerCamelCase						)
       pipe.set_progress_bar_config(disable=__lowerCamelCase						)

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

       assert image.shape == (1, 1_28, 1_28, 3)
       lowerCAmelCase__										=      np.array([0.6_5863, 0.5_9425, 0.4_9326, 0.5_6313, 0.5_3875, 0.5_6627, 0.5_1065, 0.3_9777, 0.4_6330]						)

       assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-2
     def        A__      (						self	:		List[Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,						provider='''CPUExecutionProvider'''						)
       lowerCAmelCase__										=      LMSDiscreteScheduler.from_config(pipe.scheduler.config						)
       pipe.set_progress_bar_config(disable=__lowerCamelCase						)

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

       assert image.shape == (1, 1_28, 1_28, 3)
       lowerCAmelCase__										=      np.array([0.5_3755, 0.6_0786, 0.4_7402, 0.4_9488, 0.5_1869, 0.4_9819, 0.4_7985, 0.3_8957, 0.4_4279]						)

       assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-2
     def        A__      (						self	:		Tuple						):



       """simple docstring"""
       lowerCAmelCase__										=      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,						provider='''CPUExecutionProvider'''						)
       lowerCAmelCase__										=      EulerDiscreteScheduler.from_config(pipe.scheduler.config						)
       pipe.set_progress_bar_config(disable=__lowerCamelCase						)

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

       assert image.shape == (1, 1_28, 1_28, 3)
       lowerCAmelCase__										=      np.array([0.5_3755, 0.6_0786, 0.4_7402, 0.4_9488, 0.5_1869, 0.4_9819, 0.4_7985, 0.3_8957, 0.4_4279]						)

       assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-2
     def        A__      (						self	:		int						):



       """simple docstring"""
       lowerCAmelCase__										=      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,						provider='''CPUExecutionProvider'''						)
       lowerCAmelCase__										=      EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config						)
       pipe.set_progress_bar_config(disable=__lowerCamelCase						)

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

       assert image.shape == (1, 1_28, 1_28, 3)
       lowerCAmelCase__										=      np.array([0.5_3817, 0.6_0812, 0.4_7384, 0.4_9530, 0.5_1894, 0.4_9814, 0.4_7984, 0.3_8958, 0.4_4271]						)

       assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-2
     def        A__      (						self	:		str						):



       """simple docstring"""
       lowerCAmelCase__										=      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,						provider='''CPUExecutionProvider'''						)
       lowerCAmelCase__										=      DPMSolverMultistepScheduler.from_config(pipe.scheduler.config						)
       pipe.set_progress_bar_config(disable=__lowerCamelCase						)

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

       assert image.shape == (1, 1_28, 1_28, 3)
       lowerCAmelCase__										=      np.array([0.5_3895, 0.6_0808, 0.4_7933, 0.4_9608, 0.5_1886, 0.4_9950, 0.4_8053, 0.3_8957, 0.4_4200]						)

       assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-2
     def        A__      (						self	:		Optional[Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,						provider='''CPUExecutionProvider'''						)
       pipe.set_progress_bar_config(disable=__lowerCamelCase						)

       lowerCAmelCase__										=      self.get_dummy_inputs()
       lowerCAmelCase__										=      3 * [inputs['''prompt''']]

       # forward
       lowerCAmelCase__										=      pipe(**__lowerCamelCase						)
       lowerCAmelCase__										=      output.images[0, -3:, -3:, -1]

       lowerCAmelCase__										=      self.get_dummy_inputs()
       lowerCAmelCase__										=      3 * [inputs.pop('''prompt'''						)]

       lowerCAmelCase__										=      pipe.tokenizer(
           __lowerCamelCase       ,						padding='''max_length'''       ,						max_length=pipe.tokenizer.model_max_length       ,						truncation=__lowerCamelCase       ,						return_tensors='''np'''       ,						)
       lowerCAmelCase__										=      text_inputs['''input_ids''']

       lowerCAmelCase__										=      pipe.text_encoder(input_ids=text_inputs.astype(np.intaa						)						)[0]

       lowerCAmelCase__										=      prompt_embeds

       # forward
       lowerCAmelCase__										=      pipe(**__lowerCamelCase						)
       lowerCAmelCase__										=      output.images[0, -3:, -3:, -1]

       assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()						).max() < 1e-4



     def        A__      (						self	:		Optional[Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,						provider='''CPUExecutionProvider'''						)
       pipe.set_progress_bar_config(disable=__lowerCamelCase						)

       lowerCAmelCase__										=      self.get_dummy_inputs()
       lowerCAmelCase__										=      3 * ['''this is a negative prompt''']
       lowerCAmelCase__										=      negative_prompt
       lowerCAmelCase__										=      3 * [inputs['''prompt''']]

       # forward
       lowerCAmelCase__										=      pipe(**__lowerCamelCase						)
       lowerCAmelCase__										=      output.images[0, -3:, -3:, -1]

       lowerCAmelCase__										=      self.get_dummy_inputs()
       lowerCAmelCase__										=      3 * [inputs.pop('''prompt'''						)]

       lowerCAmelCase__										=      []
       for p in [prompt, negative_prompt]:
         lowerCAmelCase__										=      pipe.tokenizer(
             __lowerCamelCase       ,						padding='''max_length'''       ,						max_length=pipe.tokenizer.model_max_length       ,						truncation=__lowerCamelCase       ,						return_tensors='''np'''       ,						)
         lowerCAmelCase__										=      text_inputs['''input_ids''']

         embeds.append(pipe.text_encoder(input_ids=text_inputs.astype(np.intaa						)						)[0]						)

       lowerCAmelCase__				,      lowerCAmelCase__										=      embeds

       # forward
       lowerCAmelCase__										=      pipe(**__lowerCamelCase						)
       lowerCAmelCase__										=      output.images[0, -3:, -3:, -1]

       assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()						).max() < 1e-4





@nightly
@require_onnxruntime
@require_torch_gpu
class     SCREAMING_SNAKE_CASE__		(unittest.TestCase			):
     @property
     def        A__      (						self	:		List[Any]						):



       """simple docstring"""
       return (
           "CUDAExecutionProvider",
           {
               "gpu_mem_limit": "15000000000",  # 15GB
               "arena_extend_strategy": "kSameAsRequested",
           },
       )
     @property
     def        A__      (						self	:		Union[str, Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      ort.SessionOptions()
       lowerCAmelCase__										=      False
       return options
     def        A__      (						self	:		Optional[Any]						):



       """simple docstring"""
       # using the PNDM scheduler by default
       lowerCAmelCase__										=      OnnxStableDiffusionPipeline.from_pretrained(
           '''CompVis/stable-diffusion-v1-4'''       ,						revision='''onnx'''       ,						safety_checker=__lowerCamelCase       ,						feature_extractor=__lowerCamelCase       ,						provider=self.gpu_provider       ,						sess_options=self.gpu_options       ,						)
       sd_pipe.set_progress_bar_config(disable=__lowerCamelCase						)

       lowerCAmelCase__										=      '''A painting of a squirrel eating a burger'''
       np.random.seed(0						)
       lowerCAmelCase__										=      sd_pipe([prompt]       ,						guidance_scale=6.0       ,						num_inference_steps=10       ,						output_type='''np'''						)
       lowerCAmelCase__										=      output.images

       lowerCAmelCase__										=      image[0, -3:, -3:, -1]

       assert image.shape == (1, 5_12, 5_12, 3)
       lowerCAmelCase__										=      np.array([0.0452, 0.0390, 0.0087, 0.0350, 0.0617, 0.0364, 0.0544, 0.0523, 0.0720]						)

       assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-3
     def        A__      (						self	:		int						):



       """simple docstring"""
       lowerCAmelCase__										=      DDIMScheduler.from_pretrained(
           '''runwayml/stable-diffusion-v1-5'''       ,						subfolder='''scheduler'''       ,						revision='''onnx'''						)
       lowerCAmelCase__										=      OnnxStableDiffusionPipeline.from_pretrained(
           '''runwayml/stable-diffusion-v1-5'''       ,						revision='''onnx'''       ,						scheduler=__lowerCamelCase       ,						safety_checker=__lowerCamelCase       ,						feature_extractor=__lowerCamelCase       ,						provider=self.gpu_provider       ,						sess_options=self.gpu_options       ,						)
       sd_pipe.set_progress_bar_config(disable=__lowerCamelCase						)

       lowerCAmelCase__										=      '''open neural network exchange'''
       lowerCAmelCase__										=      np.random.RandomState(0						)
       lowerCAmelCase__										=      sd_pipe([prompt]       ,						guidance_scale=7.5       ,						num_inference_steps=10       ,						generator=__lowerCamelCase       ,						output_type='''np'''						)
       lowerCAmelCase__										=      output.images
       lowerCAmelCase__										=      image[0, -3:, -3:, -1]

       assert image.shape == (1, 5_12, 5_12, 3)
       lowerCAmelCase__										=      np.array([0.2867, 0.1974, 0.1481, 0.7294, 0.7251, 0.6667, 0.4194, 0.5642, 0.6486]						)

       assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-3
     def        A__      (						self	:		List[str]						):



       """simple docstring"""
       lowerCAmelCase__										=      LMSDiscreteScheduler.from_pretrained(
           '''runwayml/stable-diffusion-v1-5'''       ,						subfolder='''scheduler'''       ,						revision='''onnx'''						)
       lowerCAmelCase__										=      OnnxStableDiffusionPipeline.from_pretrained(
           '''runwayml/stable-diffusion-v1-5'''       ,						revision='''onnx'''       ,						scheduler=__lowerCamelCase       ,						safety_checker=__lowerCamelCase       ,						feature_extractor=__lowerCamelCase       ,						provider=self.gpu_provider       ,						sess_options=self.gpu_options       ,						)
       sd_pipe.set_progress_bar_config(disable=__lowerCamelCase						)

       lowerCAmelCase__										=      '''open neural network exchange'''
       lowerCAmelCase__										=      np.random.RandomState(0						)
       lowerCAmelCase__										=      sd_pipe([prompt]       ,						guidance_scale=7.5       ,						num_inference_steps=10       ,						generator=__lowerCamelCase       ,						output_type='''np'''						)
       lowerCAmelCase__										=      output.images
       lowerCAmelCase__										=      image[0, -3:, -3:, -1]

       assert image.shape == (1, 5_12, 5_12, 3)
       lowerCAmelCase__										=      np.array([0.2306, 0.1959, 0.1593, 0.6549, 0.6394, 0.5408, 0.5065, 0.6010, 0.6161]						)

       assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-3
     def        A__      (						self	:		List[Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      0

       def test_callback_fn(__lowerCamelCase	:		int       ,						__lowerCamelCase	:		int       ,						__lowerCamelCase	:		np.ndarray						) -> None:
         lowerCAmelCase__										=      True
         nonlocal number_of_steps
         number_of_steps += 1
         if step == 0:
           assert latents.shape == (1, 4, 64, 64)
           lowerCAmelCase__										=      latents[0, -3:, -3:, -1]
           lowerCAmelCase__										=      np.array(
               [-0.6772, -0.3835, -1.2456, 0.1905, -1.0974, 0.6967, -1.9353, 0.0178, 1.0167]						)

           assert np.abs(latents_slice.flatten() - expected_slice						).max() < 1e-3
         elif step == 5:
           assert latents.shape == (1, 4, 64, 64)
           lowerCAmelCase__										=      latents[0, -3:, -3:, -1]
           lowerCAmelCase__										=      np.array(
               [-0.3351, 0.2241, -0.1837, -0.2325, -0.6577, 0.3393, -0.0241, 0.5899, 1.3875]						)

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

       lowerCAmelCase__										=      False

       lowerCAmelCase__										=      OnnxStableDiffusionPipeline.from_pretrained(
           '''runwayml/stable-diffusion-v1-5'''       ,						revision='''onnx'''       ,						safety_checker=__lowerCamelCase       ,						feature_extractor=__lowerCamelCase       ,						provider=self.gpu_provider       ,						sess_options=self.gpu_options       ,						)
       pipe.set_progress_bar_config(disable=__lowerCamelCase						)

       lowerCAmelCase__										=      '''Andromeda galaxy in a bottle'''

       lowerCAmelCase__										=      np.random.RandomState(0						)
       pipe(
           prompt=__lowerCamelCase       ,						num_inference_steps=5       ,						guidance_scale=7.5       ,						generator=__lowerCamelCase       ,						callback=__lowerCamelCase       ,						callback_steps=1       ,						)
       assert test_callback_fn.has_been_called
       assert number_of_steps == 6



     def        A__      (						self	:		List[Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      OnnxStableDiffusionPipeline.from_pretrained(
           '''runwayml/stable-diffusion-v1-5'''       ,						revision='''onnx'''       ,						safety_checker=__lowerCamelCase       ,						feature_extractor=__lowerCamelCase       ,						provider=self.gpu_provider       ,						sess_options=self.gpu_options       ,						)
       assert isinstance(__lowerCamelCase       ,						__lowerCamelCase						)
       assert pipe.safety_checker is None

       lowerCAmelCase__										=      pipe('''example prompt'''       ,						num_inference_steps=2						).images[0]
       assert image is not None

       # check that there's no error when saving a pipeline with one of the models being None
       with tempfile.TemporaryDirectory() as tmpdirname:
         pipe.save_pretrained(__lowerCamelCase						)
         lowerCAmelCase__										=      OnnxStableDiffusionPipeline.from_pretrained(__lowerCamelCase						)

       # sanity check that the pipeline still works
       assert pipe.safety_checker is None
       lowerCAmelCase__										=      pipe('''example prompt'''       ,						num_inference_steps=2						).images[0]
       assert image is not None







from __future__ import annotations
def 		a_							( __lowerCAmelCase   , __lowerCAmelCase = None				):
  lowerCAmelCase__										=      word_bank or []
  # create a table
  lowerCAmelCase__										=      len(__lowerCAmelCase				) + 1

  lowerCAmelCase__										=      []
  for _ in range(__lowerCAmelCase				):
    table.append([]				)
  # seed value
  lowerCAmelCase__										=      [[]]  # because empty string has empty combination

  # iterate through the indices
  for i in range(__lowerCAmelCase				):
    # condition
    if table[i] != []:
      for word in word_bank:
        # slice condition
        if target[i : i + len(__lowerCAmelCase				)] == word:
          lowerCAmelCase__										=      [
              [word, *way] for way in table[i]
          ]
          # adds the word to every combination the current position holds
          # now,push that combination to the table[i+len(word)]
          table[i + len(__lowerCAmelCase				)] += new_combinations

    # combinations are in reverse order so reverse for better output
  for combination in table[len(__lowerCAmelCase				)]:
    combination.reverse()

  return table[len(__lowerCAmelCase				)]


if __name__ == "__main__":
      print(all_construct("""jwajalapa""", ["""jwa""", """j""", """w""", """a""", """la""", """lapa"""]))
      print(all_construct("""rajamati""", ["""s""", """raj""", """amat""", """raja""", """ma""", """i""", """t"""]))
      print(
          all_construct(
              """hexagonosaurus""",
              ["""h""", """ex""", """hex""", """ag""", """ago""", """ru""", """auru""", """rus""", """go""", """no""", """o""", """s"""],
          )
      )

