Diffusers unet working

cdim/diffusion/diffusion_pipeline.py

@@ -15,7 +15,8 @@ def run_diffusion(
         num_inference_steps: int = 1000,
         K=5,
         image_dim=256,
-        image_channels=3
+        image_channels=3,
+        model_type="diffusers"
     ):
     batch_size = noisy_observation.shape[0]
     image_shape = (batch_size, image_channels, image_dim, image_dim)
@@ -26,7 +27,8 @@ def run_diffusion(
 
     for i, t in tqdm(enumerate(scheduler.timesteps), total=len(scheduler.timesteps), desc="Processing timesteps"):
          # 1. predict noise model_output
-        model_output = model(image, t.unsqueeze(0).to(device))[:, :3]
+        model_output = model(image, t.unsqueeze(0).to(device))
+        model_output = model_output.sample if model_type == "diffusers" else model_output[:, :3]
 
         # 2. compute previous image: x_t -> x_t-1
         image = scheduler.step(model_output, t, image).prev_sample
@@ -38,7 +40,8 @@ def run_diffusion(
 
             with torch.enable_grad():
                 # Calculate x^hat_0
-                model_output = model(image, (t - t_skip).unsqueeze(0).to(device))[:, :3]
+                model_output = model(image, (t - t_skip).unsqueeze(0).to(device))
+                model_output = model_output.sample if model_type == "diffusers" else model_output[:, :3]
                 x_0 = (image - beta_prod_t_prev ** (0.5) * model_output) / alpha_prod_t_prev ** (0.5)
 
                 distance = operator(x_0) - noisy_observation
@@ -48,6 +51,6 @@ def run_diffusion(
                 print(loss.mean())
                 loss.mean().backward()
 
-            image -= 10 / torch.linalg.norm(image.grad) * image.grad
+            image -= 15 / torch.linalg.norm(image.grad) * image.grad
 
     return image

inference.py

@@ -7,6 +7,8 @@ from PIL import Image
 import numpy as np
 import torch
 
+from diffusers import DiffusionPipeline
+
 from cdim.noise import get_noise
 from cdim.operators import get_operator
 from cdim.image_utils import save_to_image
@@ -51,11 +53,18 @@ def main(args):
     operator_config["device"] = device
     operator = get_operator(**operator_config)
 
-    # Load the model
-    model_config = load_yaml(args.model_config)
-    model = create_model(**model_config)
-    model = model.to(device)
-    model.eval()
+    if args.model_config.endswith(".yaml"):
+        # Local model from DPS
+        model_type = "dps"
+        model_config = load_yaml(args.model_config)
+        model = create_model(**model_config)
+        model = model.to(device)
+        model.eval()
+
+    else:
+        # Huggingface diffusers model
+        model_type = "diffusers"
+        model = DiffusionPipeline.from_pretrained(args.model_config).to("cuda").unet
 
     # All the models have the same scheduler.
     # you can change this for different models
@@ -77,7 +86,8 @@ def main(args):
         model, ddim_scheduler,
         noisy_measurement, operator, noise_function, device,
         num_inference_steps=args.T,
-        K=args.K)
+        K=args.K,
+        model_type=model_type)
     print(f"total time {time.time() - t0}")
 
     save_to_image(output_image, os.path.join(args.output_dir, "output.png"))

requirements.txt

@@ -1,2 +1,5 @@
+diffusers==0.30.3
 numpy==2.1.2
 Pillow==11.0.0
+PyYAML==6.0.2
+tqdm==4.66.5