| """Single-image inference using the unmodified official DoUNet.""" |
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser() |
| source = p.add_mutually_exclusive_group(required=True) |
| source.add_argument('--repo-id') |
| source.add_argument('--ckpt', type=Path) |
| p.add_argument('--revision', default='main', help='Prefer the uploaded commit SHA') |
| p.add_argument('--image', required=True, type=Path) |
| p.add_argument('--device', default='auto') |
| p.add_argument('--seed', type=int, default=42) |
| p.add_argument('--disable-amp', action='store_true') |
| p.add_argument('--trust-checkpoint', action='store_true', help='Allow pickle loading for a trusted official checkpoint only') |
| a = p.parse_args() |
| import torch |
| from PIL import Image |
| from huggingface_hub import hf_hub_download |
| sys.path.insert(0, str(Path(__file__).resolve().parent / 'upstream')) |
| from models.dou import DoUNet, load_clip_to_cpu |
| from utils.options import TrainCfg |
| from dataset.dataset import translate_duplicate |
| from torchvision import transforms |
| torch.manual_seed(a.seed) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed_all(a.seed) |
| device = torch.device(('cuda' if torch.cuda.is_available() else 'cpu') if a.device == 'auto' else a.device) |
| path = a.ckpt or Path(hf_hub_download(a.repo_id, 'best.pth', revision=a.revision)) |
| cfg = TrainCfg(model_name='ViT-L/14', classnames=('real', 'generated'), learn_logit_scale=False, amp=not a.disable_amp) |
| backbone = load_clip_to_cpu(cfg).eval() |
| for param in backbone.parameters(): |
| param.requires_grad_(False) |
| model = DoUNet(cfg, backbone).to(device) |
| ckpt = torch.load(path, map_location='cpu', weights_only=False) |
| if not isinstance(ckpt, dict) or not isinstance(ckpt.get('model'), dict): |
| raise ValueError("Expected a checkpoint dictionary containing ckpt['model']") |
| |
| |
| report = model.load_state_dict(ckpt['model'], strict=False) |
| if report.missing_keys or report.unexpected_keys: |
| raise RuntimeError(f'Checkpoint mismatch: {report}. Resolve against the official release; do not ignore it.') |
| model.eval() |
| |
| transform = transforms.Compose([ |
| transforms.Lambda( |
| lambda img: translate_duplicate(img, 224)), |
| transforms.CenterCrop(224), |
| transforms.ToTensor(), |
| transforms.Normalize( |
| [0.48145466, 0.4578275, 0.40821073], [0.26862954, 0.26130258, 0.27577711]), |
| ]) |
| with Image.open(a.image) as img: |
| x = transform(img.convert('RGB')).unsqueeze(0).to(device) |
| with torch.inference_mode(), torch.amp.autocast('cuda', enabled=cfg.amp and device.type == 'cuda'): |
| probs = model(x).softmax(dim=-1)[0].float().cpu() |
| print(json.dumps({'label': ['real', 'generated'][probs.argmax().item()], |
| 'real_score': probs[0].item(), 'generated_score': probs[1].item(), |
| 'seed': a.seed, 'note': 'Softmax scores are not calibrated probabilities of authenticity.'})) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
