Use onnx as backend for CPU inference

app.py

@@ -2,14 +2,21 @@ from contextlib import nullcontext
 import gradio as gr
 import torch
 from torch import autocast
-from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionPipeline, StableDiffusionOnnxPipeline
 
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 context = autocast if device == "cuda" else nullcontext
 dtype = torch.float16 if device == "cuda" else torch.float32
 
-pipe = StableDiffusionPipeline.from_pretrained("lambdalabs/sd-pokemon-diffusers", torch_dtype=dtype)
+if device == "cuda":
+    pipe = StableDiffusionPipeline.from_pretrained("lambdalabs/sd-pokemon-diffusers", torch_dtype=dtype)
+else:
+    pipe = StableDiffusionOnnxPipeline.from_pretrained(
+        "lambdalabs/sd-pokemon-diffusers",
+        revision="onnx",
+        provider="CPUExecutionProvider"
+    )
 pipe = pipe.to(device)
 
 

requirements.txt

@@ -1,7 +1,8 @@
 --extra-index-url https://download.pytorch.org/whl/cu113
-torch
-diffusers
-transformers
-scipy
-ftfy
-datasets
+torch==1.12.1+cu113
+diffusers==0.3.0
+transformers==4.22.2
+scipy==1.9.1
+ftfy==6.1.1
+datasets==2.5.1
+onnxruntime==1.12.1