fix(api): avoid setting device on schedulers that do not support it

ssube · Feb 4, 2023 · d636ce3 · d636ce3
1 parent 39a422f
commit d636ce3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/api/onnx_web/diffusion/load.py b/api/onnx_web/diffusion/load.py
@@ -68,7 +68,7 @@ def load_pipeline(pipeline: DiffusionPipeline, model: str, provider: str, schedu
             scheduler=scheduler.from_pretrained(model, subfolder='scheduler')
         )
 
-        if device is not None:
+        if device is not None and hasattr(pipe, 'to'):
             pipe = pipe.to(device)
 
         last_pipeline_instance = pipe
@@ -80,7 +80,7 @@ def load_pipeline(pipeline: DiffusionPipeline, model: str, provider: str, schedu
         scheduler = scheduler.from_pretrained(
             model, subfolder='scheduler')
 
-        if device is not None:
+        if device is not None and hasattr(scheduler, 'to'):
             scheduler = scheduler.to(device)
 
         pipe.scheduler = scheduler