Spaces:
Running on Zero
Running on Zero
File size: 5,302 Bytes
1587925 23efe23 4ad0ef7 23efe23 4ad0ef7 1587925 23efe23 1587925 23efe23 | 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | import spaces # MUST come before torch
import os
import pathlib
import tempfile
import time
import gradio as gr
import torch
from huggingface_hub import hf_hub_download
from torchvision.utils import save_image
from dit import MFDiT
MODEL_ID = "yscript/agent-genpath"
INPUT_SIZE = 256
# Backbone hyper-parameters exactly as in configs/sample.json from the paper repo.
BACKBONE = dict(
input_size=256,
patch_size=8,
in_channels=3,
dim=768,
depth=12,
num_heads=12,
num_classes=None,
)
# --- Model load (module scope, eager .to("cuda"); ZeroGPU intercepts) -----------
model = MFDiT(**BACKBONE)
_ckpt = hf_hub_download(MODEL_ID, "dit.pth")
_state = torch.load(_ckpt, map_location="cpu", weights_only=True)
_state = {k.removeprefix("module."): v for k, v in _state.items()}
model.load_state_dict(_state, strict=True)
model.to("cuda").eval()
del _state
@torch.no_grad()
def euler_sample(batch, steps, seed):
"""Reference Euler sampler from generation/sample.py (Agent-GenPath repo)."""
torch.manual_seed(seed)
device = next(model.parameters()).device
x = torch.randn(batch, 3, INPUT_SIZE, INPUT_SIZE, device=device)
for i in range(steps):
t = torch.full((batch,), 1 - i / steps, device=device)
x = x - model(x, t, torch.zeros_like(t), y=None) / steps
return (x.clamp(-1, 1) + 1) / 2
def _duration(seed, batch, steps, *args, **kwargs):
# Positional args mirror generate()'s signature (Gradio also injects progress=,
# which lands in *args). Measured on ZeroGPU: ~14 ms per image per Euler step
# (showcase batch=8 x 50 steps = 5.5 s GPU time). ~2x headroom for dispatch
# and first-call warmup; capped well below for the slider maximum.
try:
batch = max(1, min(int(batch), 8))
steps = max(1, min(int(steps), 100))
except (TypeError, ValueError):
batch, steps = 8, 50
return min(60, 4 + 0.016 * steps * batch)
@spaces.GPU(duration=_duration)
def generate(seed=6657, batch=8, steps=50, progress=gr.Progress(track_tqdm=True)):
"""Sample synthetic pathology image patches with the pretrained Agent-GenPath flow-matching DiT.
Args:
seed: RNG seed for the whole batch; same seed + settings reproduces the same patches.
batch: number of 256x256 RGB patches to synthesize in one go (1-8).
steps: number of Euler integration steps (paper showcase uses 50).
progress: Gradio progress bar.
Returns:
List of file paths to the generated PNG patches.
"""
seed = int(seed)
batch = max(1, min(int(batch), 8))
steps = max(1, min(int(steps), 100))
t0 = time.perf_counter()
imgs = euler_sample(batch, steps, seed)
dt = time.perf_counter() - t0
out = (
pathlib.Path(tempfile.gettempdir())
/ f"agent-genpath-{seed}-{batch}-{steps}-{os.getpid()}"
)
out.mkdir(parents=True, exist_ok=True)
paths = []
for i, image in enumerate(imgs.cpu()):
p = out / f"sample_{i + 1:05d}.png"
save_image(image, p)
paths.append(str(p))
print(
f"[agent-genpath] seed={seed} batch={batch} steps={steps} took {dt:.1f}s",
flush=True,
)
return paths
CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""
with gr.Blocks() as demo:
gr.Markdown(
"""
# Agent-GenPath — Pathology Patch Synthesis
**Agent-Driven Generative Priors for Pathology Image Synthesis and Segmentation** ([arXiv 2611.06701](https://arxiv.org/abs/2611.06701))
Unconditional flow-matching Diffusion Transformer (DiT) trained on **TCGA-200K**, a
217,030-patch pathology dataset curated by the paper's pathology agent system. The same
backbone serves as a generative prior for downstream semantic segmentation. This demo
samples 256×256 synthetic H&E patches from the released checkpoint
([`yscript/agent-genpath`](https://huggingface.co/yscript/agent-genpath)) with the
authors' reference Euler sampler.
*Code: [script-Yang/Agent-GenPath](https://github.com/script-Yang/Agent-GenPath) · Paper: [arXiv 2611.06701](https://arxiv.org/abs/2611.06701)*
"""
)
with gr.Column(elem_id="col-container"):
with gr.Row():
seed = gr.Number(label="Seed", value=6657, precision=0)
batch = gr.Slider(1, 8, value=8, step=1, label="Batch size")
steps = gr.Slider(10, 100, value=50, step=5, label="Euler steps")
run = gr.Button("Sample patches", variant="primary")
gallery = gr.Gallery(
label="Synthetic pathology patches",
columns=4,
height=520,
format="png",
type="filepath",
buttons=["download", "fullscreen"],
)
gr.Examples(
examples=[[6657, 8, 50]],
inputs=[seed, batch, steps],
outputs=gallery,
fn=generate,
cache_examples=True,
cache_mode="lazy",
label="Paper showcase batch (seed 6657 · 8 patches · 50 Euler steps)",
)
run.click(generate, inputs=[seed, batch, steps], outputs=gallery, api_name="generate")
demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)
|