refactor and add DDPM and DDIMSampler

.gitattributes

@@ -22,6 +22,7 @@
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
@@ -32,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+cat.png filter=lfs diff=lfs merge=lfs -text

DDIMSampler.cs

@@ -0,0 +1,34 @@
+using TorchSharp;
+
+public class DDIMSampler
+{
+    private readonly DDPM _model;
+    private const int TIME_STEPS = 1000;
+    private readonly torch.Device _device;
+
+    public DDIMSampler(DDPM model, float scale = 9.0f)
+    {
+        _model = model;
+        _device = model.Device;
+    }
+
+    public torch.Tensor Sample(torch.Tensor img, torch.Tensor condition, torch.Tensor unconditional_condition, int steps = 50, float scale = 9.0f)
+    {
+        var gap = DDIMSampler.TIME_STEPS / steps;
+        using(var context = torch.enable_grad(false))
+        {
+            for(var i = DDIMSampler.TIME_STEPS-1; i >=0; i -= gap)
+            {
+                var t_cur = torch.full(1, i, dtype: torch.ScalarType.Int64, device: _device);
+                var t_prev = torch.full(1, i - gap >= 0? i - gap: 0, dtype: torch.ScalarType.Int64, device: _device);
+                (var e_t_uncond, var e_t) = _model.DiffusionModel(img, condition, unconditional_condition, t_cur);
+                var model_output = e_t_uncond + scale * (e_t - e_t_uncond);
+                e_t = _model.PredictEPSFromZANDV(img, t_cur, model_output);
+                var pred_x0 = _model.PredictStartFromZANDV(img, t_cur, model_output);
+                img = _model.QSample(pred_x0, t_prev, e_t);
+            }
+
+            return img;
+        }
+    }
+}

DDPM.cs

@@ -0,0 +1,48 @@
+using TorchSharp;
+
+public class DDPM
+{
+    private readonly torch.jit.ScriptModule _model;
+    public torch.Device Device {get;}
+    public DDPM(string modelPath, torch.Device device)
+    {
+        _model = TorchSharp.torch.jit.load(modelPath);
+        Device = device;
+        _model.to(Device);
+        _model.eval();
+    }
+
+    public (torch.Tensor e_T_Uncondition, torch.Tensor e_T) DiffusionModel(torch.Tensor img, torch.Tensor condition, torch.Tensor unconditional_condition, torch.Tensor t)
+    {
+        var x_in = torch.cat(new[] { img, img });
+        var condition_in = torch.cat(new[] { unconditional_condition, condition });
+        var t_in = torch.cat(new[] { t, t });
+        var res = _model.invoke<torch.Tensor>("diffusion_model", x_in, t_in, condition_in).chunk(2);
+        return (res[0], res[1]);
+    }
+
+    public torch.Tensor DecodeImage(torch.Tensor img)
+    {
+        return _model.invoke<torch.Tensor>("decode_image", img);
+    }
+
+    public torch.Tensor ClipEncoder(torch.Tensor tokenTensor)
+    {
+        return _model.invoke<torch.Tensor>("clip_encoder", tokenTensor);
+    }
+
+    public torch.Tensor QSample(torch.Tensor z, torch.Tensor t, torch.Tensor v)
+    {
+        return _model.invoke<torch.Tensor>("q_sample",z, t, v);
+    }
+
+    public torch.Tensor PredictEPSFromZANDV(torch.Tensor z, torch.Tensor t, torch.Tensor v)
+    {
+        return _model.invoke<torch.Tensor>("predict_eps_from_z_and_v", z, t, v);
+    }
+
+    public torch.Tensor PredictStartFromZANDV(torch.Tensor z, torch.Tensor t, torch.Tensor v)
+    {
+        return _model.invoke<torch.Tensor>("predict_start_from_z_and_v", z, t, v);
+    }
+}

Program.cs

@@ -3,11 +3,11 @@ using System.Collections.Generic;
 using System.IO;
 using System.Linq;
 using TorchSharp;
+
 torchvision.io.DefaultImager = new torchvision.io.SkiaImager();
 var device = TorchSharp.torch.device("cuda:0");
-var ddpm_v_sampler = TorchSharp.torch.jit.load("ddim_v_sampler.ckpt");
-ddpm_v_sampler.to(device);
-ddpm_v_sampler.eval();
+var ddpm = new DDPM("ddim_v_sampler.ckpt", device);
+var ddimSampler = new DDIMSampler(ddpm);
 
 var start_token = 49406;
 var end_token = 49407;
@@ -34,30 +34,12 @@ var unconditional_tokenTensor = torch.tensor(uncontional_tokens.ToArray(), dtype
 unconditional_tokenTensor = unconditional_tokenTensor.reshape((long)batch, -1);
 var img = torch.randn(batch, 4, 96, 96, dtype: torch.ScalarType.Float32, device: device);
 var t = torch.ones(batch, dtype: torch.ScalarType.Int32, device: device);
-var condition = ddpm_v_sampler.invoke("clip_encoder", tokenTensor);
-var unconditional_condition = ddpm_v_sampler.invoke("clip_encoder", unconditional_tokenTensor);
-Console.WriteLine(condition);
-var timesteps = 1000;
+var condition = ddpm.ClipEncoder(tokenTensor);
+var unconditional_condition = ddpm.ClipEncoder(unconditional_tokenTensor);
 var ddim_steps = 50;
-int gap = timesteps / ddim_steps;
-using(var context = torch.enable_grad(false))
-{
-    for(var i = timesteps-1; i >=0; i -= gap)
-    {
-        var t_cur = torch.full(batch, i, dtype: torch.ScalarType.Int64, device: device);
-        var t_prev = torch.full(batch, i - gap >= 0? i - gap: 0, dtype: torch.ScalarType.Int64, device: device);
-        img = (torch.Tensor)ddpm_v_sampler.invoke("ddim_sampler", img, condition, unconditional_condition, t_cur, t_prev);
-        Console.WriteLine($"step {i}");
-    }
-
-    var decoded_images = (torch.Tensor)ddpm_v_sampler.invoke("decode_image", img);
-    decoded_images = torch.clamp((decoded_images + 1.0) / 2.0, 0.0, 1.0);
-
-    for(int i = 0; i!= batch; ++i)
-    {
-        // c * h * w
-        var image = decoded_images[i];
-        image = (image * 255.0).to(torch.ScalarType.Byte).cpu();
-        torchvision.io.write_image(image, $"{i}.png", torchvision.ImageFormat.Png);
-    }
-}
+img = ddimSampler.Sample(img, condition, unconditional_condition, ddim_steps);
+var decoded_images = (torch.Tensor)ddpm.DecodeImage(img);
+decoded_images = torch.clamp((decoded_images + 1.0) / 2.0, 0.0, 1.0);
+var image = decoded_images[0];
+image = (image * 255.0).to(torch.ScalarType.Byte).cpu();
+torchvision.io.write_image(image, $"0.png", torchvision.ImageFormat.Png);

ddim_v_sampler.ckpt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d32dd77b16c5ab4b584f6827037039c6d484caa6d11e541b7f6c8b93bb30c8cc
-size 5216916628
+oid sha256:22b16b2fc18c3b20c0eb74ed49a8f1834388fbfd84a49110340943f22fd30fa1
+size 5216915007