scripts/hf_jobs/hf_latency.py

# /// script
# requires-python = ">=3.12"
# dependencies = ["torch>=2.5.0", "torchvision", "Pillow", "numpy"]
# ///
"""DiffThinker C5: Latency benchmark on HF Jobs. DiT forward + pipeline on T4 and A10G."""
import torch, json, os, time, numpy as np, torch.nn as nn, torch.nn.functional as F
from PIL import Image, ImageDraw

def make_cond(n=5):
    gs=8; data=[]
    for _ in range(n):
        img = Image.new("RGB",(64,64),(255,255,255)); draw=ImageDraw.Draw(img); cw=64//gs
        grid = np.zeros((gs,gs),dtype=np.uint8)
        for _ in range(int(gs*gs*0.15)):
            wx,wy=np.random.randint(1,gs-1,2); grid[wx,wy]=1
        for r in range(gs):
            for c in range(gs):
                if grid[r,c]==1: draw.rectangle([c*cw,r*cw,(c+1)*cw,(r+1)*cw],fill=(100,100,100))
        draw.rectangle([0,0,cw,cw], fill=(0,255,0))
        draw.rectangle([(gs-1)*cw,(gs-1)*cw,gs*cw,gs*cw], fill=(255,0,0))
        data.append(img)
    return data

class SimpleDiT(nn.Module):
    def __init__(self, ic=3, ims=64, ld=64):
        super().__init__(); self.ims=ims
        self.tp=nn.Linear(ld,ic); self.te=nn.Sequential(nn.Linear(1,ld),nn.SiLU(),nn.Linear(ld,ld))
        self.ce=nn.Sequential(nn.Conv2d(ic,16,3,padding=1),nn.SiLU(),nn.Conv2d(16,32,3,padding=1),nn.SiLU(),nn.Conv2d(32,ld,3,padding=1))
        self.d1=nn.Conv2d(ic+ic,32,3,padding=1); self.d2=nn.Conv2d(32,64,3,stride=2,padding=1)
        self.d3=nn.Conv2d(64,ld,3,stride=2,padding=1)
        self.mid=nn.Sequential(nn.Conv2d(ld,ld,3,padding=1),nn.SiLU(),nn.Conv2d(ld,ld,3,padding=1))
        self.u3=nn.ConvTranspose2d(ld,64,4,stride=2,padding=1)
        self.u2=nn.ConvTranspose2d(64,32,4,stride=2,padding=1); self.u1=nn.Conv2d(32,ic,3,padding=1)
    def forward(self,x,t,c):
        B=x.shape[0]; te=self.te(t.view(-1,1).float()/100.0); te=self.tp(te).view(B,-1,1,1).expand(-1,-1,self.ims,self.ims)
        h=torch.cat([x,te],dim=1); c=self.ce(c); c=F.interpolate(c,size=(self.ims,self.ims),mode='bilinear',align_corners=False)
        h=self.d1(h);h=self.d2(h);h=self.d3(h);h=self.mid(h);h=self.u3(h);h=self.u2(h);h=self.u1(h);return h

device=torch.device("cuda"); gpu=torch.cuda.get_device_name(0); mem=torch.cuda.get_device_properties(0).total_memory/1e9
print(f"GPU: {gpu} ({mem:.1f}GB)")

imgs = make_cond(5)
cond = torch.tensor(np.array([np.array(img).transpose(2,0,1) for img in imgs]), dtype=torch.float32).to(device)/255.0

models = {"SimpleDiT": SimpleDiT().to(device)}
for name, m in models.items():
    p = sum(p.numel() for p in m.parameters())
    print(f"\n--- {name}: {p:,} params ---")
    for _ in range(10):
        _ = m(torch.randn(1,3,64,64,device=device), torch.zeros(1,device=device,dtype=torch.long), cond[:1])
    fwd=[]
    for _ in range(30):
        x=torch.randn(1,3,64,64,device=device); t=torch.randint(0,100,(1,),device=device,dtype=torch.long)
        torch.cuda.synchronize(); t0=time.time(); _=m(x,t,cond[:1]); torch.cuda.synchronize(); fwd.append(time.time()-t0)
    fwd=fwd[5:]; avg_f=sum(fwd)/len(fwd)
    pipe=[]
    for _ in range(10):
        x=torch.randn(1,3,64,64,device=device)
        torch.cuda.synchronize(); t0=time.time()
        with torch.no_grad():
            for s in range(20):
                tv=torch.full((1,),s*5,device=device,dtype=torch.long); vc=m(x,tv,cond[:1]); vu=m(x,tv,torch.zeros_like(cond[:1]))
                x=x+(1.0/20)*(vu+4.0*(vc-vu))
        torch.cuda.synchronize(); pipe.append(time.time()-t0)
    pipe=pipe[2:]; avg_p=sum(pipe)/len(pipe)
    print(f"  Forward: {avg_f*1000:.2f}ms | Pipeline (20 steps): {avg_p*1000:.1f}ms | Per-NFE: {avg_p/20*1000:.2f}ms")

print(f"\nVRAM: {torch.cuda.memory_allocated(0)/1e9:.2f}GB")
print(json.dumps({"gpu":gpu,"gpu_mem_gb":round(mem,1),"forward_avg_ms":round(avg_f*1000,2),"pipeline_avg_ms":round(avg_p*1000,1),"params":sum(p.numel() for p in models['SimpleDiT'].parameters())}))