chore: import upstream snapshot with attribution

scripts/ci/utils/diffusion/publish_diffusion_gt.py

@@ -0,0 +1,507 @@
+"""
+Publish diffusion CI ground-truth images to sgl-project/ci-data
+via the GitHub API (same pattern as publish_traces.py).
+"""
+
+import argparse
+import base64
+import hashlib
+import io
+import json
+import os
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from urllib.error import HTTPError
+
+import numpy as np
+from PIL import Image, ImageFilter
+
+# Reuse GitHub API helpers from publish_traces.
+# Support both direct script execution and package-style imports.
+if __package__:
+    from ..publish_traces import (
+        create_blobs,
+        create_commit,
+        create_tree,
+        get_branch_sha,
+        get_tree_sha,
+        is_permission_error,
+        is_rate_limit_error,
+        make_github_request,
+        update_branch_ref,
+        verify_token_permissions,
+    )
+else:
+    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+    from publish_traces import (
+        create_blobs,
+        create_commit,
+        create_tree,
+        get_branch_sha,
+        get_tree_sha,
+        is_permission_error,
+        is_rate_limit_error,
+        make_github_request,
+        update_branch_ref,
+        verify_token_permissions,
+    )
+
+REPO_OWNER = "sgl-project"
+REPO_NAME = "ci-data"
+BRANCH = "main"
+DEFAULT_TARGET_DIR = "diffusion-ci/consistency_gt/sglang_generated"
+
+IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"}
+QUALITY_MAX_SIDE = 256
+LOW_DETAIL_STD_THRESHOLD = 0.075
+LOW_DETAIL_ENTROPY_THRESHOLD = 0.55
+LOW_DETAIL_BLUR_RESIDUAL_THRESHOLD = 0.035
+LOW_DETAIL_GRADIENT_P95_THRESHOLD = 0.045
+RANDOM_NOISE_CORRELATION_THRESHOLD = 0.55
+RANDOM_NOISE_LOW_FREQUENCY_THRESHOLD = 0.20
+RANDOM_NOISE_BLUR_RESIDUAL_THRESHOLD = 0.045
+OLD_NEW_MIN_SSIM = 0.20
+OLD_NEW_MAX_MEAN_ABS_DIFF = 45.0
+
+
+@dataclass(frozen=True)
+class ImageQualityMetrics:
+    luminance_std: float
+    entropy: float
+    blur_residual: float
+    gradient_p95: float
+    neighbor_correlation: float
+    low_frequency_ratio: float
+
+
+@dataclass(frozen=True)
+class OldNewMetrics:
+    ssim: float
+    mean_abs_diff: float
+
+
+def collect_images(source_dir, target_dir):
+    """Collect image files from source_dir and return list of (repo_path, content) tuples."""
+    files = []
+    for entry in sorted(os.listdir(source_dir)):
+        ext = os.path.splitext(entry)[1].lower()
+        if ext not in IMAGE_EXTENSIONS:
+            continue
+        full_path = os.path.join(source_dir, entry)
+        if not os.path.isfile(full_path):
+            continue
+        with open(full_path, "rb") as f:
+            content = f.read()
+        repo_path = f"{target_dir}/{entry}"
+        files.append((repo_path, content))
+    return files
+
+
+def git_blob_sha(content):
+    header = f"blob {len(content)}\0".encode()
+    return hashlib.sha1(header + content).hexdigest()
+
+
+def get_remote_blob_shas(repo_owner, repo_name, target_dir, token):
+    return {
+        path: item["sha"]
+        for path, item in get_remote_image_entries(
+            repo_owner, repo_name, target_dir, token
+        ).items()
+    }
+
+
+def get_remote_image_entries(repo_owner, repo_name, target_dir, token):
+    url = (
+        f"https://api.github.com/repos/{repo_owner}/{repo_name}/contents/"
+        f"{target_dir}?ref={BRANCH}"
+    )
+    try:
+        response = make_github_request(url, token)
+    except HTTPError as e:
+        if e.code == 404:
+            return {}
+        raise
+    entries = json.loads(response)
+    return {
+        item["path"]: item
+        for item in entries
+        if item.get("type") == "file"
+        and "sha" in item
+        and os.path.splitext(item["path"])[1].lower() in IMAGE_EXTENSIONS
+    }
+
+
+def filter_changed_files(files, remote_blob_shas):
+    return [
+        (path, content)
+        for path, content in files
+        if remote_blob_shas.get(path) != git_blob_sha(content)
+    ]
+
+
+def get_remote_blob_content(repo_owner, repo_name, blob_sha, token):
+    url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/git/blobs/{blob_sha}"
+    response = make_github_request(url, token)
+    blob = json.loads(response)
+    if blob.get("encoding") != "base64":
+        raise ValueError(
+            f"Unexpected blob encoding for {blob_sha}: {blob.get('encoding')}"
+        )
+    return base64.b64decode(blob["content"])
+
+
+def _load_quality_image(content):
+    with Image.open(io.BytesIO(content)) as image:
+        image = image.convert("RGB")
+        image.thumbnail((QUALITY_MAX_SIDE, QUALITY_MAX_SIDE), Image.Resampling.BICUBIC)
+        return image.copy()
+
+
+def _image_to_rgb_array(image):
+    return np.asarray(image, dtype=np.float32)
+
+
+def _luminance(rgb):
+    return 0.299 * rgb[..., 0] + 0.587 * rgb[..., 1] + 0.114 * rgb[..., 2]
+
+
+def _neighbor_correlation(luma):
+    def corr(a, b):
+        a = a.ravel()
+        b = b.ravel()
+        if a.std() < 1e-6 or b.std() < 1e-6:
+            return 1.0
+        return float(np.corrcoef(a, b)[0, 1])
+
+    return (corr(luma[:, 1:], luma[:, :-1]) + corr(luma[1:, :], luma[:-1, :])) / 2
+
+
+def _low_frequency_ratio(luma):
+    centered = luma - luma.mean()
+    power = np.abs(np.fft.fftshift(np.fft.fft2(centered))) ** 2
+    total_power = power.sum()
+    if total_power < 1e-12:
+        return 0.0
+
+    height, width = luma.shape
+    y, x = np.ogrid[:height, :width]
+    center_y = height // 2
+    center_x = width // 2
+    radius = np.sqrt((y - center_y) ** 2 + (x - center_x) ** 2)
+    low_frequency_radius = min(height, width) * 0.08
+    return float(power[radius <= low_frequency_radius].sum() / total_power)
+
+
+def compute_image_quality_metrics(content):
+    image = _load_quality_image(content)
+    rgb = _image_to_rgb_array(image)
+    luma = _luminance(rgb) / 255.0
+
+    gradients = np.concatenate(
+        [
+            np.abs(np.diff(luma, axis=1)).ravel(),
+            np.abs(np.diff(luma, axis=0)).ravel(),
+        ]
+    )
+    histogram, _ = np.histogram(luma, bins=32, range=(0, 1))
+    probabilities = histogram / histogram.sum()
+    nonzero_probabilities = probabilities[probabilities > 0]
+    entropy = float(
+        -(nonzero_probabilities * np.log2(nonzero_probabilities)).sum() / 5.0
+    )
+    blurred = _image_to_rgb_array(image.filter(ImageFilter.GaussianBlur(radius=3)))
+
+    return ImageQualityMetrics(
+        luminance_std=float(luma.std()),
+        entropy=entropy,
+        blur_residual=float(np.mean(np.abs(rgb - blurred)) / 255.0),
+        gradient_p95=float(np.percentile(gradients, 95)),
+        neighbor_correlation=_neighbor_correlation(luma),
+        low_frequency_ratio=_low_frequency_ratio(luma),
+    )
+
+
+def get_quality_failure_reasons(metrics):
+    reasons = []
+    low_detail_static = (
+        metrics.luminance_std < LOW_DETAIL_STD_THRESHOLD
+        and metrics.entropy < LOW_DETAIL_ENTROPY_THRESHOLD
+        and (
+            metrics.blur_residual < LOW_DETAIL_BLUR_RESIDUAL_THRESHOLD
+            or metrics.gradient_p95 < LOW_DETAIL_GRADIENT_P95_THRESHOLD
+        )
+    )
+    high_frequency_noise = (
+        metrics.neighbor_correlation < RANDOM_NOISE_CORRELATION_THRESHOLD
+        and metrics.low_frequency_ratio < RANDOM_NOISE_LOW_FREQUENCY_THRESHOLD
+        and metrics.blur_residual > RANDOM_NOISE_BLUR_RESIDUAL_THRESHOLD
+    )
+    if low_detail_static:
+        reasons.append("low-contrast low-detail output")
+    if high_frequency_noise:
+        reasons.append("high-frequency random noise")
+    return reasons
+
+
+def _resize_for_old_new_compare(content, size=None):
+    with Image.open(io.BytesIO(content)) as image:
+        image = image.convert("RGB")
+        if size is None:
+            image.thumbnail(
+                (QUALITY_MAX_SIDE, QUALITY_MAX_SIDE), Image.Resampling.BICUBIC
+            )
+        else:
+            image = image.resize(size, Image.Resampling.BICUBIC)
+        return _image_to_rgb_array(image)
+
+
+def compute_old_new_metrics(old_content, new_content):
+    old_rgb = _resize_for_old_new_compare(old_content)
+    new_rgb = _resize_for_old_new_compare(
+        new_content, size=(old_rgb.shape[1], old_rgb.shape[0])
+    )
+    old_luma = _luminance(old_rgb) / 255.0
+    new_luma = _luminance(new_rgb) / 255.0
+
+    old_mean = old_luma.mean()
+    new_mean = new_luma.mean()
+    old_variance = old_luma.var()
+    new_variance = new_luma.var()
+    covariance = ((old_luma - old_mean) * (new_luma - new_mean)).mean()
+    c1 = 0.01**2
+    c2 = 0.03**2
+    ssim = (
+        (2 * old_mean * new_mean + c1)
+        * (2 * covariance + c2)
+        / ((old_mean**2 + new_mean**2 + c1) * (old_variance + new_variance + c2))
+    )
+
+    return OldNewMetrics(
+        ssim=float(ssim),
+        mean_abs_diff=float(np.abs(old_rgb - new_rgb).mean()),
+    )
+
+
+def _format_quality_metrics(metrics):
+    return (
+        f"std={metrics.luminance_std:.4f}, entropy={metrics.entropy:.4f}, "
+        f"blur_residual={metrics.blur_residual:.4f}, "
+        f"gradient_p95={metrics.gradient_p95:.4f}, "
+        f"neighbor_corr={metrics.neighbor_correlation:.4f}, "
+        f"low_freq={metrics.low_frequency_ratio:.4f}"
+    )
+
+
+def _format_old_new_metrics(metrics):
+    return f"ssim={metrics.ssim:.4f}, mean_abs_diff={metrics.mean_abs_diff:.2f}"
+
+
+def validate_gt_files(files_to_upload, changed_files, remote_image_entries, token):
+    failures = []
+    for path, content in files_to_upload:
+        quality_metrics = compute_image_quality_metrics(content)
+        quality_reasons = get_quality_failure_reasons(quality_metrics)
+        if quality_reasons:
+            failures.append(
+                f"{path}: {', '.join(quality_reasons)} "
+                f"({_format_quality_metrics(quality_metrics)})"
+            )
+
+    for path, content in changed_files:
+        remote_entry = remote_image_entries.get(path)
+        if not remote_entry:
+            continue
+
+        old_content = get_remote_blob_content(
+            REPO_OWNER, REPO_NAME, remote_entry["sha"], token
+        )
+        old_quality_metrics = compute_image_quality_metrics(old_content)
+        old_quality_reasons = get_quality_failure_reasons(old_quality_metrics)
+        if old_quality_reasons:
+            print(
+                f"Skipping old/new drift check for {path} because existing GT is "
+                f"already suspicious: {', '.join(old_quality_reasons)} "
+                f"({_format_quality_metrics(old_quality_metrics)})"
+            )
+            continue
+
+        old_new_metrics = compute_old_new_metrics(old_content, content)
+        if (
+            old_new_metrics.ssim < OLD_NEW_MIN_SSIM
+            and old_new_metrics.mean_abs_diff > OLD_NEW_MAX_MEAN_ABS_DIFF
+        ):
+            failures.append(
+                f"{path}: changed too far from existing GT "
+                f"({_format_old_new_metrics(old_new_metrics)})"
+            )
+
+    if not failures:
+        print(
+            f"GT quality gate passed for {len(files_to_upload)} generated image(s) "
+            f"and {len(changed_files)} changed image(s)."
+        )
+        return
+
+    print("GT quality gate failed; refusing to publish suspicious image updates:")
+    for failure in failures:
+        print(f"  - {failure}")
+    sys.exit(1)
+
+
+def check_quality(source_dir, target_dir=None):
+    target_dir = target_dir or DEFAULT_TARGET_DIR
+    token = os.getenv("GITHUB_TOKEN")
+    if not token:
+        print("Error: GITHUB_TOKEN environment variable not set")
+        sys.exit(1)
+
+    files_to_upload = collect_images(source_dir, target_dir)
+    if not files_to_upload:
+        print(f"No image files found in {source_dir}")
+        return
+
+    remote_image_entries = get_remote_image_entries(
+        REPO_OWNER, REPO_NAME, target_dir, token
+    )
+    remote_blob_shas = {
+        path: item["sha"] for path, item in remote_image_entries.items()
+    }
+    changed_files = filter_changed_files(files_to_upload, remote_blob_shas)
+    validate_gt_files(files_to_upload, changed_files, remote_image_entries, token)
+
+
+def publish(source_dir, target_dir=None):
+    target_dir = target_dir or DEFAULT_TARGET_DIR
+    token = os.getenv("GITHUB_TOKEN")
+    if not token:
+        print("Error: GITHUB_TOKEN environment variable not set")
+        sys.exit(1)
+
+    files_to_upload = collect_images(source_dir, target_dir)
+    if not files_to_upload:
+        print(f"No image files found in {source_dir}")
+        return
+
+    print(
+        f"Found {len(files_to_upload)} image(s) to upload to {REPO_OWNER}/{REPO_NAME}/{target_dir}"
+    )
+
+    # Verify token
+    perm = verify_token_permissions(REPO_OWNER, REPO_NAME, token)
+    if perm == "rate_limited":
+        print("GitHub API rate-limited, skipping upload.")
+        return
+    if not perm:
+        print("Token permission verification failed.")
+        sys.exit(1)
+
+    # Commit with retry (handle concurrent pushes)
+    max_retries = 5
+    for attempt in range(max_retries):
+        try:
+            branch_sha = get_branch_sha(REPO_OWNER, REPO_NAME, BRANCH, token)
+            tree_sha = get_tree_sha(REPO_OWNER, REPO_NAME, branch_sha, token)
+            remote_image_entries = get_remote_image_entries(
+                REPO_OWNER, REPO_NAME, target_dir, token
+            )
+            remote_blob_shas = {
+                path: item["sha"] for path, item in remote_image_entries.items()
+            }
+            changed_files = filter_changed_files(files_to_upload, remote_blob_shas)
+            validate_gt_files(
+                files_to_upload, changed_files, remote_image_entries, token
+            )
+            if not changed_files:
+                print("No image changes to publish.")
+                return
+
+            try:
+                tree_items = create_blobs(REPO_OWNER, REPO_NAME, changed_files, token)
+            except Exception as e:
+                if is_rate_limit_error(e):
+                    print("Rate-limited during blob creation, skipping.")
+                    return
+                if is_permission_error(e):
+                    print(
+                        f"ERROR: Token lacks write permission to {REPO_OWNER}/{REPO_NAME}. "
+                        "Update GH_PAT_FOR_NIGHTLY_CI_DATA with a token that has contents:write."
+                    )
+                    sys.exit(1)
+                raise
+
+            new_tree_sha = create_tree(
+                REPO_OWNER, REPO_NAME, tree_sha, tree_items, token
+            )
+            if new_tree_sha == tree_sha:
+                print("No tree changes to publish.")
+                return
+
+            commit_msg = f"diffusion-ci: update images in {target_dir} ({len(changed_files)} files) [automated]"
+            commit_sha = create_commit(
+                REPO_OWNER, REPO_NAME, new_tree_sha, branch_sha, commit_msg, token
+            )
+            update_branch_ref(REPO_OWNER, REPO_NAME, BRANCH, commit_sha, token)
+            print(
+                f"Successfully pushed {len(changed_files)} changed images (commit {commit_sha[:10]})"
+            )
+            return
+        except Exception as e:
+            if is_rate_limit_error(e):
+                print("Rate-limited, skipping.")
+                return
+            if is_permission_error(e):
+                print(f"ERROR: permission denied to {REPO_OWNER}/{REPO_NAME}")
+                sys.exit(1)
+
+            retryable = False
+            if hasattr(e, "error_body"):
+                if "Update is not a fast forward" in e.error_body:
+                    retryable = True
+                elif "Object does not exist" in e.error_body:
+                    retryable = True
+
+            if isinstance(e, HTTPError) and e.code in [422, 500, 502, 503, 504]:
+                retryable = True
+
+            if retryable and attempt < max_retries - 1:
+                import time
+
+                wait = 2**attempt
+                print(
+                    f"Attempt {attempt + 1}/{max_retries} failed, retrying in {wait}s..."
+                )
+                time.sleep(wait)
+            else:
+                print(f"Failed after {attempt + 1} attempts: {e}")
+                raise
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Publish diffusion GT images to GitHub"
+    )
+    parser.add_argument(
+        "--source-dir", required=True, help="Directory containing GT images"
+    )
+    parser.add_argument(
+        "--target-dir",
+        required=False,
+        default=None,
+        help=f"Target directory in the remote repo (default: {DEFAULT_TARGET_DIR})",
+    )
+    parser.add_argument(
+        "--check-only",
+        action="store_true",
+        help="Validate generated GT images without publishing them",
+    )
+    args = parser.parse_args()
+    if args.check_only:
+        check_quality(args.source_dir, args.target_dir)
+    else:
+        publish(args.source_dir, args.target_dir)
+
+
+if __name__ == "__main__":
+    main()