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,
|
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(
|
||||||
|
|
|
@ -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(
|
||||||
|
|
|
@ -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,
|
||||||
|
|
|
@ -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:
|
||||||
|
|
Loading…
Reference in New Issue