Update app.py

app.py

@@ -26,15 +26,16 @@ class ImageClassifier(nn.Module):
         output=self.pretrain_model(input)
         return output
         
-model = ImageClassifier()
-model.load_state_dict(torch.load('model-data_comet-torch-model.pth'))
-
 def predict(inp):
-    image_transform = transforms.Compose([ transforms.Resize(size=(224,224)), transforms.ToTensor()])
+    if not hasattr(predict, "model"):
+        # Load the model only if it hasn't been loaded yet
+        predict.model = ImageClassifier()
+        predict.model.load_state_dict(torch.load('model-data_comet-torch-model.pth'))
+        predict.image_transform = transforms.Compose([transforms.Resize(size=(224, 224)), transforms.ToTensor()])
     labels = ['normal', 'cancer']
-    inp = image_transform(inp).unsqueeze(dim=0)
+    inp = predict.image_transform(inp).unsqueeze(dim=0)
     with torch.no_grad():
-        prediction = torch.nn.functional.softmax(model(inp))
+        prediction = torch.nn.functional.softmax(predict.model(inp))
         confidences = {labels[i]: float(prediction.squeeze()[i]) for i in range(len(labels))}    
     return confidences
     
@@ -44,4 +45,4 @@ gr.Interface(fn=predict,
              title=title,
              description=description,
              article=article,
-             examples=['image-1.jpg', 'image-2.jpg']).launch()
+             examples=['image-1.jpg', 'image-2.jpg']).launch()