temp

jolibrain · Jul 18, 2023 · 524840e · 524840e
1 parent 6c92cab
commit 524840e
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/models/modules/diffusion_generator.py b/models/modules/diffusion_generator.py
@@ -100,6 +100,7 @@ def restoration(
                 sample_num=sample_num,
                 cls=cls,
                 guidance_scale=guidance_scale,
+                ref=ref,
             )
         elif self.sampling_method == "ddim":
             return self.restoration_ddim(
@@ -118,12 +119,13 @@ def restoration(
     def restoration_ddpm(
         self,
         y_cond,
-        y_t=None,
-        y_0=None,
-        mask=None,
-        sample_num=8,
-        cls=None,
-        guidance_scale=0.0,
+        y_t,
+        y_0,
+        mask,
+        sample_num,
+        cls,
+        guidance_scale,
+        ref,
     ):
         phase = "test"