File size: 3,498 Bytes
1a18f22 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | """Smoke: port DeCo (PixNerDiT) or PixelDiT (PixDiT) into the PixDiff mask-concat scaffolding.
Backbone-agnostic decouple: build with in=img+cond, out defaults to in, take x_pred[:, :img_ch].
Usage: python smoke_backbone.py {deco|pixeldit}"""
import os, sys
sys.path.insert(0, "/home/wzhang/LSC/Code/NPJ")
import torch, torch.nn as nn
from torch.utils.data import DataLoader
from framework.synth.pixdiff.conditioning import build_conditioner
from framework.synth.pixdiff.data import MaskCondGenDataset
BK = sys.argv[1]
DECO = "/home/wzhang/LSC/Code/NPJ/sota/DeCo"
PIXELDIT = "/home/wzhang/LSC/Code/NPJ/sota/PixelDiT"
dev = "cuda"
DR = "/home/wzhang/LSC/Dataset/Segmentation/processed_unified"
ds = MaskCondGenDataset(DR, "medsegdb_isic2018", "holdout", img_size=256,
train_fraction=0.02, fraction_seed=0)
cond = build_conditioner("onehot", ds.num_classes).to(dev)
img_ch, K = ds.in_channels, cond.cond_channels
in_tot = img_ch + K
print(f"[{BK}] ds n={len(ds)} img_ch={img_ch} K={K} in_tot={in_tot}", flush=True)
if BK == "deco":
sys.path.insert(0, os.path.join(DECO, "src", "models", "transformer"))
from dit_c2i_DeCo import PixNerDiT
net = PixNerDiT(in_channels=in_tot, patch_size=16, num_groups=12, hidden_size=768,
hidden_size_x=32, num_blocks=13, num_cond_blocks=12, num_classes=1).to(dev)
elif BK == "pixeldit":
sys.path.insert(0, PIXELDIT)
from pixdit_core.pixeldit_c2i import PixDiT
net = PixDiT(in_channels=in_tot, num_groups=12, hidden_size=768, pixel_hidden_size=16,
patch_depth=12, pixel_depth=4, patch_size=16, num_classes=1).to(dev)
else:
raise SystemExit("backbone must be deco|pixeldit")
print(f"[{BK}] params={sum(p.numel() for p in net.parameters())/1e6:.1f}M", flush=True)
opt = torch.optim.AdamW(net.parameters(), lr=1e-4)
dl = DataLoader(ds, batch_size=4, shuffle=True, drop_last=True, num_workers=2)
it = iter(dl)
def get_batch():
global it
try: b = next(it)
except StopIteration: it = iter(dl); b = next(it)
return (b["image"], b["mask"]) if isinstance(b, dict) else (b[0], b[1])
net.train()
for step in range(20):
img, msk = get_batch(); img, msk = img.to(dev), msk.to(dev)
t = torch.sigmoid(torch.randn(img.size(0), device=dev) * 0.8 - 0.8).view(-1, 1, 1, 1)
e = torch.randn_like(img)
z = t * img + (1 - t) * e
v = (img - z) / (1 - t).clamp_min(5e-2)
c = cond(msk)
y = torch.zeros(img.size(0), dtype=torch.long, device=dev)
out = net(torch.cat([z, c], dim=1), t.flatten(), y)
assert out.dim() == 4 and out.shape[1] >= img_ch, f"bad out shape {tuple(out.shape)}"
x_pred = out[:, :img_ch]
v_pred = (x_pred - z) / (1 - t).clamp_min(5e-2)
loss = ((v - v_pred) ** 2).mean()
loss.backward(); opt.step(); opt.zero_grad()
if step % 5 == 0 or step == 19:
print(f"[{BK}] step {step:2d} loss {loss.item():.4f}", flush=True)
net.eval()
with torch.no_grad():
msk0 = msk[:2]; c0 = cond(msk0)
z = torch.randn(2, img_ch, 256, 256, device=dev)
ts = torch.linspace(0, 1, 11).tolist()
for i in range(10):
tc, dt = ts[i], ts[i + 1] - ts[i]
out = net(torch.cat([z, c0], dim=1), torch.full((2,), tc, device=dev),
torch.zeros(2, dtype=torch.long, device=dev))[:, :img_ch]
z = z + (out - z) / max(1 - tc, 5e-2) * dt
print(f"[{BK}] sample ok shape={tuple(z.shape)} range=({z.min():.2f},{z.max():.2f})", flush=True)
print(f"SMOKE_{BK.upper()}_PASS", flush=True)
|