fix: start super-resolution restoration from noise

jolibrain · Sep 22, 2023 · 3bce4a9 · 3bce4a9
1 parent 8f8ccf0
commit 3bce4a9
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 4 deletions.
diff --git a/models/palette_model.py b/models/palette_model.py
@@ -457,7 +457,6 @@ def set_input(self, data):
                 else:
                     self.mask = data["B_label_mask"].to(self.device)
             else:  # e.g. super-resolution
-                self.y_t = data["A"].to(self.device)
                 self.gt_image = data["A"].to(self.device)
                 self.mask = None
 
@@ -746,7 +745,6 @@ def inference(self):
         elif self.task == "super_resolution":
             self.output, self.visuals = netG.restoration(
                 y_cond=self.cond_image[: self.inference_num],
-                y_t=self.cond_image[: self.inference_num],
                 sample_num=self.sample_num,
                 cls=None,
             )

diff --git a/scripts/gen_single_image_diffusion.py b/scripts/gen_single_image_diffusion.py
@@ -469,7 +469,7 @@ def generate(
                 img_tensor.clone().detach(), mask.clone().detach(), {}
             )
     else:
-        y_t = img_tensor.clone().detach()
+        y_t = torch.randn_like(img_tensor)
 
     if opt.alg_palette_cond_image_creation == "previous_frame":
         if previous_frame is not None:

diff --git a/scripts/run_tests.sh b/scripts/run_tests.sh
@@ -75,7 +75,7 @@ if [ $OUT != 0 ]; then
 fi
 
 ###### diffusion super-resolution process test
-python3 -m pytest -p no:cacheprovider -s "${current_dir}/../tests/test_run_sr_diffusion.py" --dataroot "$TARGET_MASK_SEM_DIR"
+python3 -m pytest -p no:cacheprovider -s "${current_dir}/../tests/test_run_sr_diffusion.py" --dataroot "$TARGET_MASK_SEM_DIR/trainA/img"
 OUT=$?
 
 if [ $OUT != 0 ]; then