fix: inference on ref cond

jolibrain · Jul 4, 2023 · 13bd2d3 · 13bd2d3
1 parent 350578e
commit 13bd2d3
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/models/palette_model.py b/models/palette_model.py
@@ -271,7 +271,10 @@ def __init__(self, opt, rank):
             for i in range(self.nb_classes_inference):
                 self.gen_visual_names.append("output_" + str(i + 1) + "_")
 
-        elif self.opt.alg_palette_cond_image_creation == "ref":
+        elif (
+            self.opt.alg_palette_cond_image_creation == "ref"
+            or "ref" in self.opt.alg_palette_conditioning
+        ):
             for i in range(self.inference_num):
                 self.gen_visual_names.append("cond_ref_" + str(i + 1) + "_")
                 self.gen_visual_names.append("output_" + str(i + 1) + "_")