Upload folder using huggingface_hub

app.py

@@ -9,14 +9,14 @@ import torch.optim as optim
 import gradio as gr
 
 from utils import preprocess_img, preprocess_img_from_path, postprocess_img
-from vgg19 import VGG_19
+from vgg16 import VGG_16
 
 if torch.cuda.is_available(): device = 'cuda'
 elif torch.backends.mps.is_available(): device = 'mps'
 else: device = 'cpu'
 print('DEVICE:', device)
 
-model = VGG_19().to(device)
+model = VGG_16().to(device)
 for param in model.parameters():
     param.requires_grad = False
 

vgg16.py

@@ -0,0 +1,25 @@
+import torch.nn as nn
+import torchvision.models as models
+
+class VGG_16(nn.Module):
+    def __init__(self):
+        super(VGG_16, self).__init__()
+        self.model = models.vgg16(weights='DEFAULT').features[:30]
+        
+        for i, _ in enumerate(self.model):
+            if i in [4, 9, 16, 23]:
+                self.model[i] = nn.AvgPool2d(kernel_size=2, stride=2, padding=0)
+                
+    def forward(self, x):
+        features = []
+        
+        for i, layer in enumerate(self.model):
+            x = layer(x)
+            if i in [0, 5, 10, 17, 24]:
+                features.append(x)
+        return features
+    
+
+if __name__ == '__main__':
+    model = VGG_16()
+    print(model)

vgg19.py

@@ -17,4 +17,9 @@ class VGG_19(nn.Module):
             x = layer(x)
             if i in [0, 5, 10, 19, 28]:
                 features.append(x)
-        return features
+        return features
+    
+
+if __name__ == '__main__':
+    model = VGG_19()
+    print(model)