fix: inference_num for reference img

jolibrain · Sep 14, 2023 · e444ba1 · e444ba1
1 parent bd51510
commit e444ba1
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/models/palette_model.py b/models/palette_model.py
@@ -656,7 +656,9 @@ def inference(self):
                         cls=cur_class,
                         ddim_num_steps=self.ddim_num_steps,
                         ddim_eta=self.ddim_eta,
-                        ref=self.ref_A if hasattr(self, "ref_A") else None,
+                        ref=self.ref_A[: self.inference_num]
+                        if hasattr(self, "ref_A")
+                        else None,
                     )
 
                     name = "output_" + str(i + 1)