Merge pull request #46 from yaosheng216/patch-4

Update ldm_sce.py
2024-09-27 09:49:27 +08:00
parent 5a94f6c4a2 4ad0f6038c
commit 8a14866562
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