127 lines
5.9 KiB
Python
127 lines
5.9 KiB
Python
import os
|
|
from pathlib import Path
|
|
from PIL import Image
|
|
import numpy as np
|
|
import torch
|
|
from typing import Tuple, Union, Dict, Any
|
|
import comfy.model_management as mm
|
|
|
|
from .image_manager import JtpImageManager
|
|
|
|
from ..redrocket.tag_manager import JtpTagManager
|
|
from ..redrocket.model_manager import JtpModelManager
|
|
from ..helpers.cache import ComfyCache
|
|
class JtpInference:
|
|
"""
|
|
A Clip Vision Classifier by RedRocket (inference code made robust by deitydurg)
|
|
"""
|
|
|
|
@classmethod
|
|
def run_classifier(
|
|
cls,
|
|
model_name: str,
|
|
tags_name: str,
|
|
device: torch.device,
|
|
image: Union[Image.Image, np.ndarray, Path],
|
|
steps: float,
|
|
threshold: float,
|
|
exclude_tags: str = "",
|
|
replace_underscore: bool = True,
|
|
trailing_comma: bool = False,
|
|
implications_mode: str = "off",
|
|
exclude_categories: str = "",
|
|
prefix: str = "",
|
|
keep_model_loaded: bool = False,
|
|
) -> Tuple[str, Dict[str, float]]:
|
|
from ..helpers.logger import ComfyLogger
|
|
|
|
model_version: int = ComfyCache.get(f'config.models.{model_name}.version')
|
|
tags_version: int = ComfyCache.get(f'config.tags.{tags_name}.version')
|
|
ComfyLogger().log(message=f"Model version: {model_version}, Tags version: {tags_version}", type="DEBUG", always=True)
|
|
|
|
# Load all the things
|
|
if JtpModelManager().is_installed(model_name) is False:
|
|
if not JtpModelManager().download(model_name):
|
|
ComfyLogger().log(f"Model {model_name} could not be downloaded", "ERROR", True)
|
|
return "", {}
|
|
if JtpTagManager().is_installed(tags_name) is False:
|
|
if not JtpTagManager().download(tags_name):
|
|
ComfyLogger().log(f"Tags {tags_name} could not be downloaded", "ERROR", True)
|
|
return "", {}
|
|
if JtpModelManager().is_loaded(model_name) is False:
|
|
if not JtpModelManager().load(model_name=model_name, version=model_version, device=device):
|
|
ComfyLogger().log(f"Model {model_name} could not be loaded", "ERROR", True)
|
|
return "", {}
|
|
if JtpTagManager().is_loaded(tags_name) is False:
|
|
if not JtpTagManager().load(tags_name=tags_name, version=tags_version):
|
|
ComfyLogger().log(f"Tags {tags_name} could not be loaded", "ERROR", True)
|
|
return "", {}
|
|
|
|
# Load implications metadata if needed
|
|
if implications_mode != "off":
|
|
if not JtpTagManager().has_metadata(tags_name):
|
|
# Download CSV if not on disk yet
|
|
csv_path = os.path.join(JtpTagManager().tags_basepath, f"{tags_name}-tags.csv")
|
|
if not os.path.exists(csv_path):
|
|
JtpTagManager().download_metadata(model_name)
|
|
JtpTagManager().load_metadata(tags_name)
|
|
|
|
# Create a unique key for the parameters to handle caching correctly
|
|
import hashlib
|
|
params_string = f"{model_name}|{tags_name}|{steps}|{threshold}|{exclude_tags}|{replace_underscore}|{trailing_comma}|{implications_mode}|{exclude_categories}|{prefix}"
|
|
params_key = hashlib.sha256(params_string.encode()).hexdigest()
|
|
|
|
tensor = JtpImageManager().load(image=image, device=device, params_key=params_key)
|
|
if tensor is None:
|
|
ComfyLogger().log(message="Image could not be loaded", type="ERROR", always=True)
|
|
return "", {}
|
|
if isinstance(tensor, tuple):
|
|
ComfyLogger().log(message="Returning cached result", type="DEBUG", always=True)
|
|
return tensor[0], tensor[1]
|
|
ComfyLogger().log(message=f"Classifying image with model {model_name} and tags {tags_name}", type="INFO", always=True)
|
|
model_data: Union[torch.nn.Module, None] = ComfyCache.get(f'model.{model_name}.model')
|
|
if model_data is None:
|
|
ComfyLogger().log(message=f"Model data for {model_name} not found in cache", type="ERROR", always=True)
|
|
return "", {}
|
|
# Move model to inference device (may have been offloaded to CPU)
|
|
model_data.to(device)
|
|
with torch.no_grad():
|
|
if f"{model_version}" == "1":
|
|
logits = model_data(tensor)
|
|
probits = torch.nn.functional.sigmoid(logits[0]).cpu()
|
|
values, indices = probits.topk(250)
|
|
elif f"{model_version}" == "2":
|
|
probits = model_data(tensor)[0].cpu()
|
|
values, indices = probits.topk(250)
|
|
else:
|
|
ComfyLogger().log(message=f"Model version {model_version} not supported", type="ERROR", always=True)
|
|
return "", {}
|
|
# Memory management based on keep_model_loaded toggle
|
|
if keep_model_loaded:
|
|
# Offload model from GPU to CPU (RAM) — keep cached for fast reuse
|
|
model_data.to(torch.device("cpu"))
|
|
mm.soft_empty_cache()
|
|
else:
|
|
# Fully unload: drop local ref, flush cache, GC, free VRAM
|
|
del model_data
|
|
JtpModelManager.unload(model_name)
|
|
|
|
ComfyLogger().log(message="Processing tags...", type="DEBUG", always=True)
|
|
tags_str, tag_scores = JtpTagManager().process_tags(
|
|
tags_name=tags_name,
|
|
values=values,
|
|
indices=indices,
|
|
threshold=threshold,
|
|
exclude_tags=exclude_tags,
|
|
replace_underscore=replace_underscore,
|
|
trailing_comma=trailing_comma,
|
|
implications_mode=implications_mode,
|
|
exclude_categories=exclude_categories,
|
|
prefix=prefix,
|
|
)
|
|
if not JtpImageManager().commit_cache(image=image, output=(tags_str, tag_scores), params_key=params_key):
|
|
ComfyLogger().log(message="Image cache could not be committed", type="WARN", always=True)
|
|
return tags_str, tag_scores
|
|
ComfyLogger().log(message=f"Classification complete: {tags_str}", type="INFO", always=True)
|
|
return tags_str, tag_scores
|