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, server,
model_file, model_file,
provider=device.provider, provider=device.provider,
provider_options=device.options, sess_options=device.options,
) )
elif params.format == "pth": elif params.format == "pth":
model = RRDBNet( model = RRDBNet(

View File

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

View File

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

View File

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