More sample images, finished readme

README.md

@@ -16,14 +16,14 @@ We solve noisy linear inverse problems with diffusion models. The method is fast
 
 ## Getting started 
 
+Recommended environment: Python 3.11, Cuda 12, Conda. For lower verions please adjust the dependencies below.
+
 ### 1) Clone the repository
 
 ```
 git clone https://github.com/vivjay30/cdim
 
 cd cdim
-
-export PYTHONPATH=$PYTHONPATH:`pwd`
 ```
 
 ### 2) Install dependencies
@@ -37,3 +37,39 @@ pip install -r requirements.txt
 
 pip install torch==2.4.1+cu124 torchvision-0.19.1+cu124 --extra-index-url https://download.pytorch.org/whl/cu124
 ```
+
+## Inference Examples
+We recommend using the ddpm models from the diffusers library. These are better and can be run without any manual downloading. Model will be downloaded on first run. The run will produce a noisy_measurement.png and output.png. The ouptut directory can be passed as an argument. You can use kl optimization instead of l2 with the `--loss kl` flag.
+
+#### CelebHQ Inpainting Example
+
+`python inference.py sample_images/celebhq_00001.jpg 50 3 operator_configs/box_inpainting_config.yaml noise_configs/gaussian_noise_config.yaml google/ddpm-celebahq-256`
+ 
+#### LSUN Churches Gaussian Deblur Example
+`python inference.py sample_images/lsun_church.png 50 3 operator_configs/gaussian_blur_config.yaml noise_configs/gaussian_noise_config.yaml google/ddpm-church-256`
+ 
+#### Poisson Noise Example
+`python inference.py sample_images/celebhq_29999.jpg 50 3 operator_configs/identity_operator_config.yaml noise_configs/poisson_noise_config.yaml google/ddpm-celebahq-256 --loss kl --eta-type gradnorm`
+
+#### Discrete KL with Bimodal Noise Example
+An example to show discrete KL on a non-standard noise distirbution
+`python inference.py sample_images/celebhq_00001.jpg 200 1 operator_configs/box_inpainting_config.yaml noise_configs/bimodal_noise_config.yaml google/ddpm-celebahq-256 --loss categorical_kl --lambda-val 2 --eta-type gradnorm`
+ 
+## FFHQ and Imagenet Models
+These models are generally not as strong as the huggingface ddpm models, but are used for comparisons with baseline methods.
+
+From [this link](https://drive.google.com/drive/folders/1jElnRoFv7b31fG0v6pTSQkelbSX3xGZh?usp=sharing), download the checkpoints "ffhq_10m.pt" and "imagenet_256.pt" to models/
+
+#### Imagenet Super Resolution Example
+`python inference.py sample_images/imagenet_val_00002.png 50 3 operator_configs/super_resolution_config.yaml noise_configs/gaussian_noise_config.yaml models/imagenet_model_config.yaml`
+
+#### FFHQ Random Inpainting (Faster)
+Here we set T=25 and K=1 to show the algorithm running faster
+`python inference.py sample_images/ffhq_00010.png 25 1 operator_configs/random_inpainting_config.yaml noise_configs/gaussian_noise_config.yaml models/ffhq_model_config.yaml`
+
+## A note on Eta and Lambda schedules
+By default the model tries to use expected gradnorm to set the step size schedule (eta). The gradient magnitudes have been precomputed on the train set and are stored in `etas.json`. However, those values are only valid for the specific tasks and number of steps T and K. When using a different task or step configuration, we fall back to `--eta-type gradnorm` which performs individual step gradient normalization as the value of eta. You can always use that flag, which is a less efficient but more general method.
+
+In addition, the step size schedule eta is scaled by a constant value lambda (the proportionality constant). Getting eta and lambda correct is vital to good convergence, especially with fewer denoising and optimization steps. If you find that the model overfits (loss oscillates wildly) or undefits (loss doesn't go to 0 for KL or sigma^2 for L2), then you should tweak the argument `--lambda-val`. The best guess is printed out for you to use as a starting point.
+
+

cdim/eta_scheduler.py

@@ -1,7 +1,8 @@
 import json
 
 class EtaScheduler:
-    def __init__(self, method, task, T, K, loss_type, lambda_val=None):
+    def __init__(self, method, task, T, K, loss_type,
+                       noise_function, lambda_val=None):
         self.task = task
         self.T = T
         self.K = K
@@ -10,11 +11,19 @@ class EtaScheduler:
         self.method = method
 
         self.precomputed_etas = self._load_precomputed_etas()
+
         # Couldn't find expected gradnorm
         if not self.precomputed_etas and method == "expected_gradnorm":
             self.method = "gradnorm"
             print("Etas for this configuration not found. Switching to gradnorm.")
 
+
+        # Precomputed gradients are only for gaussian noise
+        if noise_function.name != "gaussian" and method == "expected_gradnorm":
+            self.method = "gradnorm"
+            print("Precomputed gradients are only for gaussian noise. Switching to gradnorm.")
+
+
         # Get the best lambda_val if it's not passed
         if self.lambda_val is None:
             if self.method == "expected_gradnorm":

inference.py

@@ -17,7 +17,6 @@ from cdim.diffusion.scheduling_ddim import DDIMScheduler
 from cdim.diffusion.diffusion_pipeline import run_diffusion
 from cdim.eta_scheduler import EtaScheduler
 
-# torch.manual_seed(7)
 
 def load_image(path):
     """
@@ -83,7 +82,7 @@ def main(args):
     save_to_image(noisy_measurement, os.path.join(args.output_dir, "noisy_measurement.png"))
 
     eta_scheduler = EtaScheduler(args.eta_type, operator.name, args.T,
-        args.K, args.loss, args.lambda_val)
+        args.K, args.loss, noise_function, args.lambda_val)
 
     t0 = time.time()
     output_image = run_diffusion(

models/imagenet_model_config.yaml

@@ -0,0 +1,20 @@
+# Defaults for image training.
+
+image_size: 256
+num_channels: 256
+num_res_blocks: 2
+channel_mult: ""
+learn_sigma: True
+class_cond: False
+use_checkpoint: False
+attention_resolutions: "32,16,8"
+num_heads: 4
+num_head_channels: 64
+num_heads_upsample: -1
+use_scale_shift_norm: True
+dropout: 0.0
+resblock_updown: True
+use_fp16: False
+use_new_attention_order: False
+
+model_path: models/imagenet256.pt

noise_configs/poisson_noise_config.yaml

@@ -1,2 +1,2 @@
 name: poisson
-rate: 0.1
+rate: 0.05