세그먼트 모델 변경을 어떻게 해야 하나요.

코드를 확인해보니 현재 세그먼트 모델이 sam_vit_b를 쓰는 걸로 확인이 되는데 모델이 작아서 그런지 세그먼트가 잘 안됄때가 있어서 sam_vit_h 버전을 써보려고 하는데 어디를 손봐야 하는지 감이 안 잡히네요.

변경 코드는 아래와 같습니다. base/sam_hq.py :

def sam_init():
    #MODEL_TYPE = 'vit_b'
    MODEL_TYPE = 'vit_h'

    global sam_model
    if not sam_model:
        torch.load = unsafe_torch_load
        #sam_model = sam_model_registry[MODEL_TYPE](checkpoint='%s/sam_vit_b_01ec64.pth' % bmab_model_path)
        sam_model = sam_model_registry[MODEL_TYPE](checkpoint='%s/sam_vit_h_4b8939.pth' % bmab_model_path)
        sam_model.to(device=device)
        sam_model.eval()
        torch.load = load

    return sam_model

masking/sam.py:

class SamPredictVitB(SamPredict):

    @property
    def type(self):
        return 'vit_h'
        #return 'vit_b'

    @property
    def file(self):
        return 'sam_vit_h_4b8939.pth'
        #return 'sam_vit_b_01ec64.pth'

masking/sam_hq.py:

class SamPredictVitB(SamPredict):

    @property
    def type(self):
        return 'vit_h'
        #return 'vit_b'

    @property
    def file(self):
        return 'sam_vit_h_4b8939.pth'
        #return 'sam_vit_b_01ec64.pth'

잘못 고쳤는지 아래와 같은 에러가 발생하고요.

*** Error running before_process: extensions/sd-webui-bmab/scripts/sd_webui_bmab.py
    Traceback (most recent call last):
      File "modules/scripts.py", line 611, in before_process
        script.before_process(p, *script_args)
      File "extensions/sd-webui-bmab/sd_bmab/bmab.py", line 63, in before_process
        internal.process_img2img(ctx)
      File "extensions/sd-webui-bmab/sd_bmab/pipeline/internal/internalpipe.py", line 72, in process_img2img
        ctx.sdprocessing.init_images[0] = process_intermediate_step2(ctx, image)
      File "extensions/sd-webui-bmab/sd_bmab/pipeline/internal/internalpipe.py", line 61, in process_intermediate_step2
        ret = proc.process(context, processed)
      File "extensions/sd-webui-bmab/sd_bmab/processors/basic/img2imgmasking.py", line 37, in process
        mask = self.sam(context, self.prompt, context.sdprocessing.init_images[0])
      File "extensions/sd-webui-bmab/sd_bmab/processors/basic/img2imgmasking.py", line 28, in sam
        sam = masking.get_mask_generator()
      File "extensions/sd-webui-bmab/sd_bmab/masking/__init__.py", line 14, in get_mask_generator
        model = dmasks.get(name, dmasks[shared.opts.bmab_mask_model])
    KeyError: 'sam_vit_b'

---

portu-sim / sd-webui-bmab

세그먼트 모델 변경을 어떻게 해야 하나요. #25