File size: 767 Bytes
2cc08ea
7779efa
 
2cc08ea
7779efa
2cc08ea
7779efa
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2cc08ea
 
 
 
 
 
7779efa
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
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;
    }
}