Provably Safe Generative Sampling with Constricting Barrier Functions

This repository contains the implementation of the constricting CBF-based constrained image generation (Section 5.2.1) from our paper.

Overview

We enforce hard pixel-level constraints on a pre-trained diffusion model (DDPM-bedroom-256) by solving a minimum-norm Quadratic Program (QP) at each sampling step. The QP synthesizes a feedback control input that maintains the sample within a constricting safety tube (Algorithm 1 in the paper), guaranteeing that the final generated image satisfies the specified constraints.

Setup

pip install torch diffusers cvxpy pillow numpy

A CUDA-capable GPU is recommended. The code falls back to CPU if unavailable.

Usage

1. Place your reference image as window.png in the same directory as the script. This is the image patch that will be enforced at the target location.

2. Run:

python cbf_constricting.py

3. The output image is saved to cbf_constricting_window.png.
4. logger.py compares various constricting schedules, and saves the results in a new directory called 'out'

How it works

The script implements Algorithm 1 from the paper in a single file:

1. Load pre-trained model: DDPM-bedroom-256 from Hugging Face Diffusers (no retraining or modification).
2. Define the constraint: A reference image patch is placed at pixel coordinates (40, 150) with a spatially-varying mask v(p) that smoothly decays at boundaries (Eq. 16).
3. Initialize the safety tube: Compute the initial relaxation \epsilon_0 based on the constraint violation of the noise sample x(T) (Theorem 4.1).
4. Sample with CBF guidance: At each of the 200 sampling steps:
   • Predict noise with the pre-trained U-Net
   • Compute the nominal (uncontrolled) update
   • Solve the per-pixel minimum-norm QP to obtain the control input u (Eq. 10)
   • Apply the control: x_{k-1} = x_uncontrolled - u·\Delta t
5. Save: The final image satisfies h(x) \geq 0 for all constrained pixels.

Parameters

Parameter	Value	Reference
Error tolerance e	0.005	Eq. 16
Class-K parameter \alpha	0.5	Section 5
Safety margin c	0.01	Section 5
Sampling steps K	200	Section 5.2.1
Edge decay width	5%	Section 5.2.1