changing the root

defaults/deepfashion.py

@@ -1,4 +1,8 @@
 from yacs.config import CfgNode as CN
+import os
+
+# Use the same persistent directory where models are stored
+PERSISTENT_DIR = "/data/models"
 
 _C = CN()
 
@@ -20,10 +24,22 @@ _C.MODEL.U_COND_DOWN_BLOCK_GUIDANCE = False
 _C.MODEL.U_COND_UP_BLOCK_GUIDANCE = False
 
 _C.MODEL.FIRST_STAGE_CONFIG = CN()
-_C.MODEL.FIRST_STAGE_CONFIG.PRETRAINED_PATH = "pretrained_models/vae"
+# _C.MODEL.FIRST_STAGE_CONFIG.PRETRAINED_PATH = "pretrained_models/vae"
+_C.MODEL.FIRST_STAGE_CONFIG.PRETRAINED_PATH = os.path.join(
+    PERSISTENT_DIR,
+    "pretrained_models",
+    "vae"
+)
 
 _C.MODEL.UNET_CONFIG = CN()
-_C.MODEL.UNET_CONFIG.PRETRAINED_PATH = "pretrained_models/unet"
+# _C.MODEL.UNET_CONFIG.PRETRAINED_PATH = "pretrained_models/unet"
+_C.MODEL.UNET_CONFIG.PRETRAINED_PATH =  os.path.join(
+    PERSISTENT_DIR,
+    "pretrained_models",
+    "unet"
+)
+
+
 _C.MODEL.UNET_CONFIG.TRAINABLE_BLOCK_IDX = [11, 10, 9, 8, 7, 6, 5, 4, 3]
 _C.MODEL.UNET_CONFIG.TRAIN_SELF_ATTN_Q = False
 _C.MODEL.UNET_CONFIG.TRAIN_SELF_ATTN_K = False
@@ -35,10 +51,25 @@ _C.MODEL.UNET_CONFIG.TRAIN_CROSS_ATTN_V = True
 _C.MODEL.SCHEDULER_CONFIG = CN()
 _C.MODEL.SCHEDULER_CONFIG.NAME = "ddpm"
 _C.MODEL.SCHEDULER_CONFIG.PRETRAINED_PATH = "pretrained_models/scheduler"
+_C.MODEL.SCHEDULER_CONFIG.PRETRAINED_PATH = os.path.join(
+    PERSISTENT_DIR,
+    "pretrained_models",
+    "scheduler"
+)
+
+
 _C.MODEL.SCHEDULER_CONFIG.CUBIC_SAMPLING = True
 
 _C.MODEL.COND_STAGE_CONFIG = CN()
-_C.MODEL.COND_STAGE_CONFIG.PRETRAINED_PATH = "pretrained_models/swin/swin_base_patch4_window12_384_22kto1k.pth"
+# _C.MODEL.COND_STAGE_CONFIG.PRETRAINED_PATH = "pretrained_models/swin/swin_base_patch4_window12_384_22kto1k.pth"
+_C.MODEL.COND_STAGE_CONFIG.PRETRAINED_PATH =  os.path.join(
+    PERSISTENT_DIR,
+    "pretrained_models",
+    "swin",
+    "swin_base_patch4_window12_384_22kto1k.pth"
+)
+
+
 _C.MODEL.COND_STAGE_CONFIG.EMBED_DIM = 128
 _C.MODEL.COND_STAGE_CONFIG.DEPTHS = [2, 2, 18, 2]
 _C.MODEL.COND_STAGE_CONFIG.NUM_HEADS = [4, 8, 16, 32]

pose_transfer_train.py

@@ -25,13 +25,13 @@ from defaults import pose_transfer_C as cfg
 import os
 
 # Use persistent dir for pretrained swin weights
-PERSISTENT_DIR = "/data/models"
-cfg.MODEL.COND_STAGE_CONFIG.PRETRAINED_PATH = os.path.join(
-    PERSISTENT_DIR,
-    "pretrained_models",
-    "swin",
-    "swin_base_patch4_window12_384_22kto1k.pth"
-)
+# PERSISTENT_DIR = "/data/models"
+# cfg.MODEL.COND_STAGE_CONFIG.PRETRAINED_PATH = os.path.join(
+#     PERSISTENT_DIR,
+#     "pretrained_models",
+#     "swin",
+#     "swin_base_patch4_window12_384_22kto1k.pth"
+# )
 
 from lr_scheduler import LinearWarmupMultiStepDecayLRScheduler
 from models import (AppearanceEncoder, Decoder, PoseEncoder, UNet,