nodes/chflame163-ComfyUI_LayerStyle
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: improve model loading logic for existing BiRefNet-General model

Browse Source
This commit is contained in:
Jim Lee
2024-11-14 03:30:48 +08:00
parent c377306fe3
commit 26db8262a8
1 changed files with 12 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/birefnet_ultra_v2.py
+12 -11
View File
@@ -80,18 +80,19 @@ class LS_LoadBiRefNetModelV2:
os.makedirs(birefnet_path, exist_ok=True)
model_path = os.path.join(birefnet_path, version)
old_model_path = os.path.join(birefnet_path, 'pth')
if version == "BiRefNet-General" and os.path.exists(old_model_path):
model = "BiRefNet-general-epoch_244.pth"
from .BiRefNet_v2.models.birefnet import BiRefNet
from .BiRefNet_v2.utils import check_state_dict
model_dict = get_models()
self.birefnet = BiRefNet(bb_pretrained=False)
self.state_dict = torch.load(model_dict[model], map_location='cpu', weights_only=True)
self.state_dict = check_state_dict(self.state_dict)
self.birefnet.load_state_dict(self.state_dict)
return (self.birefnet,)
if version == "BiRefNet-General":
old_birefnet_path = os.path.join(birefnet_path, 'pth')
old_model = "BiRefNet-general-epoch_244.pth"
old_model_path = os.path.join(old_birefnet_path, old_model)
if os.path.exists(old_model_path):
from .BiRefNet_v2.models.birefnet import BiRefNet
from .BiRefNet_v2.utils import check_state_dict
self.birefnet = BiRefNet(bb_pretrained=False)
self.state_dict = torch.load(old_model_path, map_location='cpu', weights_only=True)
self.state_dict = check_state_dict(self.state_dict)
self.birefnet.load_state_dict(self.state_dict)
return (self.birefnet,)
elif not os.path.exists(model_path):
log(f"Downloading {version} model...")
repo_id = self.birefnet_model_repos[version]