Fix issue when not being able to display mask when not returning fore…

…ground image
dimitribarbot · Sep 1, 2024 · 55f949e · 55f949e
1 parent 552e0c5
commit 55f949e
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/scripts/postprocessing_rembg.py b/scripts/postprocessing_rembg.py
@@ -29,7 +29,7 @@ def get_pipeline_using_cache(model_name: BiRefNetModelName):
     return birefnet
 
 
-class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing):
+class ScriptPostprocessingBiRefNet(scripts_postprocessing.ScriptPostprocessing):
     name = "BiRefNet"
     order = 20001
     model = None
@@ -77,9 +77,9 @@ def process(self, pp: scripts_postprocessing.PostprocessedImage, enable, model,
             edge_mask_width=edge_mask_width
         )
 
-        pp.image = output_image
-        if mask:
-            pp.extra_images.append(mask)
+        pp.image = mask
+        if output_image:
+            pp.extra_images.append(output_image)
         if edge_mask:
             pp.extra_images.append(edge_mask)