start simple

wvvva · Aug 3, 2022 · 51cd488 · 51cd488
1 parent 250e040
commit 51cd488
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 10 deletions.
diff --git a/configs/stable-diffusion/inpainting/v1-edgeinpainting.yaml b/configs/stable-diffusion/inpainting/v1-edgeinpainting.yaml
@@ -17,10 +17,10 @@ model:
     scale_factor: 0.18215
     ckpt_path: "/fsx/stable-diffusion/stable-diffusion/checkpoints/v1pp/v1pp-flatlined-hr.ckpt"
 
-    concat_keys:
-      - mask
-      - masked_image
-      - smoothing_strength
+    #concat_keys:
+    #  - mask
+    #  - masked_image
+    #  - smoothing_strength
 
     c_concat_log_start: 1
     c_concat_log_end: 5
@@ -38,7 +38,7 @@ model:
       target: ldm.modules.diffusionmodules.openaimodel.UNetModel
       params:
         image_size: 32 # unused
-        in_channels: 10  # 4 data + 4 downscaled image + 1 mask + 1 strength
+        in_channels: 9  # 4 data + 4 downscaled image + 1 mask
         out_channels: 4
         model_channels: 320
         attention_resolutions: [ 4, 2, 1 ]

diff --git a/ldm/data/laion.py b/ldm/data/laion.py
@@ -264,8 +264,8 @@ def __init__(self, mode="512train", mask_edges=True):
         super().__init__()
         assert mode in list(MASK_MODES.keys()), f'unknown mask generation mode "{mode}"'
         self.make_mask = MASK_MODES[mode]
-        self.n_down_choices = [0, 1, 2]
-        self.sigma_choices = [1, 2, 3, 4, 5]
+        self.n_down_choices = [0]
+        self.sigma_choices = [1]
         self.mask_edges = mask_edges
 
     @torch.no_grad()
@@ -285,7 +285,7 @@ def __call__(self, sample):
         n_choices = len(self.n_down_choices)*len(self.sigma_choices)
         raveled_idx = np.ravel_multi_index((n_down_idx, sigma_idx),
                                            (len(self.n_down_choices), len(self.sigma_choices)))
-        normalized_idx = raveled_idx/(n_choices-1)
+        normalized_idx = raveled_idx/max(1, n_choices-1)
 
         n_down = self.n_down_choices[n_down_idx]
         sigma = self.sigma_choices[sigma_idx]
@@ -322,6 +322,7 @@ def __call__(self, sample):
             sample['masked_image'] = y * (mask < 0.5)
         else:
             sample['masked_image'] = y
+            sample['mask'] = torch.zeros_like(sample['mask'])
 
         # concat normalized idx
         sample['smoothing_strength'] = torch.ones_like(sample['mask'])*normalized_idx

diff --git a/scripts/slurm/v1_edgeinpainting/launcher.sh b/scripts/slurm/v1_edgeinpainting/launcher.sh
@@ -23,6 +23,9 @@ cd /fsx/stable-diffusion/stable-diffusion
 
 CONFIG="/fsx/stable-diffusion/stable-diffusion/configs/stable-diffusion/inpainting/v1-edgeinpainting.yaml"
 
+# start without masking
+EXTRA="data.params.train.postprocess.params.mask_edges=false"
+
 # resume and set new seed to reshuffle data
 #EXTRA="--seed 543 --resume_from_checkpoint ..."
 
@@ -38,4 +41,4 @@ CONFIG="/fsx/stable-diffusion/stable-diffusion/configs/stable-diffusion/inpainti
 # detect bad gpus early on
 /bin/bash /fsx/stable-diffusion/stable-diffusion/scripts/test_gpu.sh
 
-python main.py --base $CONFIG --gpus 0,1,2,3,4,5,6,7 -t --num_nodes ${WORLD_SIZE} --scale_lr False
+python main.py --base $CONFIG --gpus 0,1,2,3,4,5,6,7 -t --num_nodes ${WORLD_SIZE} --scale_lr False $EXTRA
diff --git a/scripts/slurm/v1_edgeinpainting/sbatch.sh b/scripts/slurm/v1_edgeinpainting/sbatch.sh
@@ -1,7 +1,7 @@
 #!/bin/bash
 #SBATCH --partition=compute-od-gpu
 #SBATCH --job-name=stable-diffusion-v1-edgeinpainting
-#SBATCH --nodes 24
+#SBATCH --nodes 8
 #SBATCH --ntasks-per-node 1
 #SBATCH --cpus-per-gpu=4
 #SBATCH --gres=gpu:8