diff --git a/api/onnx_web/diffusion.py b/api/onnx_web/diffusion.py
index 439a6681..dfe56c49 100644
--- a/api/onnx_web/diffusion.py
+++ b/api/onnx_web/diffusion.py
@@ -12,6 +12,9 @@ import gc
 import numpy as np
 import torch
 
+from .chain import (
+    StageParams,
+)
 from .image import (
     expand_image,
 )
@@ -25,7 +28,6 @@ from .utils import (
     ImageParams,
     Border,
     ServerContext,
-    StageParams,
     Size,
 )