Bethie's picture
Define path to ONNX weights
b768d73 verified
raw
history blame contribute delete
912 Bytes
import onnxruntime as ort
VAE_CONFIGS = 'onnx_1/vae_decoder/config.json'
UNET_CONFIGS = 'onnx_1/unet/config.json'
TEXT_ENCODER_2_CONFIGS = 'onnx_1/text_encoder_2/config.json'
VAE_ONNX_PATH = 'onnx_1/vae_decoder/model.onnx'
UNET_ONNX_PATH = 'onnx_1/unet/canny_2/model.onnx'
TOKENIZER_PATH = 'onnx_1/tokenizer'
TOKENIZER_PATH2 = 'onnx_1/tokenizer_2'
TEXT_ENCODER_PATH = 'onnx_1/text_encoder/model.onnx'
TEXT_ENCODER_PATH2 = 'onnx_1/text_encoder_2/model.onnx'
SCHEDULER_PATH = 'onnx_1/scheduler'
CONTROLNEXT_ONNX_PATH = 'onnx_1/controlnext/model.onnx'
IMAGE_ENCODER_ONNX_PATH = 'onnx_1/image_encoder/model.onnx'
PROJ_ONNX_PATH = 'onnx_1/proj/model.onnx'
providers = ['CUDAExecutionProvider']
provider_options_0 = [{'device_id': 0}]
provider_options_1 = [{'device_id': 1}]
session_options = ort.SessionOptions()
session_options.enable_mem_pattern = False
session_options.enable_cpu_mem_arena = False