From fd29c628cfdaad871e990e7b2a60691aeeff2799 Mon Sep 17 00:00:00 2001
From: Rey Christian <rcsaquino.md@gmail.com>
Date: Tue, 19 Dec 2023 01:18:17 +0800
Subject: [PATCH] Fixed broken models for bg remover

---
 nodes.py | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/nodes.py b/nodes.py
index 8cfcedd..73c2a6d 100644
--- a/nodes.py
+++ b/nodes.py
@@ -58,16 +58,20 @@ class vae_loader:
 class background_remover:
     @classmethod
     def INPUT_TYPES(cls):
-        models = []
-        for session in rembg.session_factory.sessions_class:
-            model_name = session.name()
-            if not model_name == "u2net_custom":
-                models.append(model_name)
-
         return {
             "required": {
                 "image": ("IMAGE",),
-                "model": (models,),
+                "model": (
+                    [
+                        "u2net",
+                        "u2netp",
+                        "u2net_human_seg",
+                        "u2net_cloth_seg",
+                        "silueta",
+                        "isnet-general-use",
+                        "isnet-anime",
+                    ],
+                ),
             }
         }