split model into small ones

.gitattributes

@@ -34,3 +34,5 @@ saved_model/**/* 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
+autoencoder_kl.ckpt filter=lfs diff=lfs merge=lfs -text
+clip_encoder.ckpt filter=lfs diff=lfs merge=lfs -text

AutoencoderKL.cs

@@ -0,0 +1,24 @@
+using TorchSharp;
+
+public class AutoencoderKL
+{
+    private readonly torch.jit.ScriptModule _model;
+    private readonly float _scale;
+    public torch.Device Device {get;}
+
+    public AutoencoderKL(string modelPath, torch.Device device, float scale = 0.18215f)
+    {
+        _model = TorchSharp.torch.jit.load(modelPath);
+        Device = device;
+        _model.to(Device);
+        _model.eval();
+        _scale = scale;
+    }
+
+    public torch.Tensor Forward(torch.Tensor tokenTensor)
+    {   
+        var context = torch.enable_grad(false);
+        tokenTensor = 1.0f / _scale * tokenTensor;
+        return (torch.Tensor)_model.forward(tokenTensor);
+    }
+}

ClipEnocder.cs

@@ -0,0 +1,20 @@
+using TorchSharp;
+
+public class ClipEncoder
+{
+    private readonly torch.jit.ScriptModule _model;
+    public torch.Device Device {get;}
+
+    public ClipEncoder(string modelPath, torch.Device device)
+    {
+        _model = TorchSharp.torch.jit.load(modelPath);
+        Device = device;
+        _model.to(Device);
+        _model.eval();
+    }
+
+    public torch.Tensor Forward(torch.Tensor tokenTensor)
+    {
+        return (torch.Tensor)_model.forward(tokenTensor);
+    }
+}

DDIMSampler.cs

@@ -1,3 +1,4 @@
+using System;
 using TorchSharp;
 
 public class DDIMSampler
@@ -15,17 +16,20 @@ public class DDIMSampler
     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;
+        var batch = img.shape[0];
+        
         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 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);
                 (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);
+                Console.WriteLine(img);
             }
 
             return img;

DDPM.cs

@@ -1,3 +1,4 @@
+using System;
 using TorchSharp;
 
 public class DDPM
@@ -21,16 +22,6 @@ public class DDPM
         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);

Program.cs

@@ -8,7 +8,8 @@ torchvision.io.DefaultImager = new torchvision.io.SkiaImager();
 var device = TorchSharp.torch.device("cuda:0");
 var ddpm = new DDPM("ddim_v_sampler.ckpt", device);
 var ddimSampler = new DDIMSampler(ddpm);
-
+var autoencoderKL = new AutoencoderKL("autoencoder_kl.ckpt", device);
+var clipEncoder = new ClipEncoder("clip_encoder.ckpt", device);
 var start_token = 49406;
 var end_token = 49407;
 var dictionary = new Dictionary<string, long>(){
@@ -20,7 +21,7 @@ var dictionary = new Dictionary<string, long>(){
     {"green", 1901},
 };
 
-var batch = 1;
+var batch = 2;
 
 var prompt = "a wild cute green cat";
 var tokens = prompt.Split(' ').Select(x => dictionary[x]).ToList();
@@ -29,17 +30,21 @@ tokens = tokens.Append(end_token).ToList();
 tokens = tokens.Concat(Enumerable.Repeat<long>(0, 77 - tokens.Count)).ToList();
 var uncontional_tokens = new[]{start_token, end_token}.Concat(Enumerable.Repeat(0, 75)).ToList();
 var tokenTensor = torch.tensor(tokens.ToArray(), dtype: torch.ScalarType.Int64, device: device);
-tokenTensor = tokenTensor.reshape((long)batch, -1);
+tokenTensor = tokenTensor.repeat(batch, 1);
 var unconditional_tokenTensor = torch.tensor(uncontional_tokens.ToArray(), dtype: torch.ScalarType.Int64, device: device);
-unconditional_tokenTensor = unconditional_tokenTensor.reshape((long)batch, -1);
+unconditional_tokenTensor = unconditional_tokenTensor.repeat(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.ClipEncoder(tokenTensor);
-var unconditional_condition = ddpm.ClipEncoder(unconditional_tokenTensor);
+var t = torch.full(new[]{batch, 1L}, value: batch, dtype: torch.ScalarType.Int32, device: device);
+var condition = clipEncoder.Forward(tokenTensor);
+var unconditional_condition = clipEncoder.Forward(unconditional_tokenTensor);
 var ddim_steps = 50;
 img = ddimSampler.Sample(img, condition, unconditional_condition, ddim_steps);
-var decoded_images = (torch.Tensor)ddpm.DecodeImage(img);
+var decoded_images = (torch.Tensor)autoencoderKL.Forward(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);
+
+for(int i = 0; i!= batch; ++i)
+{
+    var image = decoded_images[i];
+    image = (image * 255.0).to(torch.ScalarType.Byte).cpu();
+    torchvision.io.write_image(image, $"{i}.png", torchvision.ImageFormat.Png);
+}

autoencoder_kl.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0f5b15ed1a0f81a0ec4a274ac368a5f4fb84f0ce7c3676e683de527e69a59840
+size 334940269

clip_encoder.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ef9706f02a78b2cf93acff22f3036bc3e629d0a5b595c640ada1f73788826f37
+size 1416615515

ddim_v_sampler.ckpt

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