| """4x4 grid of chest X-rays from the final global model of the X-ray diffusion tutorial. |
| |
| The first n samples drawn from one seed, not picked, with the tutorial's own scheduler and sampler (the full |
| 1000-step DDPM), shown as the tutorial's debug grids are: one min-max over the whole grid. Runs on the CPU so it |
| leaves the GPU to the queue. Also saves the samples as a plain tensor. |
| |
| python sample_xray.py <app custom dir> <FL_global_model.pt> <out.png> [n=16] [seed=0] [threads=8] |
| """ |
|
|
| import json |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import torch |
| from torchvision.utils import make_grid, save_image |
|
|
| app, ckpt, out = Path(sys.argv[1]), Path(sys.argv[2]), Path(sys.argv[3]) |
| n, seed, threads = (int(sys.argv[i]) if len(sys.argv) > i else d for i, d in ((4, 16), (5, 0), (6, 8))) |
| sys.path.insert(0, str(app)) |
| from models import get_model |
| from monai.inferers import DiffusionInferer |
| from monai.networks.schedulers import DDPMScheduler |
|
|
| torch.set_num_threads(threads) |
| config = json.loads((app / "config.json").read_text()) |
| model = get_model() |
| blob = torch.load(ckpt, map_location="cpu", weights_only=False) |
| model.load_state_dict(blob["model"] if isinstance(blob.get("model"), dict) else blob, strict=True) |
| model.eval() |
|
|
| scheduler = DDPMScheduler(num_train_timesteps=1000, schedule="scaled_linear_beta", clip_sample=False, |
| prediction_type="epsilon", beta_start=0.0015, beta_end=0.0195) |
| inferer = DiffusionInferer(scheduler=scheduler) |
| shape = [n, config["net_config"]["diffusion_model"]["in_channels"], *config["spatial_shape"]] |
| generator = torch.Generator().manual_seed(seed) |
| start = time.time() |
| with torch.no_grad(): |
| samples = inferer.sample(input_noise=torch.randn(shape, generator=generator), diffusion_model=model.diffusion_model, |
| scheduler=scheduler, save_intermediates=False, verbose=False) |
| print(f"sampled {n} in {time.time() - start:.0f} s; range {samples.min():.3f}..{samples.max():.3f}") |
| torch.save(samples.as_subclass(torch.Tensor).clone(), out.with_suffix(".pt")) |
| save_image(make_grid(samples, nrow=int(n**0.5), normalize=True, scale_each=False, padding=2), str(out)) |
| print("wrote", out) |
|
|
