fix: fix split module feature not used

Author: rubi-du

nodes.py

@@ -1,5 +1,6 @@
 from transformers import AutoModelForImageSegmentation, AutoConfig
 from transformers.dynamic_module_utils import get_class_from_dynamic_module
+from transformers.models.auto.auto_factory import add_generation_mixin_to_remote_model
 import torch
 from torchvision import transforms
 import numpy as np
@@ -147,7 +148,7 @@ def get_device_by_name(device):
             if torch.cuda.is_available():
                 device = "cuda"
             elif torch.backends.mps.is_available():
-                device = "mps"
+                device = "cpu"
             elif torch.xpu.is_available():
                 device = "xpu"
         except:
@@ -352,7 +353,12 @@ def background_remove(self,
                     except Exception as e:
                         print('No need to delete:', e)
 
-                birefnet = AutoModelForImageSegmentation.from_pretrained(local_model_path,trust_remote_code=True, **spare_params)
+                AutoModelForImageSegmentation.register(config.__class__, model_class, exist_ok=True)
+                model_class = add_generation_mixin_to_remote_model(model_class)
+                birefnet = model_class.from_pretrained(
+                    local_model_path, config=config, **spare_params
+                )
+                # birefnet = AutoModelForImageSegmentation.from_pretrained(local_model_path,trust_remote_code=True, **spare_params)
                 if cached:
                     _birefnet_model = birefnet
                     _birefnet_model_name = local_model_path

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-birefnet-super"
 description = "This repository packages the latest BiRefNet model as a ComfyUI node for use, supporting chunked loading on both CPU and GPU, as well as model caching features."
-version = "1.0.6"
+version = "1.0.7"
 license = {file = "LICENSE"}
 
 [project.urls]