399 lines
17 KiB
Python
399 lines
17 KiB
Python
from collections import defaultdict
|
|
from pathlib import Path
|
|
from PIL import Image
|
|
import numpy as np
|
|
import torch
|
|
from typing import Optional, Tuple, Union, Dict, Any, List
|
|
import hashlib
|
|
from torch.nn import Parameter
|
|
import comfy.model_management as mm
|
|
|
|
from .image_v3_manager import JtpImageV3Manager
|
|
from .tag_v3_manager import CATEGORY_ID_TO_NAME, JtpTagV3Manager
|
|
from .model_v3_manager import JtpModelV3Manager
|
|
from ..helpers.cache import ComfyCache
|
|
from ..helpers.logger import ComfyLogger
|
|
|
|
class JtpInferenceV3:
|
|
"""
|
|
Inference for JTP-3 Hydra.
|
|
"""
|
|
|
|
@classmethod
|
|
def run_cam(cls, model, display_image, features, tag_idx, cam_depth, patch_coords, patch_valid):
|
|
"""
|
|
Runs Class Activation Mapping (CAM) for a specific tag.
|
|
"""
|
|
intermediates = features.get("image_intermediates")
|
|
if intermediates is None:
|
|
intermediates = features.get("intermediates")
|
|
|
|
if intermediates is None:
|
|
raise KeyError(f"Features dict missing 'image_intermediates' or 'intermediates'. Keys: {list(features.keys())}")
|
|
|
|
# Handle cam_depth logic relative to available intermediates
|
|
# JTP-3 app.py logic:
|
|
if len(intermediates) > cam_depth:
|
|
intermediates = intermediates[-cam_depth:]
|
|
|
|
# Ensure intermediates is a mutable list for in-place updates
|
|
intermediates = list(intermediates)
|
|
|
|
saved_q = model.attn_pool.q
|
|
saved_p = model.attn_pool.out_proj.weight
|
|
|
|
try:
|
|
# Slice weights for the specific tag
|
|
model.attn_pool.q = Parameter(saved_q[:, [tag_idx], :], requires_grad=False)
|
|
model.attn_pool.out_proj.weight = Parameter(saved_p[[tag_idx], :, :], requires_grad=False)
|
|
|
|
with torch.enable_grad():
|
|
# We need to process intermediates one by one or all together?
|
|
# The original implementation seems to iterate.
|
|
# BUT, forward_head likely expects a specific input structure.
|
|
# Wait, 'intermediate' in forward_head(x) is usually the output of the transformer blocks.
|
|
|
|
# If we iterate, we are doing multiple forward/backward passes.
|
|
# Each intermediate is a tensor [B, N, C].
|
|
|
|
# The issue reported is "element 0 of tensors does not require grad".
|
|
# This happens when we call backward() on a tensor that has no history.
|
|
|
|
for i in range(len(intermediates)):
|
|
# Ensure the input tensor has requires_grad=True and is a leaf tensor for this subgraph.
|
|
# We detach from any previous graph (inference was no_grad anyway).
|
|
intermediate_tensor = intermediates[i].detach()
|
|
intermediate_tensor.requires_grad_(True)
|
|
|
|
# We might need to retain grad if it's not a leaf, but here it IS the leaf of our new graph.
|
|
# However, for safety and access later:
|
|
intermediate_tensor.retain_grad()
|
|
|
|
# Forward pass through the head with sliced weights
|
|
# Note: forward_head usually takes the *last* hidden state.
|
|
# If 'intermediates' are layers, does forward_head handle them all?
|
|
# Assuming forward_head(x) expects the transformer output.
|
|
|
|
logits = model.forward_head(intermediate_tensor, patch_valid=patch_valid)
|
|
|
|
# Backward pass
|
|
# We want to maximize the target class score (logits[0, 0] since we sliced the head to size 1)
|
|
if logits.requires_grad:
|
|
logits[0, 0].backward()
|
|
else:
|
|
# This should not happen if intermediate_tensor requires grad and is used in forward_head
|
|
ComfyLogger().log("CAM: Logits do not require grad, skipping backward.", "WARNING")
|
|
|
|
# Update the list so we can access .grad later
|
|
intermediates[i] = intermediate_tensor
|
|
finally:
|
|
model.attn_pool.q = saved_q
|
|
model.attn_pool.out_proj.weight = saved_p
|
|
|
|
cam_1d = None
|
|
for intermediate in intermediates:
|
|
if intermediate.grad is None: continue
|
|
patch_grad = (intermediate.grad.float() * intermediate.sign()).sum(dim=(0, 2))
|
|
intermediate.grad = None # Clear grad to free memory
|
|
|
|
if cam_1d is None:
|
|
cam_1d = patch_grad
|
|
else:
|
|
cam_1d.add_(patch_grad)
|
|
|
|
if cam_1d is None:
|
|
return display_image # Should not happen
|
|
|
|
cam_2d = JtpImageV3Manager.unpatchify(cam_1d, patch_coords, patch_valid).cpu().numpy()
|
|
return JtpImageV3Manager.cam_composite(display_image, cam_2d)
|
|
|
|
@classmethod
|
|
def run_classifier(
|
|
cls,
|
|
model_name: str,
|
|
device: torch.device,
|
|
image: Union[Image.Image, np.ndarray, Path],
|
|
threshold: float,
|
|
cam_depth: int,
|
|
seqlen: int,
|
|
implications_mode: str,
|
|
exclude_tags: str,
|
|
exclude_categories: str,
|
|
prefix: str,
|
|
original_tags: bool,
|
|
seed: int,
|
|
replace_underscore: bool = True,
|
|
trailing_comma: bool = False,
|
|
cam_mode: str = "auto",
|
|
cam_tag: str = "",
|
|
mode: str = "topk",
|
|
topk: int = 40,
|
|
max_tags: int = 0,
|
|
category_config: Optional[Dict[str, Any]] = None,
|
|
) -> Tuple[str, Dict[str, float], Optional[Image.Image]]:
|
|
|
|
# Hash params for caching
|
|
cc_hash = ""
|
|
if category_config:
|
|
cc_hash = "|".join(f"{k}={v}" for k, v in sorted(category_config.items()))
|
|
params_string = f"{model_name}|{threshold}|{cam_depth}|{seqlen}|{implications_mode}|{exclude_tags}|{exclude_categories}|{prefix}|{original_tags}|{seed}|{replace_underscore}|{trailing_comma}|{cam_mode}|{cam_tag}|{mode}|{topk}|{max_tags}|{cc_hash}"
|
|
params_key = hashlib.sha256(params_string.encode()).hexdigest()
|
|
|
|
# Load Model
|
|
if JtpModelV3Manager().is_installed(model_name) is False:
|
|
if not JtpModelV3Manager().download(model_name):
|
|
return "", {}, None
|
|
|
|
if JtpTagV3Manager.is_loaded(model_name) is False:
|
|
JtpTagV3Manager.download(model_name)
|
|
if not JtpTagV3Manager.load(model_name):
|
|
ComfyLogger().log("Failed to load tags", "ERROR", True)
|
|
return "", {}, None
|
|
|
|
loaded, tags = JtpModelV3Manager.load(model_name, device=device)
|
|
if not loaded:
|
|
return "", {}, None
|
|
|
|
model = ComfyCache.get(f'model_v3.{model_name}.model')
|
|
|
|
# Move model to inference device (may have been offloaded to CPU)
|
|
model.to(device)
|
|
|
|
# Load/Process Image
|
|
result = JtpImageV3Manager.load(image, device=device, seqlen=seqlen, params_key=params_key)
|
|
|
|
if isinstance(result, tuple) and len(result) == 3 and isinstance(result[1], Dict): # Cached output (tags, scores, cam)
|
|
return result[0], result[1], result[2]
|
|
|
|
if result is None:
|
|
return "", {}, None
|
|
|
|
# It's tensors: patches, coords, valid
|
|
patches, coords, valid = result
|
|
|
|
# Prepare inputs
|
|
p_d = patches.unsqueeze(0).to(dtype=torch.bfloat16).div_(127.5).sub_(1.0)
|
|
pc_d = coords.unsqueeze(0).to(dtype=torch.int32)
|
|
pv_d = valid.unsqueeze(0)
|
|
|
|
# Inference
|
|
# We need gradients for CAM if we are going to run it?
|
|
# No, JTP-3 app.py runs forward with no_grad, then enables grad only for the backward pass on intermediates.
|
|
|
|
try:
|
|
with torch.no_grad():
|
|
features = model.forward_intermediates(
|
|
p_d,
|
|
patch_coord=pc_d,
|
|
patch_valid=pv_d,
|
|
indices=cam_depth,
|
|
output_dict=True,
|
|
output_fmt='NLC'
|
|
)
|
|
logits = model.forward_head(features["image_features"], patch_valid=pv_d)
|
|
|
|
probits = torch.sigmoid(logits[0].to(dtype=torch.float32))
|
|
probits.mul_(2.0).sub_(1.0) # Scale to -1..1
|
|
|
|
values, indices = probits.cpu().topk(250)
|
|
|
|
predictions = {
|
|
tags[idx.item()]: val.item()
|
|
for idx, val in sorted(
|
|
zip(indices, values),
|
|
key=lambda item: item[1].item(),
|
|
reverse=True
|
|
)
|
|
}
|
|
|
|
# ----------------------------------------------------------
|
|
# Top-K + threshold selection
|
|
# ----------------------------------------------------------
|
|
metadata = ComfyCache.get(f'tags_v3.{model_name}')
|
|
|
|
# Check if per-category selection should be used
|
|
use_per_cat = (
|
|
category_config is not None
|
|
and metadata is not None
|
|
and any(
|
|
category_config.get(f"{cat}_topk", 0) != 0
|
|
or category_config.get(f"{cat}_threshold", 0.0) != 0.0
|
|
for cat in CATEGORY_ID_TO_NAME.values()
|
|
)
|
|
)
|
|
|
|
if use_per_cat:
|
|
predictions = cls._per_category_select(
|
|
predictions=predictions,
|
|
metadata=metadata,
|
|
category_config=category_config,
|
|
global_mode=mode,
|
|
global_topk=topk,
|
|
global_threshold=threshold,
|
|
)
|
|
# Use a very low threshold so process_tags doesn't re-filter
|
|
effective_threshold = -2.0
|
|
else:
|
|
# Global selection: apply mode-based filtering on predictions
|
|
if mode == "topk":
|
|
# Filter by threshold first, then cap to topk
|
|
filtered = {t: s for t, s in predictions.items() if s >= threshold}
|
|
sorted_preds = sorted(filtered.items(), key=lambda x: x[1], reverse=True)
|
|
predictions = dict(sorted_preds[:topk])
|
|
effective_threshold = -2.0 # Already filtered
|
|
else:
|
|
effective_threshold = threshold
|
|
|
|
# Process tags (Implications, Thresholding)
|
|
tag_str, final_scores, top_tag_raw = JtpTagV3Manager.process_tags(
|
|
predictions, model_name, effective_threshold, implications_mode, exclude_tags,
|
|
exclude_categories, prefix, original_tags,
|
|
replace_underscore=replace_underscore, trailing_comma=trailing_comma,
|
|
max_tags=max_tags
|
|
)
|
|
|
|
cam_image = None
|
|
|
|
# CAM Logic
|
|
if cam_mode != "none":
|
|
target_idx = -1
|
|
|
|
if cam_mode == "auto":
|
|
# Use top 1 valid tag from filtered results
|
|
if top_tag_raw and top_tag_raw in tags:
|
|
try:
|
|
target_idx = tags.index(top_tag_raw)
|
|
except ValueError:
|
|
ComfyLogger().log(f"Auto CAM: '{top_tag_raw}' found in tags list but index retrieval failed.", "WARNING")
|
|
|
|
# Fallback to raw model top 1 if filtering removed everything but auto was requested?
|
|
# Or if process_tags returned nothing.
|
|
if target_idx == -1 and indices.numel() > 0:
|
|
# If we have no valid tags after filtering, should we show the raw top tag?
|
|
# The user likely wants to see *why* it triggered or just the strongest signal.
|
|
# But if they excluded it, maybe they don't want to see it.
|
|
# Let's fallback to raw top 1 to avoid black square confusion,
|
|
# but ideally 'top_tag_raw' covers the "best valid" case.
|
|
target_idx = indices[0].item()
|
|
|
|
elif cam_mode == "specific_tag" and cam_tag:
|
|
# Find tag index
|
|
# Clean up input (remove trailing commas, whitespace)
|
|
clean_tag = cam_tag.strip().rstrip(",")
|
|
|
|
try:
|
|
# Normalized check
|
|
if clean_tag in tags:
|
|
target_idx = tags.index(clean_tag)
|
|
else:
|
|
# Try replacing spaces with underscores
|
|
cam_tag_norm = clean_tag.replace(" ", "_")
|
|
if cam_tag_norm in tags:
|
|
target_idx = tags.index(cam_tag_norm)
|
|
else:
|
|
ComfyLogger().log(f"CAM Tag '{cam_tag}' (normalized: '{cam_tag_norm}') not found in model tags.", "WARNING")
|
|
except ValueError:
|
|
ComfyLogger().log(f"CAM Tag '{cam_tag}' caused a ValueError during lookup.", "WARNING")
|
|
|
|
if target_idx >= 0:
|
|
# Need original PIL image for composite
|
|
# If 'image' is PIL, use it. If tensor/numpy, we need to convert or use 'result' input?
|
|
# JtpImageV3Manager.load processes image.
|
|
# We can use the original 'image' if it's PIL.
|
|
display_img = None
|
|
if isinstance(image, Image.Image):
|
|
display_img = image.convert("RGB")
|
|
elif isinstance(image, np.ndarray):
|
|
display_img = Image.fromarray(image).convert("RGB")
|
|
elif isinstance(image, Path) and image.exists():
|
|
display_img = Image.open(str(image)).convert("RGB")
|
|
|
|
if display_img:
|
|
try:
|
|
cam_image = cls.run_cam(
|
|
model, display_img, features, target_idx, cam_depth,
|
|
pc_d, pv_d
|
|
)
|
|
except Exception as e:
|
|
import traceback
|
|
ComfyLogger().log(f"CAM generation failed: {e}\n{traceback.format_exc()}", "ERROR", True)
|
|
finally:
|
|
# Offload model to free GPU memory
|
|
model.to("cpu")
|
|
mm.soft_empty_cache()
|
|
|
|
# Cache result
|
|
image_key = None
|
|
if isinstance(image, Path): image_key = str(image)
|
|
elif isinstance(image, np.ndarray): image_key = hashlib.sha256(image.tobytes()).hexdigest()
|
|
elif isinstance(image, Image.Image):
|
|
img_arr = np.array(image.convert("RGB"))
|
|
image_key = hashlib.sha256(img_arr.tobytes()).hexdigest()
|
|
|
|
JtpImageV3Manager.commit_cache(image_key, (tag_str, final_scores, cam_image), params_key)
|
|
|
|
return tag_str, final_scores, cam_image
|
|
|
|
# ------------------------------------------------------------------
|
|
# Per-category tag selection
|
|
# ------------------------------------------------------------------
|
|
|
|
@staticmethod
|
|
def _per_category_select(
|
|
predictions: Dict[str, float],
|
|
metadata: Dict[str, Tuple],
|
|
category_config: Dict[str, Any],
|
|
global_mode: str,
|
|
global_topk: int,
|
|
global_threshold: float,
|
|
) -> Dict[str, float]:
|
|
"""Select tags using per-category topk/threshold overrides.
|
|
|
|
Groups predictions by category (from metadata CSV), then applies
|
|
per-category settings from *category_config*. Categories where
|
|
both topk and threshold are 0 fall back to global settings.
|
|
|
|
Returns a filtered predictions dict.
|
|
"""
|
|
# Group predictions by category ID
|
|
groups: Dict[int, List[Tuple[str, float]]] = defaultdict(list)
|
|
for tag, score in predictions.items():
|
|
cat_id = metadata[tag][0] if tag in metadata else 0
|
|
groups[cat_id].append((tag, score))
|
|
|
|
result: Dict[str, float] = {}
|
|
|
|
for cat_id, tag_scores in groups.items():
|
|
cat_name = CATEGORY_ID_TO_NAME.get(cat_id, "general")
|
|
cat_topk = category_config.get(f"{cat_name}_topk", 0)
|
|
cat_threshold = category_config.get(f"{cat_name}_threshold", 0.0)
|
|
|
|
# Determine effective settings for this category
|
|
if cat_topk == 0 and cat_threshold == 0.0:
|
|
# Use global settings
|
|
eff_threshold = global_threshold
|
|
eff_topk = global_topk if global_mode == "topk" else 0
|
|
elif cat_topk == 0:
|
|
# Threshold-only mode for this category
|
|
eff_threshold = cat_threshold
|
|
eff_topk = 0
|
|
else:
|
|
# TopK mode for this category, threshold as floor
|
|
eff_threshold = cat_threshold if cat_threshold > 0.0 else global_threshold
|
|
eff_topk = cat_topk
|
|
|
|
# Sort by score descending
|
|
tag_scores.sort(key=lambda x: x[1], reverse=True)
|
|
|
|
# Apply threshold
|
|
filtered = [(t, s) for t, s in tag_scores if s >= eff_threshold]
|
|
|
|
# Apply topk cap
|
|
if eff_topk > 0:
|
|
filtered = filtered[:eff_topk]
|
|
|
|
for t, s in filtered:
|
|
result[t] = s
|
|
|
|
return result
|