apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

feat(api): add option to switch between AMD hardware and CPU software rendering

This commit is contained in:
Sean Sube 2023-01-05 17:24:14 -06:00
parent 64c78b887e
commit 668e46a404
1 changed files with 8 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
api/serve.py
View File

@ -30,6 +30,12 @@ max_width = 512
model_path = environ.get('ONNX_WEB_MODEL_PATH', "../models/stable-diffusion-onnx-v1-5") model_path = environ.get('ONNX_WEB_MODEL_PATH', "../models/stable-diffusion-onnx-v1-5")
output_path = environ.get('ONNX_WEB_OUTPUT_PATH', "../outputs") output_path = environ.get('ONNX_WEB_OUTPUT_PATH', "../outputs")
# platforms
platform_providers = {
'amd': 'DmlExecutionProvider',
'cpu': 'CPUExecutionProvider',
}
# schedulers # schedulers
pipeline_schedulers = { pipeline_schedulers = {
'ddim': DDIMScheduler.from_pretrained(model_path, subfolder="scheduler"), 'ddim': DDIMScheduler.from_pretrained(model_path, subfolder="scheduler"),
@ -85,6 +91,7 @@ def txt2img():
user = request.remote_addr user = request.remote_addr
prompt = request.args.get('prompt', default_prompt) prompt = request.args.get('prompt', default_prompt)
provider = get_from_map(request.args, 'provider', platform_providers, 'amd')
scheduler = get_from_map(request.args, 'scheduler', scheduler = get_from_map(request.args, 'scheduler',
pipeline_schedulers, 'euler-a') pipeline_schedulers, 'euler-a')
cfg = get_and_clamp(request.args, 'cfg', default_cfg, max_cfg, 0) cfg = get_and_clamp(request.args, 'cfg', default_cfg, max_cfg, 0)
@ -103,7 +110,7 @@ def txt2img():
pipe = OnnxStableDiffusionPipeline.from_pretrained( pipe = OnnxStableDiffusionPipeline.from_pretrained(
model_path, model_path,
provider="DmlExecutionProvider", provider=provider,
safety_checker=None, safety_checker=None,
scheduler=scheduler scheduler=scheduler
) )