Spaces:
Sleeping
Sleeping
| """Compare two uploaded images with the gated SushiUI style encoders.""" | |
| from __future__ import annotations | |
| import json | |
| import threading | |
| from pathlib import Path | |
| import gradio as gr | |
| import spaces | |
| import torch | |
| from huggingface_hub import get_hf_file_metadata, hf_hub_download, hf_hub_url | |
| from huggingface_hub.errors import GatedRepoError, HfHubHTTPError | |
| from PIL import Image | |
| from model import load_encoder | |
| MODEL_REPO = "celstk/style_embedder" | |
| MODEL_REVISION = "dbf4e0b4530876cdf150d4006625e8638bc338db" | |
| CALIBRATION = json.loads(Path(__file__).with_name("thresholds.json").read_text(encoding="utf-8")) | |
| MODELS: dict[str, tuple] = {} | |
| MODEL_LOCK = threading.RLock() | |
| def _check_image(image: Image.Image | None, name: str) -> Image.Image: | |
| if image is None: | |
| raise gr.Error(f"{name} の画像を入力してください。") | |
| if image.width * image.height > 40_000_000: | |
| raise gr.Error(f"{name} は 4,000 万画素以下の画像にしてください。") | |
| return image.convert("RGB") | |
| def _require_model_access(variant: str, token: str | None) -> None: | |
| if not token: | |
| raise gr.Error("Hugging Face にログインし、モデルへのアクセス承認を受けてください。") | |
| url = hf_hub_url(MODEL_REPO, f"{variant}/best_style_encoder.safetensors", | |
| revision=MODEL_REVISION) | |
| try: | |
| # Check every request: a shared model cache must not bypass the model's gate. | |
| get_hf_file_metadata(url, token=token) | |
| except (GatedRepoError, HfHubHTTPError) as exc: | |
| raise gr.Error("モデルへのアクセス権を確認できません。celstk/style_embedder でアクセスを申請し、" | |
| "承認後に再度ログインしてください。") from exc | |
| def _model_for(variant: str, token: str) -> tuple: | |
| cached = MODELS.get(variant) | |
| if cached is not None: | |
| return cached | |
| filename = f"{variant}/best_style_encoder.safetensors" | |
| path = hf_hub_download(MODEL_REPO, filename, revision=MODEL_REVISION, token=token) | |
| hf_hub_download(MODEL_REPO, f"{variant}/best_metadata.json", | |
| revision=MODEL_REVISION, token=token) | |
| model, processor, max_patches = load_encoder(path) | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| dtype = torch.bfloat16 if device.type == "cuda" else torch.float32 | |
| model = model.to(device, dtype=dtype).eval() | |
| cached = (model, processor, max_patches, device, dtype) | |
| MODELS[variant] = cached | |
| return cached | |
| def _embedding(image: Image.Image, model, processor, max_patches, device, dtype): | |
| inputs = processor(images=[image], max_num_patches=max_patches, return_tensors="pt") | |
| with torch.inference_mode(): | |
| return model(inputs["pixel_values"].to(device, dtype=dtype), | |
| inputs["pixel_attention_mask"].to(device), | |
| inputs["spatial_shapes"].to(device))[0].float() | |
| def compare(image_a: Image.Image | None, image_b: Image.Image | None, | |
| variant: str, target_recall: str, | |
| oauth_token: gr.OAuthToken | None = None) -> tuple[float, str]: | |
| if variant not in CALIBRATION or target_recall not in CALIBRATION[variant]["targets"]: | |
| raise gr.Error("無効なモデルまたは判定設定です。") | |
| token = oauth_token.token if oauth_token is not None else None | |
| _require_model_access(variant, token) | |
| first = _check_image(image_a, "A") | |
| second = _check_image(image_b, "B") | |
| with MODEL_LOCK: | |
| model, processor, max_patches, device, dtype = _model_for(variant, token) | |
| a = _embedding(first, model, processor, max_patches, device, dtype) | |
| b = _embedding(second, model, processor, max_patches, device, dtype) | |
| cosine = float((a * b).sum().clamp(-1, 1)) | |
| operating_point = CALIBRATION[variant]["targets"][target_recall] | |
| threshold = operating_point["threshold"] | |
| verdict = "同一 style とみなす" if cosine >= threshold else "同一 style と判定しない" | |
| result = (f"### {verdict}\n\n" | |
| f"**Cosine 類似度:** `{cosine:.4f}` | **判定閾値:** `{threshold:.4f}`\n\n" | |
| f"この閾値は未学習作者の同作者画像ペア **{operating_point['same_recall']:.1%}** を拾い、" | |
| f"異作者ペアも **{operating_point['different_acceptance']:.1%}** 受け入れた設定です。" | |
| "作者ラベルで校正した参考判定であり、画風の同一性を証明するものではありません。") | |
| return round(cosine, 4), result | |
| CSS = """ | |
| .gradio-container {max-width: 1060px !important; margin: auto !important;} | |
| .intro {line-height: 1.55;} | |
| .result-box {min-height: 170px;} | |
| """ | |
| def create_ui() -> gr.Blocks: | |
| with gr.Blocks(title="Style Embedder · A/B 比較") as demo: | |
| gr.Markdown("# Style Embedder · A/B 比較") | |
| gr.Markdown( | |
| "画像 A と B の画風ベクトルの cosine 類似度を測ります。まず " | |
| "[モデルのアクセス申請](https://huggingface.co/celstk/style_embedder)を行い、承認後にログインしてください。" | |
| "入力画像はこの Space の推論に使い、モデルリポジトリには送信しません。", | |
| elem_classes="intro", | |
| ) | |
| gr.LoginButton() | |
| with gr.Row(): | |
| image_a = gr.Image(type="pil", sources=["upload"], label="画像 A", height=320) | |
| image_b = gr.Image(type="pil", sources=["upload"], label="画像 B", height=320) | |
| with gr.Row(): | |
| variant = gr.Dropdown(choices=[("Baseline（推奨）", "baseline"), ("HSIC", "hsic")], | |
| value="baseline", label="モデル") | |
| target_recall = gr.Dropdown( | |
| choices=[("再現率 90%", "0.9"), ("再現率 95%", "0.95"), | |
| ("再現率 99%・見逃し重視", "0.99"), ("検証ペアの見逃し 0", "1.0")], | |
| value="0.99", label="判定設定", | |
| info="同作者ペアの見逃しを減らすほど、異作者も同一と判定しやすくなります。", | |
| ) | |
| run = gr.Button("比較する", variant="primary") | |
| cosine = gr.Number(label="Style cosine 類似度", precision=4) | |
| result = gr.Markdown(elem_classes="result-box") | |
| run.click(compare, inputs=[image_a, image_b, variant, target_recall], | |
| outputs=[cosine, result], concurrency_limit=1) | |
| gr.Markdown( | |
| "判定は dataset 39 の未学習作者376人・490枚（同作者ペア298組）で校正しました。" | |
| "best checkpoint の選択にもこの検証集合を使用したため、独立した外部テスト結果ではありません。" | |
| "作者名や画像は Space のコード・モデルには含まれません。" | |
| ) | |
| demo.queue() | |
| return demo | |
| if __name__ == "__main__": | |
| create_ui().launch(css=CSS) | |