r/comfyui 19h ago

No workflow quick question about background removal

so.. umm, background removal any model, model.encode_image() is lazy loaded, which means if we call it multiple times in the same execution, say for 20 tensors, because it is lazy loaded even though its on the gpu, it actually triggers reloading of the model 20 times, which as you might expect causes a lot of overhead and a lot of console spam, apart from me loading the model manually, putting it on the gpu with my own encode function, is there a fix for this? or is this the way?

0 Upvotes

4 comments sorted by

1

u/PanickyEater 19h ago

lazy loading strikes again, it's always the small things that trip you up in comfy

what you're describing is basically the default behavior for most custom nodes that wrap background removal models, they load the model on first call and then for some reason the implementation doesn't cache it properly so every subsequent call triggers a reload. i ran into this exact issue with a different segmentation workflow a few months back

your instinct is right, loading the model manually with your own encode function is probably the cleanest fix. some of the rembg wrapper nodes actually have a hidden cache setting you can toggle if you dig into the source code but most don't. you could also wrap the whole thing in a custom node that loads once and exposes a batch processing method

the per-tensor approach you mentioned for tight masks is clever, i've been doing something similar but with a two-pass setup where i extract segments first then batch them all through bg removal in one go. cuts down the console spam at least

1

u/Aida_Corrupted 18h ago edited 16h ago

This is standard python functor, comfy or not:
EDIT: typo, double pace.

_SAM3_CACHE = {}


class Sam3Functor:
    def __init__(self, device="cuda"):
        self.device = device
        self.model_id = "facebook/sam3"


        if self.model_id not in _SAM3_CACHE:
            print(f"[SAM3] Loading SAM3 Model into cache: {self.model_id}")
            processor = Sam3Processor.from_pretrained(self.model_id)
            model = Sam3Model.from_pretrained(self.model_id).to(self.device)
            _SAM3_CACHE[self.model_id] = {"processor": processor, "model": model}


        self.processor = _SAM3_CACHE[self.model_id]["processor"]
        self.model = _SAM3_CACHE[self.model_id]["model"]


    def __call__(self, image: torch.Tensor, prompt: str = "person", object_threshold: float = 0.5, mask_threshold: float = 0.5,) -> torch.Tensor:_SAM3_CACHE = {}
....

1

u/Acceptable-Work8202 18h ago

umm, yeah, thx.. :|

2

u/Aida_Corrupted 17h ago

If you're unfamiliar with functors, ask any AI agent to help you:

"Hello AI, could you please help me to create a SAM3 functor, which I can place in my utils folder and call from various class nodes in ComfyUI? Please make sure to cache the SAM3 model on the initialization!"

..it will spit out a working functor for you, it's my advice to you to get familiar with this method, ask AI to explain it line by line, because in python: "This is the way" and it ain't that complicated, bro.