remove-ai-watermarks/scripts/chroma_text_restoration_study.py

# /// 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}")