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