From 4ad0f6038ccd62c83a3cfa78ba5949b3c7787d25 Mon Sep 17 00:00:00 2001
From: Great <44393369+yaosheng216@users.noreply.github.com>
Date: Fri, 27 Sep 2024 09:39:22 +0800
Subject: [PATCH] Update ldm_sce.py

---
 scepter/modules/model/network/ldm/ldm_sce.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/scepter/modules/model/network/ldm/ldm_sce.py b/scepter/modules/model/network/ldm/ldm_sce.py
index e60a371..50ab81c 100644
--- a/scepter/modules/model/network/ldm/ldm_sce.py
+++ b/scepter/modules/model/network/ldm/ldm_sce.py
@@ -105,7 +105,10 @@ class LatentDiffusionSCEControl(LatentDiffusion):
         hints = []
         for ctr in control:
             hint = self.control_processor(ctr)
-            hint = TT.ToTensor()(hint)
+            if len(hint.shape) == 4:
+                hint = TT.ToTensor()(hint.squeeze())
+            else:
+                hint = TT.ToTensor()(hint)
             hints.append(hint)
         hints = torch.stack(hints).to(control.device)
         return hints