some polishes

app.py

@@ -3,20 +3,6 @@ from torch.nn import functional as F
 
 from model_loader import ModelType, type_to_transforms, type_to_loaded_model
 
-def predict_from_model(model_type, raw_image):
-    tfm = type_to_transforms[model_type]
-    model = type_to_loaded_model[model_type]
-    img = tfm(raw_image)
-    y = None
-    if model_type == ModelType.SYNTHETIC_DETECTOR_V2:
-        y = model.forward(img.unsqueeze(0).to("cpu"))
-    else:
-        y = model.forward(img[None, ...])
-    y_1 = F.softmax(y, dim=1)[:, 1].cpu().detach().numpy()
-    y_2 = F.softmax(y, dim=1)[:, 0].cpu().detach().numpy()
-    return {'created by AI': y_1.tolist(),
-            'created by human': y_2.tolist()}
-
 def get_y(model_type, model, image):
     if model_type == ModelType.SYNTHETIC_DETECTOR_V2:
         return model.forward(image.unsqueeze(0).to("cpu"))

requirements.txt

@@ -1,3 +1,4 @@
 timm==0.9.5
 torch==2.0.1
-torchvision==0.15.2
+torchvision==0.15.2
+pytorch-lightning==2.0.9