Spaces:
Running on Zero
Running on Zero
| 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 | |
| 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) | |
| 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) | |