Update ldm_sce.py

2024-09-27 09:39:22 +08:00
parent 30afc0a1de
commit 4ad0f6038c
1 changed files with 4 additions and 1 deletions
@@ -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