fix(api): pass device options to ORT session (#38)
This commit is contained in:
parent
f2d0025545
commit
d473a0fd2d
|
@ -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(
|
||||
|
|
|
@ -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(
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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:
|
||||
|
|
Loading…
Reference in New Issue