fix(api): pass both device and session options to ORT (#38)
This commit is contained in:
parent
feb4603171
commit
8a2a9174ba
|
@ -52,8 +52,8 @@ def load_resrgan(
|
|||
model = OnnxNet(
|
||||
server,
|
||||
model_file,
|
||||
provider=device.provider,
|
||||
sess_options=device.options,
|
||||
provider=device.ort_provider(),
|
||||
sess_options=device.sess_options(),
|
||||
)
|
||||
elif params.format == "pth":
|
||||
model = RRDBNet(
|
||||
|
|
|
@ -34,7 +34,9 @@ def load_stable_diffusion(
|
|||
device.provider,
|
||||
)
|
||||
pipe = OnnxStableDiffusionUpscalePipeline.from_pretrained(
|
||||
model_path, provider=device.provider, sess_options=device.options
|
||||
model_path,
|
||||
provider=device.ort_provider(),
|
||||
sess_options=device.sess_options(),
|
||||
)
|
||||
else:
|
||||
logger.debug(
|
||||
|
@ -43,7 +45,8 @@ def load_stable_diffusion(
|
|||
device.provider,
|
||||
)
|
||||
pipe = StableDiffusionUpscalePipeline.from_pretrained(
|
||||
model_path, provider=device.provider
|
||||
model_path,
|
||||
provider=device.provider,
|
||||
)
|
||||
|
||||
server.cache.set("diffusion", cache_key, pipe)
|
||||
|
|
|
@ -109,8 +109,8 @@ def load_pipeline(
|
|||
logger.debug("loading new diffusion scheduler")
|
||||
scheduler = scheduler_type.from_pretrained(
|
||||
model,
|
||||
provider=device.provider,
|
||||
sess_options=device.options,
|
||||
provider=device.ort_provider(),
|
||||
sess_options=device.sess_options(),
|
||||
subfolder="scheduler",
|
||||
)
|
||||
|
||||
|
@ -134,15 +134,15 @@ def load_pipeline(
|
|||
logger.debug("loading new diffusion pipeline from %s", model)
|
||||
scheduler = scheduler_type.from_pretrained(
|
||||
model,
|
||||
provider=device.provider,
|
||||
sess_options=device.options,
|
||||
provider=device.ort_provider(),
|
||||
sess_options=device.sess_options(),
|
||||
subfolder="scheduler",
|
||||
)
|
||||
pipe = pipeline.from_pretrained(
|
||||
model,
|
||||
custom_pipeline=custom_pipeline,
|
||||
provider=device.provider,
|
||||
sess_options=device.options,
|
||||
provider=device.ort_provider(),
|
||||
sess_options=device.sess_options(),
|
||||
revision="onnx",
|
||||
safety_checker=None,
|
||||
scheduler=scheduler,
|
||||
|
|
|
@ -3,7 +3,7 @@ from typing import Any, Optional
|
|||
|
||||
import numpy as np
|
||||
import torch
|
||||
from onnxruntime import InferenceSession
|
||||
from onnxruntime import InferenceSession, SessionOptions
|
||||
|
||||
from ..utils import ServerContext
|
||||
|
||||
|
@ -47,7 +47,7 @@ class OnnxNet:
|
|||
server: ServerContext,
|
||||
model: str,
|
||||
provider: str = "DmlExecutionProvider",
|
||||
sess_options: Optional[dict] = None,
|
||||
sess_options: Optional[SessionOptions] = None,
|
||||
) -> None:
|
||||
model_path = path.join(server.model_path, model)
|
||||
self.session = InferenceSession(
|
||||
|
|
|
@ -1,4 +1,5 @@
|
|||
from enum import IntEnum
|
||||
from onnxruntime import SessionOptions
|
||||
from typing import Any, Dict, Literal, Optional, Tuple, Union
|
||||
|
||||
|
||||
|
@ -79,6 +80,15 @@ class DeviceParams:
|
|||
def __str__(self) -> str:
|
||||
return "%s - %s (%s)" % (self.device, self.provider, self.options)
|
||||
|
||||
def ort_provider(self) -> Tuple[str, Any]:
|
||||
if self.options is None:
|
||||
return self.provider
|
||||
else:
|
||||
return (self.provider, self.options)
|
||||
|
||||
def sess_options(self) -> SessionOptions:
|
||||
return SessionOptions()
|
||||
|
||||
def torch_device(self) -> str:
|
||||
if self.device.startswith("cuda"):
|
||||
return self.device
|
||||
|
|
Loading…
Reference in New Issue