Update unet/conditional_unet_model.py

unet/conditional_unet_model.py

@@ -13,6 +13,13 @@ from diffusers.models.embeddings import GaussianFourierProjection, TimestepEmbed
 from diffusers.models.modeling_utils import ModelMixin
 from diffusers.models.unets.unet_2d_blocks import UNetMidBlock2D, get_down_block, get_up_block
 
+NUM_CLASSES_FLOOR_HUE = 10
+NUM_CLASSES_OBJECT_HUE = 10
+NUM_CLASSES_ORIENTATION = 15
+NUM_CLASSES_SCALE = 8
+NUM_CLASSES_SHAPE = 4
+NUM_CLASSES_WALL_HUE = 10
+
 @dataclass
 class UNet2DOutput(BaseOutput):
     """