using StableDiffusionV2;
using System;
using System.IO;
using TorchSharp;

var batch = 1;
var device = torch.device("cuda:0");
torchvision.io.DefaultImager = new torchvision.io.SkiaImager();
var prompt = "a wild cute green cat";
var outputFolder = "Output";
if(!Directory.Exists(outputFolder))
{
    Directory.CreateDirectory(outputFolder);
}

var clipTokenizer = new ClipTokenizer("vocab.json", "merges.txt");
var tokens = clipTokenizer.Tokenize(prompt);
var uncontional_tokens = clipTokenizer.Tokenize("");
var tokenTensor = torch.tensor(tokens, dtype: torch.ScalarType.Int64, device: device);
var unconditional_tokenTensor = torch.tensor(uncontional_tokens, dtype: torch.ScalarType.Int64, device: device);
tokenTensor = tokenTensor.repeat(batch, 1);
unconditional_tokenTensor = unconditional_tokenTensor.repeat(batch, 1);

var clipEncoder = new ClipEncoder("clip_encoder.ckpt", device);
var img = torch.randn(batch, 4, 64, 64, dtype: torch.ScalarType.Float32, device: device);
var condition = clipEncoder.Forward(tokenTensor);
var unconditional_condition = clipEncoder.Forward(unconditional_tokenTensor);
clipEncoder.Dispose();

var ddpm = new DDPM("ddim_v_sampler.ckpt", device);
var ddimSampler = new DDIMSampler(ddpm);
var ddim_steps = 50;
img = ddimSampler.Sample(img, condition, unconditional_condition, ddim_steps);
ddpm.Dispose();

var vae = new AutoencoderKL("autoencoder_kl.ckpt", device);
var decoded_images = vae.Forward(img);
decoded_images = torch.clamp((decoded_images + 1.0) / 2.0, 0.0, 1.0);

for(int i = 0; i!= batch; ++i)
{
    var savedPath = Path.Join(outputFolder, $"{i}.png");
    var image = decoded_images[i];
    image = (image * 255.0).to(torch.ScalarType.Byte).cpu();
    torchvision.io.write_image(image, savedPath, torchvision.ImageFormat.Png);

    Console.WriteLine($"save image to {savedPath}, enjoy");
}