137 lines
4.3 KiB
Python
137 lines
4.3 KiB
Python
# SimpleSyrup - workflow-focused ComfyUI extensions for image generation
|
|
# Copyright (C) 2026 Artificial Sweetener and contributors
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
"""Resolve the BERT text encoder used by GroundingDINO."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from types import ModuleType
|
|
|
|
from .model_catalog import BERT_ENTRY
|
|
from .model_downloads import (
|
|
DownloadRequest,
|
|
ModelDownloader,
|
|
ProgressReporter,
|
|
)
|
|
from .model_folders import get_primary_model_folder
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BertResolution:
|
|
"""Resolved local BERT directory and how it was obtained."""
|
|
|
|
path: Path
|
|
source: str
|
|
downloaded: bool
|
|
|
|
|
|
class BertResolver:
|
|
"""Find or download a Hugging Face-style BERT directory."""
|
|
|
|
def __init__(
|
|
self,
|
|
downloader: ModelDownloader | None = None,
|
|
folder_paths_module: ModuleType | None = None,
|
|
) -> None:
|
|
"""Create a resolver with injectable runtime dependencies."""
|
|
|
|
self._downloader = downloader or ModelDownloader()
|
|
self._folder_paths_module = folder_paths_module
|
|
|
|
def resolve(
|
|
self,
|
|
auto_download: bool,
|
|
progress: ProgressReporter | None = None,
|
|
) -> BertResolution:
|
|
"""Return a usable BERT directory or raise an actionable error."""
|
|
|
|
layerstyle_path = self._models_dir() / "bert-base-uncased"
|
|
if is_valid_bert_directory(layerstyle_path):
|
|
return BertResolution(
|
|
path=layerstyle_path,
|
|
source="models/bert-base-uncased",
|
|
downloaded=False,
|
|
)
|
|
|
|
text_encoder_path = self._text_encoder_bert_path()
|
|
if is_valid_bert_directory(text_encoder_path):
|
|
return BertResolution(
|
|
path=text_encoder_path,
|
|
source="models/text_encoders/bert",
|
|
downloaded=False,
|
|
)
|
|
|
|
if not auto_download:
|
|
raise FileNotFoundError(
|
|
"BERT text encoder was not found. Checked: "
|
|
f"{layerstyle_path}; {text_encoder_path}. "
|
|
"Enable auto_download or install a Hugging Face "
|
|
"bert-base-uncased snapshot."
|
|
)
|
|
|
|
self._download_bert(text_encoder_path, progress)
|
|
if not is_valid_bert_directory(text_encoder_path):
|
|
raise FileNotFoundError(
|
|
f"Downloaded BERT files in '{text_encoder_path}' are incomplete."
|
|
)
|
|
return BertResolution(
|
|
path=text_encoder_path,
|
|
source="downloaded: google-bert/bert-base-uncased",
|
|
downloaded=True,
|
|
)
|
|
|
|
def _download_bert(
|
|
self,
|
|
target_directory: Path,
|
|
progress: ProgressReporter | None,
|
|
) -> None:
|
|
"""Download the known BERT artifact set into the target directory."""
|
|
|
|
for artifact in BERT_ENTRY.artifacts:
|
|
destination = target_directory / artifact.filename
|
|
self._downloader.download(
|
|
DownloadRequest(
|
|
source_url=artifact.source_url,
|
|
destination_path=destination,
|
|
expected_folder=target_directory,
|
|
description=artifact.description,
|
|
),
|
|
progress,
|
|
)
|
|
|
|
def _text_encoder_bert_path(self) -> Path:
|
|
"""Return SimpleSyrup's BERT path under ComfyUI text encoders."""
|
|
|
|
return (
|
|
get_primary_model_folder("text_encoders", self._folder_paths_module)
|
|
/ "bert"
|
|
)
|
|
|
|
def _models_dir(self) -> Path:
|
|
"""Return ComfyUI's models directory."""
|
|
|
|
import importlib
|
|
|
|
folder_paths = self._folder_paths_module or importlib.import_module(
|
|
"folder_paths"
|
|
)
|
|
return Path(str(folder_paths.models_dir))
|
|
|
|
|
|
def is_valid_bert_directory(path: Path) -> bool:
|
|
"""Return whether a directory has enough files to load BERT locally."""
|
|
|
|
if not path.is_dir():
|
|
return False
|
|
has_config = (path / "config.json").is_file()
|
|
has_tokenizer = (path / "tokenizer.json").is_file() or (
|
|
path / "vocab.txt"
|
|
).is_file()
|
|
has_weights = (path / "model.safetensors").is_file() or (
|
|
path / "pytorch_model.bin"
|
|
).is_file()
|
|
return has_config and has_tokenizer and has_weights
|