File size: 1,315 Bytes
2c966e2 |
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 32 33 34 35 36 37 38 39 |
import torch
from modelfile import CViT
from huggingface_hub import hf_hub_download
def predict_with_model(saved_frames):
print("PyTorch Version:", torch.__version__)
print("Is CUDA Available:", torch.cuda.is_available())
if torch.cuda.is_available():
print("CUDA Version:", torch.version.cuda)
print("Available GPU:", torch.cuda.get_device_name(0))
else:
print("CUDA is not available. Ensure you have installed a CUDA-enabled version of PyTorch.")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
input_data = torch.tensor(saved_frames, dtype=torch.float32).to(device)
model_path = hf_hub_download(
repo_id="mhamza-007/cvit_deepfake_detection",
filename="cvit2_deepfake_detection_ep_50.pth"
)
model = CViT()
model.load_state_dict(torch.load(model_path, map_location=device, weights_only=True)['state_dict'])
model = model.to(device)
with torch.no_grad():
output = model(input_data)
predictions = torch.softmax(output, dim=1)
predicted_classes = torch.argmax(predictions, dim=1)
output = output.cpu()
predictions = predictions.cpu()
predicted_classes = predicted_classes.cpu()
print("Predicted Classes:", predicted_classes)
return predicted_classes |