mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-10-14 05:20:46 +02:00
Add measured auto engine selection
This commit is contained in:
63 files changed
+1144
-74
No files matched your search
@@ -45,13 +45,17 @@ explicitly names a tracked canonical result.
|
||||
| `visible_sheets.py` | Build blinded contact sheets for relaxation candidates. |
|
||||
| `retrain_photo_classify.py` | CPU-retrain the 2026-08-31 photo heads from a sha256-keyed cache pack, no images. |
|
||||
| `publish_photo_classify_hf.py` | Upload the photo-classify card and freeze weights to Hugging Face `wiltodelta/raiw-models`. Manual; the Action `publish-photo-classify-hf.yml` is the write-token path. |
|
||||
| `verify_engine_selection_fixtures.py` | Verify hashes, dimensions, and prompt pairing in the tracked auto-engine content matrix. |
|
||||
|
||||
## Research and diagnostic prototypes
|
||||
|
||||
| Script | Purpose |
|
||||
| --- | --- |
|
||||
| `chroma_text_restoration_study.py` | Compare verified-text restoration for Qwen and Chroma on H100 without running both engines in production. |
|
||||
| `cjk_tail_probe.py` | Test a generic template for otherwise uncovered CJK labels. |
|
||||
| `controlnet_sweep.py` | Sweep the historical ControlNet removal prototype. |
|
||||
| `engine_selection_study.py` | Run one global engine at a time over the tracked prompt-matched content matrix. |
|
||||
| `analyze_engine_selection_study.py` | Reduce engine-study output to paired deltas and exact sign tests. |
|
||||
| `infer_text_lines.py` | Draft stable source-text lines without modifying pixels. |
|
||||
| `qwen_scrub_prototype.py` | Probe low-strength Qwen regeneration on a GPU. |
|
||||
| `selective_text_restoration.py` | Evaluate text restoration over a scrubbed image. |
|
||||
|
||||
@@ -0,0 +1,121 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Summarize paired fidelity results from the engine-selection study.
|
||||
|
||||
Run from the repository root after ``engine_selection_study.py``:
|
||||
|
||||
uv run scripts/analyze_engine_selection_study.py
|
||||
|
||||
The report uses per-image paired differences and exact two-sided sign tests.
|
||||
It does not turn the discovery set into a production routing rule.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import statistics
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
_METRICS: dict[str, tuple[Callable[[float], float], str]] = {
|
||||
"lpips_512": (lambda value: value, "lower"),
|
||||
"ssim": (lambda value: -value, "higher"),
|
||||
"psnr_db": (lambda value: -value, "higher"),
|
||||
"mae": (lambda value: value, "lower"),
|
||||
"edge_f1": (lambda value: -value, "higher"),
|
||||
"laplacian_ratio": (lambda value: abs(math.log(max(value, 1e-12))), "closer_to_one"),
|
||||
}
|
||||
|
||||
|
||||
def _sign_test_two_sided(wins: int, losses: int) -> float:
|
||||
"""Return the exact two-sided sign-test p-value, excluding ties."""
|
||||
trials = wins + losses
|
||||
if trials == 0:
|
||||
return 1.0
|
||||
tail = sum(math.comb(trials, index) for index in range(min(wins, losses) + 1))
|
||||
return min(1.0, 2.0 * tail / (2**trials))
|
||||
|
||||
|
||||
def _index(rows: list[dict[str, Any]]) -> dict[tuple[str, str], dict[str, Any]]:
|
||||
return {(row["provider"], row["pair_id"]): row for row in rows}
|
||||
|
||||
|
||||
def analyze(report: dict[str, Any]) -> dict[str, Any]:
|
||||
qwen = _index(report["qwen-zimage"]["rows"])
|
||||
chroma = _index(report["chroma-zimage"]["rows"])
|
||||
if qwen.keys() != chroma.keys():
|
||||
raise ValueError("Qwen and Chroma result keys differ")
|
||||
|
||||
by_provider: dict[str, list[tuple[dict[str, Any], dict[str, Any]]]] = defaultdict(list)
|
||||
for key in sorted(qwen):
|
||||
by_provider[key[0]].append((qwen[key], chroma[key]))
|
||||
|
||||
providers: dict[str, Any] = {}
|
||||
for provider, pairs in sorted(by_provider.items()):
|
||||
metrics: dict[str, Any] = {}
|
||||
for metric, (loss_transform, direction) in _METRICS.items():
|
||||
deltas: list[float] = []
|
||||
chroma_wins = 0
|
||||
qwen_wins = 0
|
||||
ties = 0
|
||||
for qwen_row, chroma_row in pairs:
|
||||
qwen_value = float(qwen_row["fidelity"][metric])
|
||||
chroma_value = float(chroma_row["fidelity"][metric])
|
||||
delta = loss_transform(chroma_value) - loss_transform(qwen_value)
|
||||
deltas.append(delta)
|
||||
if math.isclose(delta, 0.0, abs_tol=1e-12):
|
||||
ties += 1
|
||||
elif delta < 0:
|
||||
chroma_wins += 1
|
||||
else:
|
||||
qwen_wins += 1
|
||||
metrics[metric] = {
|
||||
"preferred_direction": direction,
|
||||
"delta_definition": "chroma_loss_minus_qwen_loss; negative favors Chroma",
|
||||
"chroma_wins": chroma_wins,
|
||||
"qwen_wins": qwen_wins,
|
||||
"ties": ties,
|
||||
"mean_delta": statistics.fmean(deltas),
|
||||
"median_delta": statistics.median(deltas),
|
||||
"sign_test_two_sided_p": _sign_test_two_sided(chroma_wins, qwen_wins),
|
||||
}
|
||||
providers[provider] = {"pairs": len(pairs), "metrics": metrics}
|
||||
|
||||
return {
|
||||
"method": "paired per-image differences with exact two-sided sign tests; ties excluded",
|
||||
"warning": "Discovery-set evidence only; challenge any proposed rule on unused generations.",
|
||||
"providers": providers,
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--input",
|
||||
type=Path,
|
||||
default=Path("out/engine-selection-study/summary.json"),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=Path,
|
||||
default=Path("out/engine-selection-study/analysis.json"),
|
||||
)
|
||||
args = parser.parse_args()
|
||||
report = json.loads(args.input.read_text(encoding="utf-8"))
|
||||
analysis = analyze(report)
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.output.write_text(json.dumps(analysis, indent=2) + "\n", encoding="utf-8")
|
||||
log.info("Wrote %s", args.output)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
|
||||
main()
|
||||
@@ -0,0 +1,232 @@
|
||||
# /// script
|
||||
# requires-python = ">=3.11"
|
||||
# dependencies = ["modal>=1.5"]
|
||||
# ///
|
||||
"""Paired offline study of verified-text restoration on Qwen and Chroma.
|
||||
|
||||
Each generative profile runs exactly once per fixture. Its already loaded VAE
|
||||
produces the restoration donor, then the shared production compositor creates
|
||||
the restored candidate from that one base output. The experiment therefore
|
||||
measures four outputs (Qwen/Chroma x base/restored) without accidentally timing
|
||||
or evaluating a second generation for either restoration arm.
|
||||
|
||||
Run from the repository root:
|
||||
|
||||
uvx modal run scripts/chroma_text_restoration_study.py
|
||||
|
||||
Outputs stay under ``out/text-restoration-engine-study``. Submit exact output
|
||||
hashes to the matching provider oracle before using fidelity metrics to choose
|
||||
between engines; fidelity only ranks candidates after the scrub verdict passes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import modal
|
||||
|
||||
_ROOT = Path(__file__).resolve().parent.parent
|
||||
_GPU = "H100"
|
||||
|
||||
app = modal.App("chroma-text-restoration-study")
|
||||
|
||||
image = (
|
||||
modal.Image.debian_slim(python_version="3.12")
|
||||
.pip_install(
|
||||
"torch",
|
||||
"torchvision>=0.20.0",
|
||||
"diffusers>=0.38,<1",
|
||||
"diffsynth>=2.0.17,<3",
|
||||
"transformers>=4.53",
|
||||
"accelerate",
|
||||
"sentencepiece",
|
||||
"protobuf",
|
||||
"pillow",
|
||||
"numpy",
|
||||
"opencv-python-headless<5",
|
||||
"onnxruntime>=1.24.0",
|
||||
"huggingface-hub>=0.20.0",
|
||||
)
|
||||
.env({"HF_HOME": "/cache"})
|
||||
.add_local_dir(_ROOT / "src" / "remove_ai_watermarks", remote_path="/pkg/remove_ai_watermarks")
|
||||
)
|
||||
|
||||
cache = modal.Volume.from_name("chroma-hf-cache", create_if_missing=True)
|
||||
|
||||
|
||||
def _png_bytes(image: Any) -> bytes:
|
||||
import io
|
||||
|
||||
buffer = io.BytesIO()
|
||||
image.save(buffer, format="PNG")
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def _metric_inputs(source: Any, boxes: list[tuple[int, int, int, int]]) -> tuple[Any, Any]:
|
||||
import numpy as np
|
||||
|
||||
source_rgb = np.asarray(source.convert("RGB"), dtype=np.uint8)
|
||||
mask = np.zeros(source_rgb.shape[:2], dtype=bool)
|
||||
for x1, y1, x2, y2 in boxes:
|
||||
mask[y1:y2, x1:x2] = True
|
||||
return source_rgb, mask
|
||||
|
||||
|
||||
def _metrics(source_rgb: Any, candidate: Any, mask: Any) -> dict[str, float]:
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
|
||||
candidate_rgb = np.asarray(candidate.convert("RGB"), dtype=np.uint8)
|
||||
signed_delta = source_rgb.astype(np.float32) - candidate_rgb
|
||||
absolute_delta = np.abs(signed_delta)
|
||||
mse = float((signed_delta**2).mean())
|
||||
return {
|
||||
"whole_mae": round(float(absolute_delta.mean()), 6),
|
||||
"text_box_mae": round(float(absolute_delta[mask].mean()), 6),
|
||||
"outside_box_mae": round(float(absolute_delta[~mask].mean()), 6),
|
||||
"psnr_db": round(99.0 if mse <= 1e-12 else 10.0 * math.log10(255.0**2 / mse), 6),
|
||||
}
|
||||
|
||||
|
||||
@app.function(image=image, gpu=_GPU, volumes={"/cache": cache}, timeout=3600)
|
||||
def run(payload: dict[str, dict[str, Any]]) -> dict[str, dict[str, Any]]:
|
||||
import gc
|
||||
import io
|
||||
import sys
|
||||
import time
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
sys.path.insert(0, "/pkg")
|
||||
|
||||
from remove_ai_watermarks._internal.chroma_zimage_pipeline import ChromaZImagePipeline
|
||||
from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline
|
||||
from remove_ai_watermarks._internal.text_restoration import VerifiedTextLine, restore_verified_text
|
||||
from remove_ai_watermarks._internal.two_stage_pipeline import detect_faces
|
||||
from remove_ai_watermarks._internal.watermark_profiles import resolve_strength
|
||||
|
||||
sources: dict[str, tuple[Image.Image, str, tuple[VerifiedTextLine, ...]]] = {}
|
||||
for name, item in payload.items():
|
||||
source = Image.open(io.BytesIO(item["source"])).convert("RGB")
|
||||
lines = tuple(
|
||||
VerifiedTextLine(
|
||||
tuple(line["box"]),
|
||||
line.get("text"),
|
||||
line.get("script"),
|
||||
float(line.get("angle", 0.0)),
|
||||
)
|
||||
for line in item["lines"]
|
||||
)
|
||||
sources[name] = source, str(item["vendor"]), lines
|
||||
|
||||
results: dict[str, dict[str, Any]] = {
|
||||
name: {"vendor": vendor, "engines": {}} for name, (_, vendor, _) in sources.items()
|
||||
}
|
||||
metric_inputs = {
|
||||
name: _metric_inputs(source, [line.box for line in lines]) for name, (source, _vendor, lines) in sources.items()
|
||||
}
|
||||
|
||||
for profile, pipeline_type in (
|
||||
("qwen-zimage", QwenZImagePipeline),
|
||||
("chroma-zimage", ChromaZImagePipeline),
|
||||
):
|
||||
print(f"[{profile}] loading...", flush=True)
|
||||
pipeline = pipeline_type(
|
||||
device="cuda",
|
||||
torch_dtype=torch.bfloat16,
|
||||
keep_global_models_on_device=True,
|
||||
keep_face_models_on_device=True,
|
||||
progress_callback=lambda message, p=profile: print(f"[{p}] {message}", flush=True),
|
||||
)
|
||||
for name, (source, vendor, lines) in sources.items():
|
||||
face_count = len(detect_faces(source)) if profile == "chroma-zimage" else 0
|
||||
strength = resolve_strength(
|
||||
None,
|
||||
vendor,
|
||||
profile,
|
||||
size=source.size,
|
||||
face_count=face_count,
|
||||
)
|
||||
print(f"[{profile}] {name}: strength={strength:.6f}, faces={face_count}", flush=True)
|
||||
|
||||
started = time.perf_counter()
|
||||
donor = pipeline._vae_roundtrip(source)
|
||||
donor_seconds = time.perf_counter() - started
|
||||
|
||||
started = time.perf_counter()
|
||||
base = pipeline.run(source, strength=strength, seed=0)
|
||||
generation_seconds = time.perf_counter() - started
|
||||
|
||||
started = time.perf_counter()
|
||||
restored = restore_verified_text(source, base, donor, lines)
|
||||
restoration_seconds = time.perf_counter() - started
|
||||
|
||||
source_rgb, text_mask = metric_inputs[name]
|
||||
base_png = _png_bytes(base)
|
||||
restored_png = _png_bytes(restored)
|
||||
donor_png = _png_bytes(donor)
|
||||
results[name]["engines"][profile] = {
|
||||
"strength": strength,
|
||||
"face_count": face_count,
|
||||
"donor_seconds": round(donor_seconds, 3),
|
||||
"generation_seconds": round(generation_seconds, 3),
|
||||
"restoration_seconds": round(restoration_seconds, 3),
|
||||
"base_sha256": hashlib.sha256(base_png).hexdigest(),
|
||||
"restored_sha256": hashlib.sha256(restored_png).hexdigest(),
|
||||
"donor_sha256": hashlib.sha256(donor_png).hexdigest(),
|
||||
"base_metrics": _metrics(source_rgb, base, text_mask),
|
||||
"restored_metrics": _metrics(source_rgb, restored, text_mask),
|
||||
"base_png": base_png,
|
||||
"restored_png": restored_png,
|
||||
"donor_png": donor_png,
|
||||
}
|
||||
|
||||
del pipeline
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
cache.commit()
|
||||
return results
|
||||
|
||||
|
||||
@app.local_entrypoint()
|
||||
def main(out: str = "out/text-restoration-engine-study") -> None:
|
||||
annotations = json.loads((_ROOT / "data/evaluations/fidelity/text-lines.json").read_text(encoding="utf-8"))
|
||||
vendors = {
|
||||
"ChatGPT Image May 31, 2026, 02_02_23 PM.png": "openai",
|
||||
"ChatGPT Image May 31, 2026, 02_03_55 PM.png": "openai",
|
||||
"Gemini_Generated_Image_633uuy633uuy633u.png": "google",
|
||||
}
|
||||
source_root = _ROOT / "data/synthid/originals"
|
||||
payload = {
|
||||
name: {
|
||||
"source": (source_root / name).read_bytes(),
|
||||
"vendor": vendor,
|
||||
"lines": annotations[name],
|
||||
}
|
||||
for name, vendor in vendors.items()
|
||||
}
|
||||
|
||||
results = run.remote(payload)
|
||||
out_root = Path(out)
|
||||
summary: dict[str, Any] = {"gpu": _GPU, "fixtures": {}}
|
||||
for name, fixture in results.items():
|
||||
fixture_dir = out_root / Path(name).stem
|
||||
fixture_dir.mkdir(parents=True, exist_ok=True)
|
||||
summary["fixtures"][name] = {"vendor": fixture["vendor"], "engines": {}}
|
||||
for profile, result in fixture["engines"].items():
|
||||
for variant in ("base", "restored", "donor"):
|
||||
(fixture_dir / f"{profile}_{variant}.png").write_bytes(result[f"{variant}_png"])
|
||||
summary["fixtures"][name]["engines"][profile] = {
|
||||
key: value for key, value in result.items() if not key.endswith("_png")
|
||||
}
|
||||
|
||||
out_root.mkdir(parents=True, exist_ok=True)
|
||||
(out_root / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(json.dumps(summary, ensure_ascii=False, indent=2))
|
||||
print(f"outputs in {out_root}")
|
||||
@@ -0,0 +1,221 @@
|
||||
# /// script
|
||||
# requires-python = ">=3.11"
|
||||
# dependencies = ["modal>=1.5"]
|
||||
# ///
|
||||
"""Compare one engine at a time on the tracked content-selection matrix.
|
||||
|
||||
This is an offline research harness, not a production router. Each Modal call
|
||||
loads exactly one global engine and runs all 38 fixed inputs at that engine's
|
||||
shipped provider floor. The shared Z-Image face repair is deliberately omitted:
|
||||
the model choice changes only the global stage, while the face stage regenerates
|
||||
the same original crops for both profiles.
|
||||
|
||||
Run from the repository root:
|
||||
|
||||
uvx modal run scripts/engine_selection_study.py
|
||||
|
||||
The combined metrics land in the gitignored
|
||||
``out/engine-selection-study/summary.json``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import modal
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
_ROOT = Path(__file__).resolve().parent.parent
|
||||
_FIXTURE_ROOT = _ROOT / "data" / "evaluations" / "engine-selection"
|
||||
_GPU = "H100"
|
||||
|
||||
app = modal.App("engine-selection-study")
|
||||
|
||||
image = (
|
||||
modal.Image.debian_slim(python_version="3.12")
|
||||
.pip_install(
|
||||
"torch",
|
||||
"torchvision>=0.20.0",
|
||||
"diffusers>=0.38,<1",
|
||||
"diffsynth>=2.0.17,<3",
|
||||
"transformers>=4.53",
|
||||
"accelerate",
|
||||
"sentencepiece",
|
||||
"protobuf",
|
||||
"pillow",
|
||||
"numpy",
|
||||
"opencv-python-headless",
|
||||
"scikit-image",
|
||||
"lpips",
|
||||
)
|
||||
.env({"HF_HOME": "/cache", "TORCH_HOME": "/cache/torch"})
|
||||
.add_local_dir(
|
||||
_ROOT / "src" / "remove_ai_watermarks",
|
||||
remote_path="/pkg/remove_ai_watermarks",
|
||||
)
|
||||
.add_local_dir(_FIXTURE_ROOT, remote_path="/fixtures")
|
||||
)
|
||||
|
||||
cache = modal.Volume.from_name("chroma-hf-cache", create_if_missing=True)
|
||||
|
||||
_STRENGTHS = {
|
||||
"qwen-zimage": {"openai": 0.07675, "meta": 0.10},
|
||||
"chroma-zimage": {"openai": 0.09, "meta": 0.17},
|
||||
}
|
||||
|
||||
|
||||
@app.function(image=image, gpu=_GPU, volumes={"/cache": cache}, timeout=3600)
|
||||
def run_engine(engine: str) -> dict[str, Any]:
|
||||
import csv
|
||||
import hashlib
|
||||
import sys
|
||||
import time
|
||||
|
||||
import cv2
|
||||
import lpips
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
from skimage.metrics import peak_signal_noise_ratio, structural_similarity
|
||||
|
||||
sys.path.insert(0, "/pkg")
|
||||
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
|
||||
|
||||
if engine == "qwen-zimage":
|
||||
from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline
|
||||
|
||||
pipeline_type: Any = QwenZImagePipeline
|
||||
elif engine == "chroma-zimage":
|
||||
from remove_ai_watermarks._internal.chroma_zimage_pipeline import ChromaZImagePipeline
|
||||
|
||||
pipeline_type = ChromaZImagePipeline
|
||||
else:
|
||||
raise ValueError(f"unknown engine: {engine}")
|
||||
pipeline = pipeline_type(
|
||||
device="cuda",
|
||||
torch_dtype=torch.bfloat16,
|
||||
keep_global_models_on_device=True,
|
||||
keep_face_models_on_device=False,
|
||||
)
|
||||
|
||||
perceptual = lpips.LPIPS(net="alex", verbose=False).to("cuda").eval()
|
||||
|
||||
def rgb_array(pil_image: Image.Image) -> np.ndarray:
|
||||
return np.asarray(pil_image.convert("RGB"), dtype=np.uint8)
|
||||
|
||||
def gray(rgb: np.ndarray) -> np.ndarray:
|
||||
return cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
|
||||
|
||||
def lpips_distance(source: np.ndarray, candidate: np.ndarray) -> float:
|
||||
height, width = source.shape[:2]
|
||||
scale = min(1.0, 512.0 / max(height, width))
|
||||
size = (max(1, round(width * scale)), max(1, round(height * scale)))
|
||||
|
||||
def tensor(array: np.ndarray) -> torch.Tensor:
|
||||
resized = cv2.resize(array, size, interpolation=cv2.INTER_AREA)
|
||||
normalized = resized.astype(np.float32) / 127.5 - 1.0
|
||||
return torch.from_numpy(normalized).permute(2, 0, 1).unsqueeze(0).to("cuda")
|
||||
|
||||
with torch.inference_mode():
|
||||
return float(perceptual(tensor(source), tensor(candidate)).item())
|
||||
|
||||
def edge_f1(source_edges: np.ndarray, candidate_gray: np.ndarray) -> float:
|
||||
candidate_edges = cv2.Canny(candidate_gray, 100, 200) > 0
|
||||
kernel = np.ones((3, 3), np.uint8)
|
||||
source_near = cv2.dilate(source_edges.astype(np.uint8), kernel) > 0
|
||||
candidate_near = cv2.dilate(candidate_edges.astype(np.uint8), kernel) > 0
|
||||
precision = float((candidate_edges & source_near).sum()) / max(1, int(candidate_edges.sum()))
|
||||
recall = float((source_edges & candidate_near).sum()) / max(1, int(source_edges.sum()))
|
||||
return 2.0 * precision * recall / max(1e-9, precision + recall)
|
||||
|
||||
def source_features(
|
||||
source: np.ndarray,
|
||||
source_gray: np.ndarray,
|
||||
source_edges: np.ndarray,
|
||||
source_lapvar: float,
|
||||
) -> dict[str, float]:
|
||||
histogram = np.bincount(source_gray.reshape(-1), minlength=256).astype(np.float64)
|
||||
probabilities = histogram[histogram > 0] / histogram.sum()
|
||||
hsv = cv2.cvtColor(source, cv2.COLOR_RGB2HSV)
|
||||
return {
|
||||
"edge_density": round(float(source_edges.mean()), 6),
|
||||
"laplacian_variance": round(source_lapvar, 6),
|
||||
"entropy_bits": round(float(-(probabilities * np.log2(probabilities)).sum()), 6),
|
||||
"mean_saturation": round(float(hsv[..., 1].mean() / 255.0), 6),
|
||||
}
|
||||
|
||||
def fidelity(
|
||||
source: np.ndarray,
|
||||
candidate: np.ndarray,
|
||||
source_gray: np.ndarray,
|
||||
source_edges: np.ndarray,
|
||||
source_lapvar: float,
|
||||
) -> dict[str, float]:
|
||||
candidate_gray = gray(candidate)
|
||||
candidate_lapvar = float(cv2.Laplacian(candidate_gray, cv2.CV_64F).var())
|
||||
return {
|
||||
"lpips_512": round(lpips_distance(source, candidate), 6),
|
||||
"ssim": round(float(structural_similarity(source_gray, candidate_gray)), 6),
|
||||
"psnr_db": round(float(peak_signal_noise_ratio(source, candidate)), 6),
|
||||
"mae": round(float(np.abs(source.astype(np.float32) - candidate).mean()), 6),
|
||||
"edge_f1": round(edge_f1(source_edges, candidate_gray), 6),
|
||||
"laplacian_ratio": round(candidate_lapvar / max(source_lapvar, 1e-9), 6),
|
||||
}
|
||||
|
||||
with Path("/fixtures/content-manifest.csv").open(newline="", encoding="utf-8") as stream:
|
||||
rows = list(csv.DictReader(stream))
|
||||
|
||||
results: list[dict[str, Any]] = []
|
||||
for index, row in enumerate(rows, start=1):
|
||||
source_path = Path("/fixtures") / row["file"]
|
||||
source_image = Image.open(source_path).convert("RGB")
|
||||
strength = _STRENGTHS[engine][row["provider"]]
|
||||
started = time.perf_counter()
|
||||
output = pipeline._run_global(source_image, strength, 0)
|
||||
elapsed = time.perf_counter() - started
|
||||
source = rgb_array(source_image)
|
||||
candidate = rgb_array(output)
|
||||
source_gray = gray(source)
|
||||
source_edges = cv2.Canny(source_gray, 100, 200) > 0
|
||||
source_lapvar = float(cv2.Laplacian(source_gray, cv2.CV_64F).var())
|
||||
result = {
|
||||
"pair_id": row["pair_id"],
|
||||
"provider": row["provider"],
|
||||
"content_stratum": row["content_stratum"],
|
||||
"strength": strength,
|
||||
"seconds": round(elapsed, 3),
|
||||
"output_rgb_sha256": hashlib.sha256(output.tobytes()).hexdigest(),
|
||||
"source_features": source_features(source, source_gray, source_edges, source_lapvar),
|
||||
"fidelity": fidelity(source, candidate, source_gray, source_edges, source_lapvar),
|
||||
}
|
||||
results.append(result)
|
||||
log.info(
|
||||
"%s %s/%s (%s/%s)",
|
||||
engine,
|
||||
index,
|
||||
len(rows),
|
||||
row["provider"],
|
||||
row["content_stratum"],
|
||||
)
|
||||
|
||||
cache.commit()
|
||||
return {"engine": engine, "gpu": _GPU, "rows": results}
|
||||
|
||||
|
||||
@app.local_entrypoint()
|
||||
def main(out: str = "out/engine-selection-study") -> None:
|
||||
qwen = run_engine.remote("qwen-zimage")
|
||||
chroma = run_engine.remote("chroma-zimage")
|
||||
report = {
|
||||
"protocol": "one engine per sequential H100 invocation; global-stage factorial ablation",
|
||||
"qwen-zimage": qwen,
|
||||
"chroma-zimage": chroma,
|
||||
}
|
||||
destination = Path(out) / "summary.json"
|
||||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||||
destination.write_text(json.dumps(report, indent=2), encoding="utf-8")
|
||||
log.info("Wrote %s", destination)
|
||||
@@ -0,0 +1,128 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Verify the tracked content matrix used by the auto-engine study."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import hashlib
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
_ROOT = Path(__file__).resolve().parent.parent
|
||||
_FIXTURE_ROOT = _ROOT / "data" / "evaluations" / "engine-selection"
|
||||
_MANIFEST = _FIXTURE_ROOT / "content-manifest.csv"
|
||||
_CARRIER_MANIFEST = _FIXTURE_ROOT / "carrier-manifest.csv"
|
||||
_EXPECTED_PROVIDERS = {"meta", "openai"}
|
||||
|
||||
|
||||
def _sha256(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as stream:
|
||||
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
errors: list[str] = []
|
||||
pairs: dict[str, list[dict[str, str]]] = defaultdict(list)
|
||||
paths: set[Path] = set()
|
||||
hashes: set[str] = set()
|
||||
|
||||
with _MANIFEST.open(newline="", encoding="utf-8") as stream:
|
||||
rows = list(csv.DictReader(stream))
|
||||
with _CARRIER_MANIFEST.open(newline="", encoding="utf-8") as stream:
|
||||
carrier_rows = list(csv.DictReader(stream))
|
||||
|
||||
for row in rows:
|
||||
pair_id = row["pair_id"]
|
||||
relative_path = Path(row["file"])
|
||||
path = _FIXTURE_ROOT / relative_path
|
||||
pairs[pair_id].append(row)
|
||||
|
||||
if relative_path in paths:
|
||||
errors.append(f"duplicate path: {relative_path}")
|
||||
paths.add(relative_path)
|
||||
|
||||
expected_hash = row["sha256"]
|
||||
if expected_hash in hashes:
|
||||
errors.append(f"duplicate file hash: {expected_hash}")
|
||||
hashes.add(expected_hash)
|
||||
|
||||
if not path.is_file():
|
||||
errors.append(f"missing file: {relative_path}")
|
||||
continue
|
||||
actual_hash = _sha256(path)
|
||||
if actual_hash != expected_hash:
|
||||
errors.append(f"hash mismatch for {relative_path}: expected {expected_hash}, got {actual_hash}")
|
||||
|
||||
with Image.open(path) as image:
|
||||
actual_size = image.size
|
||||
actual_format = image.format
|
||||
expected_size = (int(row["width"]), int(row["height"]))
|
||||
if actual_size != expected_size:
|
||||
errors.append(f"size mismatch for {relative_path}: expected {expected_size}, got {actual_size}")
|
||||
if actual_format != "PNG":
|
||||
errors.append(f"format mismatch for {relative_path}: expected PNG, got {actual_format}")
|
||||
|
||||
for pair_id, pair_rows in sorted(pairs.items()):
|
||||
providers = {row["provider"] for row in pair_rows}
|
||||
prompts = {row["prompt"] for row in pair_rows}
|
||||
strata = {row["content_stratum"] for row in pair_rows}
|
||||
if providers != _EXPECTED_PROVIDERS:
|
||||
errors.append(f"pair {pair_id} providers: expected {_EXPECTED_PROVIDERS}, got {providers}")
|
||||
if len(prompts) != 1:
|
||||
errors.append(f"pair {pair_id} has mismatched prompts")
|
||||
if len(strata) != 1:
|
||||
errors.append(f"pair {pair_id} has mismatched content strata")
|
||||
|
||||
expected_pair_ids = {f"{index:02d}" for index in range(19)}
|
||||
if set(pairs) != expected_pair_ids:
|
||||
errors.append(f"pair ids: expected {sorted(expected_pair_ids)}, got {sorted(pairs)}")
|
||||
|
||||
committed_files = {
|
||||
path.relative_to(_FIXTURE_ROOT) for path in (_FIXTURE_ROOT / "originals").glob("**/*") if path.is_file()
|
||||
}
|
||||
unlisted = committed_files - paths
|
||||
if unlisted:
|
||||
errors.append(f"unlisted files: {sorted(str(path) for path in unlisted)}")
|
||||
|
||||
carrier_providers: set[str] = set()
|
||||
for row in carrier_rows:
|
||||
relative_path = Path(row["file"])
|
||||
path = (_FIXTURE_ROOT / relative_path).resolve()
|
||||
carrier_providers.add(row["provider"])
|
||||
if not path.is_relative_to(_ROOT / "data"):
|
||||
errors.append(f"carrier outside data/: {relative_path}")
|
||||
continue
|
||||
if not path.is_file():
|
||||
errors.append(f"missing carrier: {relative_path}")
|
||||
continue
|
||||
actual_hash = _sha256(path)
|
||||
if actual_hash != row["sha256"]:
|
||||
errors.append(f"carrier hash mismatch for {relative_path}: expected {row['sha256']}, got {actual_hash}")
|
||||
if carrier_providers != {"google", "meta", "openai"}:
|
||||
errors.append(f"carrier providers: expected google/meta/openai; got {sorted(carrier_providers)}")
|
||||
|
||||
if errors:
|
||||
for error in errors:
|
||||
log.error("%s", error)
|
||||
return 1
|
||||
|
||||
log.info(
|
||||
"Verified %s content files in %s matched pairs and %s canonical carriers",
|
||||
len(rows),
|
||||
len(pairs),
|
||||
len(carrier_rows),
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
|
||||
raise SystemExit(main())
|
||||
Reference in new issue
Block a user