XiaoYun Zhang
clean up
9c54c90
raw
history blame contribute delete
No virus
767 Bytes
using System;
using TorchSharp;
public class AutoencoderKL : IDisposable
{
private 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);
}
public void Dispose()
{
_model.Dispose();
_model = null;
}
}