From a716f6db1f64184167765e250ae74ad5cd342894 Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Mon, 18 Dec 2023 22:21:33 -0600
Subject: [PATCH] fix(api): add loading path for controlnet

---
 api/onnx_web/diffusers/load.py | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/api/onnx_web/diffusers/load.py b/api/onnx_web/diffusers/load.py
index b6620e8b2..7f28523bf 100644
--- a/api/onnx_web/diffusers/load.py
+++ b/api/onnx_web/diffusers/load.py
@@ -229,7 +229,23 @@ def load_pipeline(
                 tokenizer_2=components.get("tokenizer_2", None),
             )
         else:
-            if "vae" in components:
+            if "controlnet" in components:
+                logger.debug(
+                    "assembling SD pipeline for %s with ControlNet",
+                    pipeline_class.__name__,
+                )
+                pipe = pipeline_class(
+                    components["vae"],
+                    components["text_encoder"],
+                    components["tokenizer"],
+                    components["unet"],
+                    components["controlnet"],
+                    scheduler,
+                    None,
+                    None,
+                    requires_safety_checker=False,
+                )
+            elif "vae" in components:
                 # upscale uses a single VAE
                 logger.debug(
                     "assembling SD pipeline for %s with single VAE",