from pathlib import Path
import argparse,hashlib,json,time
from PIL import Image
p=argparse.ArgumentParser();p.add_argument('--model',required=True);p.add_argument('--resolution',type=int,choices=[1024,2048],default=2048);p.add_argument('--output',default='output');p.add_argument('--revision',default='main');p.add_argument('--device',default='cuda');a=p.parse_args()
import torch
from torchvision import transforms
from transformers import AutoModelForImageSegmentation
ROOT=Path(__file__).resolve().parent;out=ROOT/a.output;out.mkdir(exist_ok=True)
source=Image.open(ROOT/'input.png').convert('RGB');transform=transforms.Compose([transforms.Resize((a.resolution,a.resolution)),transforms.ToTensor(),transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])])
model=AutoModelForImageSegmentation.from_pretrained(a.model,revision=a.revision,trust_remote_code=True).to(a.device).eval()
start=time.perf_counter()
with torch.inference_mode():pred=model(transform(source).unsqueeze(0).to(a.device))[-1].sigmoid().cpu()[0].squeeze()
alpha=transforms.ToPILImage()(pred).resize(source.size,Image.Resampling.BILINEAR);alpha.save(out/'alpha.png')
source.putalpha(alpha);source.save(out/'cutout.png')
record={'model_id':a.model,'revision_requested':a.revision,'preprocess_resolution':a.resolution,'input_size':list(source.size),'device':a.device,'seconds':time.perf_counter()-start,'model_run':True,'source_sha256':hashlib.sha256((ROOT/'input.png').read_bytes()).hexdigest()}
(out/'run.json').write_text(json.dumps(record,indent=2)+'\n')