Upload folder using huggingface_hub

inference.py

@@ -3,12 +3,17 @@ import torch.nn.functional as F
 from PIL import Image
 import numpy as np
 import torchvision.transforms as T
+from transformers import AutoModel, AutoConfig
 from unet import UNet
 
 # Define model
-model = UNet(in_channels=1, out_channels=3)
-state_dict = torch.load("unet_epoch20.pth", map_location="cpu")
-model.load_state_dict(state_dict)
+# model = UNet(in_channels=1, out_channels=3)
+# state_dict = torch.load("unet_epoch20.pth", map_location="cpu")
+# model.load_state_dict(state_dict)
+# model.eval()
+
+config = AutoConfig.from_pretrained(".")
+model = AutoModel.from_pretrained(".", config=config)
 model.eval()
 
 transform = T.Compose([

model.safetensors

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