Ultrasound EDM2 XL + S Autoguidance

This repository contains only the EMA weights used for generation: edm2_xl_ema.pt and edm2_s_ema.pt. The XL model is sampled with the S model as the autoguidance network.

This model is presented in A Foundational EDM2-Based Generative Model for High-Resolution Synthetic Fetal Ultrasound Imaging from Open Datasets.

Code: xfetus/fetal-ultrasound-edm2

Colab

# Clone EDM2 generation code
import os, sys, subprocess
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "huggingface_hub", "matplotlib"])
if not os.path.isdir("edm2"):
    subprocess.check_call(["git", "clone", "--depth", "1", "https://github.com/NVlabs/edm2.git"])
sys.path.insert(0, "edm2")

# Import libraries
from pathlib import Path
import json, torch, matplotlib.pyplot as plt
from huggingface_hub import snapshot_download
import dnnlib
from training.encoders import StandardRGBEncoder, StabilityVAEEncoder

# Helper function for loading diffusion UNet
def make_net(meta, dropout, dev):
    return dnnlib.util.construct_class_by_name(
        class_name="training.networks_edm2.Precond", model_channels=meta["model_channels"],
        dropout=dropout, use_fp16=True,
        **dict(interface, img_resolution=meta["img_resolution"], img_channels=meta["img_channels"], label_dim=meta["label_dim"]),
    ).to(dev).eval().requires_grad_(False)

# EDM sampling (image generation function) 
@torch.no_grad()
def edm_sampler(net, noise, labels, gnet, guidance=2.25, num_steps=32, sigma_min=0.002, sigma_max=80, rho=7):
    def denoise(x, t):
        dx = net(x, t, labels)
        gx = gnet(x.to(device), t.to(device), labels.to(device)).to(x.device)
        return gx.lerp(dx, guidance)
    steps = torch.arange(num_steps, device=noise.device, dtype=torch.float32)
    t = (sigma_max ** (1 / rho) + steps / (num_steps - 1) * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho
    t = torch.cat([t, t[:1] * 0])
    x = noise * t[0]
    for i, (tc, tn) in enumerate(zip(t[:-1], t[1:])):
        d = (x - denoise(x, tc)) / tc
        x_next = x + (tn - tc) * d
        if i < num_steps - 1:
            d2 = (x_next - denoise(x_next, tn)) / tn
            x = x + (tn - tc) * (d + d2) / 2
        else:
            x = x_next
    return x

# Huggingface repo
repo_dir = Path(snapshot_download("harveymannering/ultrasound-edm2"))
config = json.loads((repo_dir / "config.json").read_text())
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
interface = dict(img_resolution=None, img_channels=None, label_dim=None)

# Load the diffusion models
net = make_net(config["files"]["xl"], 0.10, device)
gnet = make_net(config["files"]["s"], 0.0, device)
net.load_state_dict(torch.load(repo_dir / config["files"]["xl"]["filename"], map_location="cpu"))
gnet.load_state_dict(torch.load(repo_dir / config["files"]["s"]["filename"], map_location="cpu"))
net.to(device)
gnet.to(device)

# Generate and display image
class_idx, num_steps =  1, 32
noise = torch.randn(1, 4, 64, 64, device=device)
labels = torch.eye(9, device=device)[[class_idx]]
latents = edm_sampler(net, noise, labels, gnet, guidance=2.25, num_steps=num_steps)
encoder = StabilityVAEEncoder(batch_size=1)
img = encoder.decode(latents)[0].permute(1, 2, 0).cpu().numpy()
plt.figure(figsize=(4, 4))
plt.imshow(img) 
plt.axis("off")
plt.show()