83 lines
2.9 KiB
Python
83 lines
2.9 KiB
Python
"""ComfyUI custom node for character consistency auditing."""
|
|
|
|
import base64
|
|
import json
|
|
from pathlib import Path
|
|
|
|
from auditor import CHECKLIST_ITEMS, compare_images
|
|
from report_generator import generate_html
|
|
|
|
|
|
class CharacterConsistencyAuditor:
|
|
"""ComfyUI node that compares a reference image against a generated image using Claude vision."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"reference_image": ("IMAGE",),
|
|
"generated_image": ("IMAGE",),
|
|
},
|
|
"optional": {
|
|
"checklist_items": ("STRING", {
|
|
"default": ", ".join(CHECKLIST_ITEMS),
|
|
"multiline": False,
|
|
"tooltip": "Comma-separated checklist dimensions to compare.",
|
|
}),
|
|
"threshold": ("FLOAT", {
|
|
"default": 0.3,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.05,
|
|
"tooltip": "Drift threshold. Scores above this are flagged.",
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "FLOAT", "STRING")
|
|
RETURN_NAMES = ("json_report", "overall_drift", "html_dashboard")
|
|
FUNCTION = "audit"
|
|
CATEGORY = "image/analysis"
|
|
|
|
def _save_tensor_as_image(self, tensor, path: Path):
|
|
"""Convert a ComfyUI IMAGE tensor to a PNG file."""
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
arr = tensor.squeeze(0).cpu().numpy()
|
|
if arr.ndim == 3 and arr.shape[2] == 4:
|
|
arr = arr[:, :, :3]
|
|
arr = (arr * 255).clip(0, 255).astype("uint8")
|
|
Image.fromarray(arr).save(str(path))
|
|
|
|
def audit(self, reference_image, generated_image, checklist_items=DEFAULT, threshold=DEFAULT):
|
|
from auditor import compare_images as _compare
|
|
|
|
items = [c.strip() for c in checklist_items.split(",") if c.strip()]
|
|
|
|
import tempfile, os
|
|
ref_path = os.path.join(tempfile.gettempdir(), "comfy_ref.png")
|
|
gen_path = os.path.join(tempfile.gettempdir(), "comfy_gen.png")
|
|
|
|
self._save_tensor_as_image(reference_image, Path(ref_path))
|
|
self._save_tensor_as_image(generated_image, Path(gen_path))
|
|
|
|
result = _compare(ref_path, gen_path, checklist_items=items)
|
|
|
|
report = {"results": [result], "summary": {"drift_threshold": threshold}}
|
|
json_report = json.dumps(report, indent=2)
|
|
|
|
overall = float(result.get("overall_score", 0.0))
|
|
|
|
report["reference"] = ref_path
|
|
report["generated_folder"] = str(Path(gen_path).parent)
|
|
report["summary"]["total_images"] = 1
|
|
report["summary"]["average_drift"] = overall
|
|
report["summary"]["max_drift"] = overall
|
|
report["summary"]["min_drift"] = overall
|
|
report["summary"]["flagged_images"] = [result.get("filename", "input")] if overall > threshold else []
|
|
|
|
html = generate_html(report)
|
|
|
|
return (json_report, overall, html)
|