1
0
Fork 0

fix(api): pass device options to ORT session (#38)

This commit is contained in:
Sean Sube 2023-02-14 17:12:52 -06:00
parent f2d0025545
commit d473a0fd2d
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
4 changed files with 7 additions and 7 deletions

View File

@ -53,7 +53,7 @@ def load_resrgan(
server,
model_file,
provider=device.provider,
provider_options=device.options,
sess_options=device.options,
)
elif params.format == "pth":
model = RRDBNet(

View File

@ -34,7 +34,7 @@ def load_stable_diffusion(
device.provider,
)
pipe = OnnxStableDiffusionUpscalePipeline.from_pretrained(
model_path, provider=device.provider, provider_options=device.options
model_path, provider=device.provider, sess_options=device.options
)
else:
logger.debug(

View File

@ -110,7 +110,7 @@ def load_pipeline(
scheduler = scheduler_type.from_pretrained(
model,
provider=device.provider,
provider_options=device.options,
sess_options=device.options,
subfolder="scheduler",
)
@ -135,14 +135,14 @@ def load_pipeline(
scheduler = scheduler_type.from_pretrained(
model,
provider=device.provider,
provider_options=device.options,
sess_options=device.options,
subfolder="scheduler",
)
pipe = pipeline.from_pretrained(
model,
custom_pipeline=custom_pipeline,
provider=device.provider,
provider_options=device.options,
sess_options=device.options,
revision="onnx",
safety_checker=None,
scheduler=scheduler,

View File

@ -47,11 +47,11 @@ class OnnxNet:
server: ServerContext,
model: str,
provider: str = "DmlExecutionProvider",
provider_options: Optional[dict] = None,
sess_options: Optional[dict] = None,
) -> None:
model_path = path.join(server.model_path, model)
self.session = InferenceSession(
model_path, providers=[provider], provider_options=provider_options
model_path, providers=[provider], provider_options=sess_options
)
def __call__(self, image: Any) -> Any: