fix une2td ignoring class_labels (huggingface#5401)

* fix une2td ignoring class_labels * unet2.py error message updated * style and quality changes --------- Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
yoonseokjin · Oct 20, 2023 · e34e866 · e34e866
1 parent e5e535f
commit e34e866
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/models/unet_2d.py b/models/unet_2d.py
@@ -291,6 +291,8 @@ def forward(
 
             class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
             emb = emb + class_emb
+        elif self.class_embedding is None and class_labels is not None:
+            raise ValueError("class_embedding needs to be initialized in order to use class conditioning")
 
         # 2. pre-process
         skip_sample = sample