AlbeRota commited on
Commit
955ea7d
·
verified ·
1 Parent(s): b012b06

Upload weights, notebooks, sample images

Browse files
configs/end2end02.yaml ADDED
@@ -0,0 +1,280 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ### BASELINE: CONVERGES AFTER LONG
2
+
3
+ parameters:
4
+
5
+ ### MODEL ARCHITECTURE
6
+ MODEL:
7
+ value:
8
+ MODEL_CLASS: "UnReflect_Model_TokenInpainter" # Main model class name (must match class in models.py)
9
+ MODEL_MODULE: "models" # Module name to import model classes from (default: "models")
10
+ RGB_ENCODER:
11
+ ENCODER: "facebook/dinov3-vitl16-pretrain-lvd1689m" # DINOv3 encoder model name (HuggingFace format)
12
+ IMAGE_SIZE: 896 # Input image size (height and width in pixels)
13
+ RETURN_SELECTED_LAYERS: [3, 6, 9 , 12] # Transformer layer indices to extract features from (0-indexed)
14
+ RGB_ENCODER_LR: 0.0 # Learning rate for RGB encoder (0.0 = frozen, must be explicitly set)
15
+ DECODERS:
16
+ diffuse:
17
+ FEATURE_DIM: 1024 # Feature dimension for decoder (should match encoder output)
18
+ REASSEMBLE_OUT_CHANNELS: [768,1024,1536,2048] # Output channels for each decoder stage (DPT-style reassembly)
19
+ REASSEMBLE_FACTORS: [4.0, 2.0, 1.0, 0.5] # Spatial upsampling factors for each stage
20
+ READOUT_TYPE: "ignore" # Readout type for DPT decoder ("ignore", "project", etc.)
21
+ FROM_PRETRAINED: "diffuse_decoder.pt" # Path to pretrained decoder weights (optional)
22
+ USE_BN: False # Use batch normalization in decoder
23
+ DROPOUT: 0.1 # Dropout rate in decoder layers
24
+ OUTPUT_IMAGE_SIZE: [896,896] # Output image resolution [height, width]
25
+ OUTPUT_CHANNELS: 3 # Number of output channels (3 for RGB diffuse image)
26
+ DECODER_LR: 1.0e-5 # Custom learning rate for decoder (0.0 = frozen, 1.0 = same as base LR)
27
+ NUM_FUSION_BLOCKS_TRAINABLE: 1 # Number of fusion blocks to train (0-4, null = train all if DECODER_LR != 0)
28
+ TRAIN_RGB_HEAD: True # Whether to train RGB head (true/false, null = train if DECODER_LR != 0)
29
+ highlight:
30
+ FEATURE_DIM: 1024 # Feature dimension for highlight decoder
31
+ REASSEMBLE_OUT_CHANNELS: [96,192,384,768] # Output channels for each decoder stage
32
+ REASSEMBLE_FACTORS: [4.0, 2.0, 1.0, 0.5] # Spatial upsampling factors for each stage
33
+ READOUT_TYPE: "ignore" # Readout type for DPT decoder
34
+ FROM_PRETRAINED: "highlight_decoder.pt" # Path to pretrained decoder weights (optional)
35
+ USE_BN: False # Use batch normalization in decoder
36
+ DROPOUT: 0.1 # Dropout rate in decoder layers
37
+ OUTPUT_IMAGE_SIZE: [896,896] # Output image resolution [height, width]
38
+ OUTPUT_CHANNELS: 1 # Number of output channels (1 for highlight mask)
39
+ DECODER_LR: 0.0 #1.0e-5 # Custom learning rate for decoder (0.0 = frozen, 1.0 = same as base LR)
40
+ NUM_FUSION_BLOCKS_TRAINABLE: null # Number of fusion blocks to train (0-4, null = train all if DECODER_LR != 0)
41
+ TOKEN_INPAINTER:
42
+ TOKEN_INPAINTER_CLASS: "TokenInpainter_Prior" # Token inpainter class name
43
+ TOKEN_INPAINTER_MODULE: "token_inpainters" # Module name to import token inpainter from
44
+ # FROM_PRETRAINED: "token_inpainter.pth" # Path to pretrained token inpainter weights (optional)
45
+ TOKEN_INPAINTER_LR: 1.0e-4 # Learning rate for token inpainter (can differ from base LR)
46
+ DEPTH: 6 # Number of transformer blocks
47
+ HEADS: 16 # Number of attention heads
48
+ DROP: 0.05 # Dropout rate
49
+ USE_POSITIONAL_ENCODING: True # Enable 2D sinusoidal positional encodings
50
+ USE_FINAL_NORM: True # Enable final LayerNorm before output projection
51
+ USE_LOCAL_PRIOR: True # Blend local mean prior for masked seeds
52
+ LOCAL_PRIOR_WEIGHT: 0.2 # Weight for local prior blending (1.0 = only mask_token, 0.0 = only local mean)
53
+ LOCAL_PRIOR_KERNEL: 5 # Kernel size for local prior blending (> 1)
54
+ SEED_NOISE_STD: 0.02 # Standard deviation of noise added to masked seeds during training
55
+
56
+ # FORWARD PASS PARAMETERS
57
+ INPAINT_MASK_THRESHOLD:
58
+ value: 0.2 # Threshold for inpaint mask
59
+ INPAINT_MASK_DILATION:
60
+ value: 40 # Dilation kernel size (pixels) for inpaint mask - Must be odd
61
+
62
+ USE_TORCH_COMPILE: # Enable PyTorch 2.0 torch.compile for faster training (experimental)
63
+ value: False
64
+ DISTRIBUTE:
65
+ value: "ddp"
66
+
67
+ ### DATA
68
+ DATASETS:
69
+ value:
70
+ # Reserved key: key-value pairs here override the same keys for every dataset (per-dataset entries still override this).
71
+ SCRREAM:
72
+ VAL_SCENES: ["scene10", "scene04"] # List of validation scene names
73
+ RESIZE_MODE: "resize+crop" # Image resizing mode: "resize", "crop", "resize+crop", or "pad"
74
+ SAMPLE_EVERY_N: 6 # Load every Nth frame from each scene (1 = all frames, 4 = every 4th frame)
75
+
76
+ HOUSECAT6D:
77
+ VAL_SCENES: ["val_scene1","val_scene2"] # Validation scene names
78
+ RESIZE_MODE: "resize+crop" # Image resizing mode
79
+ SAMPLE_EVERY_N: 4 # Load every Nth frame
80
+
81
+ CROMO:
82
+ TRAIN_SCENES: ["kitchen"] # Training scene names (list or string)
83
+ # VAL_SCENES: "station" # Validation scene names (optional)
84
+ RESIZE_MODE: "resize" # Image resizing mode
85
+ SAMPLE_EVERY_N: 2 # Load every Nth frame
86
+
87
+ PSD:
88
+ TRAIN_SCENES: "PSD_Train" # Training scene name (string or list)
89
+ VAL_SCENES: "PSD_Val" # Validation scene name (string or list)
90
+ RESIZE_MODE: "resize+crop" # Image resizing mode
91
+ SAMPLE_EVERY_N: 1 # Load every Nth frame (1 = all frames)
92
+
93
+ SCARED:
94
+ VAL_SCENES: ["v22","v23","v24","v25","v26","v27","v28","v29","v30","v31","v32","v33","v34"] # Validation scene names
95
+ RESIZE_MODE: "resize+crop" # Image resizing mode
96
+ SAMPLE_EVERY_N: 8 # Load every Nth frame
97
+
98
+ STEREOMIS_TRACKING:
99
+ VAL_SCENES: ["P2_2"] # Validation scene names
100
+ RESIZE_MODE: "resize+crop" # Image resizing mode
101
+ SAMPLE_EVERY_N: 2 # Load every Nth frame
102
+
103
+ CHOLEC80:
104
+ TRAIN_SCENES: ["train"] # Validation scene names
105
+ VAL_SCENES: ["test"] # Validation scene names
106
+ RESIZE_MODE: "resize+crop" # Image resizing mode
107
+ SAMPLE_EVERY_N: 40 # Load every Nth frame
108
+
109
+ SUNRGBD:
110
+ VAL_SCENES: ["realsense"] # Validation scene names
111
+ RESIZE_MODE: "resize+crop" # Image resizing mode
112
+ SAMPLE_EVERY_N: 4 # Load every Nth frame
113
+
114
+ SCANNET:
115
+ TRAIN_SCENES: ["train"]
116
+ VAL_SCENES: ["val"]
117
+ RESIZE_MODE: "resize+crop"
118
+ SAMPLE_EVERY_N: 5
119
+
120
+ OPENIMAGESV7:
121
+ TRAIN_SCENES: ["thescene"]
122
+ # VAL_SCENES: [""]
123
+ RESIZE_MODE: "resize+crop"
124
+ SAMPLE_EVERY_N: 5
125
+
126
+ ENDOSYNTH:
127
+ TRAIN_SCENES: ["scene"]
128
+ # VAL_SCENES: ["val"]
129
+ RESIZE_MODE: "resize+crop"
130
+ SAMPLE_EVERY_N: 1
131
+
132
+ ALL_DATASETS:
133
+ FEW_IMAGES: False # Override FEW_IMAGES for all datasets (for quick debugging set True)
134
+ TARGET_SIZE: [896,896] # Override target image size [height, width] for all datasets
135
+ LOAD_RGB_ONLY: True
136
+
137
+ BATCH_SIZE: # Max batch size with img size 896 is 32
138
+ value: 4 # Number of samples per batch (adjust based on GPU memory)
139
+ NUM_WORKERS:
140
+ value: 12 # Number of data loading worker processes (0 = main process only, "auto" = 90% of CPU affinity)
141
+ SHUFFLE:
142
+ value: True # Shuffle training data each epoch (False for validation/test)
143
+ PIN_MEMORY:
144
+ value: True # Pin memory in DataLoader for faster GPU transfer (recommended: True)
145
+ PREFETCH_FACTOR:
146
+ value: 2 # Number of batches to prefetch per worker (higher = more memory usage)
147
+
148
+ ### HIGHLIGHTS
149
+ MOGE_MODEL:
150
+ value: "Ruicheng/moge-2-vits-normal" # MoGe model name for normal estimation (HuggingFace format)
151
+ SURFACE_ROUGHNESS:
152
+ value: [100.0, 2000.0] # Blinn-Phong surface roughness exponent (higher = sharper highlights)
153
+ INTENSITY:
154
+ value: [0.5, 1.0] # Specular highlight intensity multiplier
155
+ LIGHT_DISTANCE_RANGE:
156
+ value: [0.0, 1] # Range for light source distance sampling [min, max] (normalized)
157
+ LIGHT_LEFT_RIGHT_ANGLE:
158
+ value: [0, 360] # Range for light source horizontal angle [min, max] in degrees
159
+ LIGHT_ABOVE_BELOW_ANGLE:
160
+ value: [0, 360] # Range for light source vertical angle [min, max] in degrees
161
+ DATASET_HIGHLIGHT_DILATION:
162
+ value: 25 # Dilation kernel size (pixels) for dataset highlight masks
163
+ DATASET_HIGHLIGHT_THRESHOLD:
164
+ value: 0.85 # Brightness/luminance threshold (0-1) for detecting highlights in dataset images
165
+ DATASET_HIGHLIGHT_USE_LUMINANCE:
166
+ value: True # If True, use perceptually-weighted luminance (0.299*R + 0.587*G + 0.114*B) for dataset highlights; if False, use simple mean brightness
167
+ HIGHLIGHT_COLOR:
168
+ value: [1.0, 1.0, 1.0] # RGB color for synthetic highlights (normalized 0-1)
169
+ CLAMP_RECONSTRUCTION:
170
+ value: True # Clamp reconstructed images to [0, 1] range if True
171
+
172
+ ### OPTIMIZATION
173
+ LEARNING_RATE:
174
+ value: 1.0e-3 # Base learning rate for optimizer # DEPRECATED. SETTING LR FOR EACH MODULE ABOVE
175
+ WEIGHT_DECAY:
176
+ value: 0.0 # L2 regularization weight (0.0 = no weight decay)
177
+ EPOCHS:
178
+ value: 10 # Maximum number of training epochs
179
+ GRADIENT_ACCUMULATION_STEPS:
180
+ value: 1 # Number of steps to accumulate gradients before optimizer step (1 = no accumulation)
181
+ WARMUP:
182
+ value: 100 # Number of warmup steps for learning rate schedule (linear warmup from 0 to LR)
183
+ GRADIENT_CLIPPING_MAX_NORM:
184
+ value: 8 # Maximum gradient norm for clipping (set to -1 to disable clipping)
185
+ LR_SCHEDULER:
186
+ value:
187
+ ONPLATEAU: # ReduceLROnPlateau scheduler (reduces LR when validation metric plateaus)
188
+ PATIENCE: 5 # Number of epochs to wait before reducing LR
189
+ FACTOR: 0.1 # Factor by which LR is reduced (new_lr = old_lr * factor)
190
+ COSINE: # CosineAnnealingLR scheduler (cosine annealing schedule)
191
+ N_PERIODS: 0.5 # Number of cosine periods over training
192
+ # STEPWISE: # StepLR scheduler (reduces LR at fixed step intervals)
193
+ # N_STEPS: 4 # Number of times to reduce LR during training
194
+ # GAMMA: 0.5 # Factor by which LR is reduced at each step (new_lr = old_lr * gamma)
195
+ # EXPONENTIAL: # ExponentialLR scheduler (exponential decay)
196
+ # GAMMA: 0.5 # Multiplicative factor for exponential decay
197
+
198
+ OPTIMIZER_BOOTSTRAP_NAME:
199
+ value: "AdamW" # Optimizer name for initial training phase ("Adam", "SGD", etc.)
200
+ EARLY_STOPPING_PATIENCE:
201
+ value: 10 # Number of epochs without improvement before stopping training
202
+ SAVE_INTERVAL:
203
+ value: 1000 # Number of training steps between model checkpoints
204
+
205
+ DATASET_HIGHLIGHT_SUPERVISION_THRESHOLD:
206
+ value: 0.1 # Pixel highlights above this threshold (should be low) are excluded from supervision
207
+
208
+ ### LOSS WEIGHTS (relative to the total loss, NOT NORMALIZED LATER)
209
+ SPECULAR_LOSS_WEIGHT:
210
+ value: 0.0 # Weight for specular component reconstruction loss
211
+ DIFFUSE_LOSS_WEIGHT:
212
+ value: 1.0 # Weight for diffuse component reconstruction loss
213
+ HIGHLIGHT_LOSS_WEIGHT:
214
+ value: 1.0 # Weight for highlight mask regression loss
215
+ TOKEN_INPAINT_LOSS_WEIGHT:
216
+ value: 1.0 # Weight for token-space inpainting loss (L1 + cosine similarity in feature space)
217
+ IMAGE_RECONSTRUCTION_LOSS_WEIGHT:
218
+ value: 0.0 # Weight for full image reconstruction loss
219
+
220
+ HLREG_W_L1:
221
+ value: 1.0 # Weight for L1 loss in highlight regression
222
+ HLREG_USE_CHARB:
223
+ value: True # Use Charbonnier loss (smooth L1) instead of standard L1 if True
224
+ HLREG_W_DICE:
225
+ value: 0.2 # Weight for Dice loss in highlight regression (for mask overlap)
226
+ HLREG_W_SSIM:
227
+ value: 0.0 # Weight for SSIM loss in highlight regression
228
+ HLREG_W_GRAD:
229
+ value: 0.0 # Weight for gradient loss in highlight regression
230
+ HLREG_W_TV:
231
+ value: 0.0 # Weight for total variation loss in highlight regression
232
+ HLREG_BALANCE_MODE:
233
+ value: "auto" # Class balancing mode for highlight regression: 'none' | 'auto' | 'pos_weight'
234
+ HLREG_POS_WEIGHT:
235
+ value: 1.0 # Positive class weight (used only if BALANCE_MODE == 'pos_weight')
236
+ HLREG_FOCAL_GAMMA:
237
+ value: 2.0 # Focal loss gamma parameter (0.0 = standard BCE, 1.0-2.0 helps with gradient vanishing)
238
+
239
+ WEIGHT_SEAM:
240
+ value: 0.5 # Weight for gradient matching loss on saturation ring
241
+ RING_DILATE_KERNEL:
242
+ value: 17 # Dilation kernel size (odd number) for creating ring mask around highlights
243
+ SEAM_USE_CHARB:
244
+ value: True # Use Charbonnier loss instead of L1 in seam loss (smooth L1 for boundary consistency)
245
+ SEAM_WEIGHT_GRAD:
246
+ value: 0.0 # Weight for gradient matching term inside seam loss (0.0 = disable gradient term)
247
+ TOKEN_FEAT_ALPHA:
248
+ value: 0.5 # Mixing factor for token feature loss: alpha * L1 + (1-alpha) * (1-cosine_sim)
249
+
250
+ ### DIFFUSE HIGHLIGHT PENALTY
251
+ WEIGHT_DIFFUSE_HIGHLIGHT_PENALTY:
252
+ value: 0.75 # Weight for penalty loss on highlights in diffuse decoder output (0.0 = disabled)
253
+ DIFFUSE_HL_THRESHOLD:
254
+ value: 0.8 # Brightness/luminance threshold for detecting highlights in diffuse (0.0-1.0)
255
+ DIFFUSE_HL_USE_CHARB:
256
+ value: True # Use Charbonnier loss instead of L1 for diffuse highlight penalty
257
+ DIFFUSE_HL_PENALTY_MODE:
258
+ value: "brightness" # Penalty mode: "brightness" (penalize brightness/luminance above threshold) or "pixel" (penalize RGB values directly)
259
+ DIFFUSE_HL_TARGET_BRIGHTNESS:
260
+ value: null # Target brightness/luminance for penalized pixels (null = use threshold value)
261
+ DIFFUSE_HL_USE_LUMINANCE:
262
+ value: False # If True, use perceptually-weighted luminance (0.299*R + 0.587*G + 0.114*B); if False, use simple mean brightness
263
+
264
+ ### LOGGING, RESULTS AND WANDB
265
+ LOG_INTERVAL:
266
+ value: 1 # Number of training steps between console log outputs
267
+ WANDB_LOG_INTERVAL:
268
+ value: 1 # Number of training steps between WandB metric logs
269
+ IMAGE_LOG_INTERVAL:
270
+ value: 5 # Number of training steps between image logging to WandB
271
+ NO_WANDB:
272
+ value: False # Disable WandB logging if True (useful for local debugging)
273
+ MODEL_WATCHER_FREQ_WANDB:
274
+ value: 50 # Frequency (in steps) for logging model parameter histograms to WandB
275
+ WANDB_ENTITY:
276
+ value: "unreflect-anything" # WandB organization/entity name
277
+ WANDB_PROJECT:
278
+ value: "UnReflectAnything" # WandB project name
279
+ NOTES:
280
+ value: "LOCAL_PRIOR_WEIGHT 0.2" # Notes/description for this training run
configs/{end2end.yaml → end2end08.yaml} RENAMED
@@ -36,7 +36,7 @@ parameters:
36
  DROPOUT: 0.1 # Dropout rate in decoder layers
37
  OUTPUT_IMAGE_SIZE: [896,896] # Output image resolution [height, width]
38
  OUTPUT_CHANNELS: 1 # Number of output channels (1 for highlight mask)
39
- DECODER_LR: 1.0e-5 # Custom learning rate for decoder (0.0 = frozen, 1.0 = same as base LR)
40
  NUM_FUSION_BLOCKS_TRAINABLE: null # Number of fusion blocks to train (0-4, null = train all if DECODER_LR != 0)
41
  TOKEN_INPAINTER:
42
  TOKEN_INPAINTER_CLASS: "TokenInpainter_Prior" # Token inpainter class name
@@ -135,7 +135,7 @@ parameters:
135
  LOAD_RGB_ONLY: True
136
 
137
  BATCH_SIZE: # Max batch size with img size 896 is 32
138
- value: 4 # Number of samples per batch (adjust based on GPU memory)
139
  NUM_WORKERS:
140
  value: 12 # Number of data loading worker processes (0 = main process only, "auto" = 90% of CPU affinity)
141
  SHUFFLE:
@@ -149,9 +149,9 @@ parameters:
149
  MOGE_MODEL:
150
  value: "Ruicheng/moge-2-vits-normal" # MoGe model name for normal estimation (HuggingFace format)
151
  SURFACE_ROUGHNESS:
152
- value: 100.0 # Blinn-Phong surface roughness exponent (higher = sharper highlights)
153
  INTENSITY:
154
- value: 0.8 # Specular highlight intensity multiplier
155
  LIGHT_DISTANCE_RANGE:
156
  value: [0.0, 1] # Range for light source distance sampling [min, max] (normalized)
157
  LIGHT_LEFT_RIGHT_ANGLE:
@@ -277,4 +277,4 @@ parameters:
277
  WANDB_PROJECT:
278
  value: "UnReflectAnything" # WandB project name
279
  NOTES:
280
- value: "LOCAL_PRIOR_WEIGHT 0.9" # Notes/description for this training run
 
36
  DROPOUT: 0.1 # Dropout rate in decoder layers
37
  OUTPUT_IMAGE_SIZE: [896,896] # Output image resolution [height, width]
38
  OUTPUT_CHANNELS: 1 # Number of output channels (1 for highlight mask)
39
+ DECODER_LR: 0.0 #1.0e-5 # Custom learning rate for decoder (0.0 = frozen, 1.0 = same as base LR)
40
  NUM_FUSION_BLOCKS_TRAINABLE: null # Number of fusion blocks to train (0-4, null = train all if DECODER_LR != 0)
41
  TOKEN_INPAINTER:
42
  TOKEN_INPAINTER_CLASS: "TokenInpainter_Prior" # Token inpainter class name
 
135
  LOAD_RGB_ONLY: True
136
 
137
  BATCH_SIZE: # Max batch size with img size 896 is 32
138
+ value: 4 # Number of samples per batch (adjust based on GPU memory)
139
  NUM_WORKERS:
140
  value: 12 # Number of data loading worker processes (0 = main process only, "auto" = 90% of CPU affinity)
141
  SHUFFLE:
 
149
  MOGE_MODEL:
150
  value: "Ruicheng/moge-2-vits-normal" # MoGe model name for normal estimation (HuggingFace format)
151
  SURFACE_ROUGHNESS:
152
+ value: [100.0, 2000.0] # Blinn-Phong surface roughness exponent (higher = sharper highlights)
153
  INTENSITY:
154
+ value: [0.5, 1.0] # Specular highlight intensity multiplier
155
  LIGHT_DISTANCE_RANGE:
156
  value: [0.0, 1] # Range for light source distance sampling [min, max] (normalized)
157
  LIGHT_LEFT_RIGHT_ANGLE:
 
277
  WANDB_PROJECT:
278
  value: "UnReflectAnything" # WandB project name
279
  NOTES:
280
+ value: "LOCAL_PRIOR_WEIGHT 0.8" # Notes/description for this training run
configs/highlight_decoder_pretrain.yaml CHANGED
@@ -41,7 +41,7 @@ parameters:
41
  LOCAL_PRIOR_KERNEL: 5 # Kernel size for local prior blending (> 1)
42
  SEED_NOISE_STD: 0.02 # Standard deviation of noise added to masked seeds during training
43
  INPAINT_MASK_DILATION:
44
- value: None # Dilation kernel size (pixels) for inpaint mask (None = compute based on image size)
45
  INPAINT_MASK_THRESHOLD:
46
  value: 0.2 # Inpaint selection threshold
47
  USE_TORCH_COMPILE: # Enable PyTorch 2.0 torch.compile for faster training (experimental)
 
41
  LOCAL_PRIOR_KERNEL: 5 # Kernel size for local prior blending (> 1)
42
  SEED_NOISE_STD: 0.02 # Standard deviation of noise added to masked seeds during training
43
  INPAINT_MASK_DILATION:
44
+ value: 40 # Dilation kernel size (pixels) for inpaint mask (None = compute based on image size)
45
  INPAINT_MASK_THRESHOLD:
46
  value: 0.2 # Inpaint selection threshold
47
  USE_TORCH_COMPILE: # Enable PyTorch 2.0 torch.compile for faster training (experimental)
configs/pretrained_config.yaml ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ### BASELINE: CONVERGES AFTER LONG
2
+
3
+ parameters:
4
+
5
+ ### MODEL ARCHITECTURE
6
+ MODEL:
7
+ value:
8
+ MODEL_CLASS: "UnReflect_Model_TokenInpainter" # Main model class name (must match class in models.py)
9
+ MODEL_MODULE: "models" # Module name to import model classes from (default: "models")
10
+ RGB_ENCODER:
11
+ ENCODER: "facebook/dinov3-vitl16-pretrain-lvd1689m" # DINOv3 encoder model name (HuggingFace format)
12
+ IMAGE_SIZE: 896 # Input image size (height and width in pixels)
13
+ RETURN_SELECTED_LAYERS: [3, 6, 9, 12] # Transformer layer indices to extract features from (0-indexed)
14
+ RGB_ENCODER_LR: 0.0 # Learning rate for RGB encoder (0.0 = frozen, must be explicitly set)
15
+ DECODERS:
16
+ diffuse:
17
+ FEATURE_DIM: 1024 # Feature dimension for decoder (should match encoder output)
18
+ REASSEMBLE_OUT_CHANNELS: [768,1024,1536,2048] # Output channels for each decoder stage (DPT-style reassembly)
19
+ REASSEMBLE_FACTORS: [4.0, 2.0, 1.0, 0.5] # Spatial upsampling factors for each stage
20
+ READOUT_TYPE: "ignore" # Readout type for DPT decoder ("ignore", "project", etc.)
21
+ # FROM_PRETRAINED: "diffuse_decoder.pt" # Path to pretrained decoder weights (optional)
22
+ USE_BN: False # Use batch normalization in decoder
23
+ DROPOUT: 0.1 # Dropout rate in decoder layers
24
+ OUTPUT_IMAGE_SIZE: [896,896] # Output image resolution [height, width]
25
+ OUTPUT_CHANNELS: 3 # Number of output channels (3 for RGB diffuse image)
26
+ DECODER_LR: 0.0 # Custom learning rate for decoder (0.0 = frozen, 1.0 = same as base LR)
27
+ NUM_FUSION_BLOCKS_TRAINABLE: 1 # Number of fusion blocks to train (0-4, null = train all if DECODER_LR != 0)
28
+ TRAIN_RGB_HEAD: True # Whether to train RGB head (true/false, null = train if DECODER_LR != 0)
29
+ highlight:
30
+ FEATURE_DIM: 1024 # Feature dimension for highlight decoder
31
+ REASSEMBLE_OUT_CHANNELS: [96,192,384,768] # Output channels for each decoder stage
32
+ REASSEMBLE_FACTORS: [4.0, 2.0, 1.0, 0.5] # Spatial upsampling factors for each stage
33
+ READOUT_TYPE: "ignore" # Readout type for DPT decoder
34
+ # FROM_PRETRAINED: "highlight_decoder.pt" # Path to pretrained decoder weights (optional)
35
+ USE_BN: False # Use batch normalization in decoder
36
+ DROPOUT: 0.1 # Dropout rate in decoder layers
37
+ OUTPUT_IMAGE_SIZE: [896,896] # Output image resolution [height, width]
38
+ OUTPUT_CHANNELS: 1 # Number of output channels (1 for highlight mask)
39
+ DECODER_LR: 5.0e-4 # Custom learning rate for decoder (0.0 = frozen, 1.0 = same as base LR)
40
+ NUM_FUSION_BLOCKS_TRAINABLE: null # Number of fusion blocks to train (0-4, null = train all if DECODER_LR != 0)
41
+ TOKEN_INPAINTER:
42
+ TOKEN_INPAINTER_CLASS: "TokenInpainter_Prior" # Token inpainter class name
43
+ TOKEN_INPAINTER_MODULE: "token_inpainters" # Module name to import token inpainter from
44
+ # FROM_PRETRAINED: "token_inpainter.pth" # Path to pretrained token inpainter weights (optional)
45
+ TOKEN_INPAINTER_LR: 1.0e-4 # Learning rate for token inpainter (can differ from base LR)
46
+ DEPTH: 6 # Number of transformer blocks
47
+ HEADS: 16 # Number of attention heads
48
+ DROP: 0.05 # Dropout rate
49
+ USE_POSITIONAL_ENCODING: True # Enable 2D sinusoidal positional encodings
50
+ USE_FINAL_NORM: True # Enable final LayerNorm before output projection
51
+ USE_LOCAL_PRIOR: True # Blend local mean prior for masked seeds
52
+ LOCAL_PRIOR_WEIGHT: 0.5 # Weight for local prior blending (1.0 = only mask_token, 0.0 = only local mean)
53
+ LOCAL_PRIOR_KERNEL: 5 # Kernel size for local prior blending (> 1)
54
+ SEED_NOISE_STD: 0.02 # Standard deviation of noise added to masked seeds during training