fix: strings, add device check

app.py

@@ -1,3 +1,4 @@
+import ast
 import gradio as gr
 import torch
 import torch.nn as nn
@@ -25,7 +26,7 @@ def set_seed(seed):
 set_seed(42)
 
 flag = os.environ["FLAG"] if "FLAG" in os.environ else "fakeflag{placeholder}"
-device = "cpu"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 config = ViTConfig.from_pretrained("google/vit-base-patch16-224")
 config.num_labels = 2  # Binary classification
 
@@ -81,20 +82,57 @@ def predict_and_compare(image):
 
     model.eval()
     with torch.no_grad():
-        output = model(image_tensor).logits
+        output = model(image_tensor).logits  # type: ignore
         probabilities = torch.softmax(output, dim=1)[0]
         predicted_class_index = torch.argmax(probabilities).item()
 
         class_names = ["False", "True"]  # Assuming 0 index is False, 1 is True
-        predicted_class = class_names[predicted_class_index]
-        probability = probabilities[predicted_class_index].item()
-
-    return f"""{'SUCCESSFULLY AUTHENTICATED!!\nFLAG: '+flag if ssim_value>0.9 and predicted_class == 'True' else 'FAILED TO AUTHENTICATE :('}
-=====================
-
-Predicted: {predicted_class}
-Probability: {probability:.4f}
-SSIM with reference: {ssim_value:.4f}"""
+        predicted_class = class_names[predicted_class_index]  # type: ignore
+        probability = probabilities[predicted_class_index].item()  # type: ignore
+
+    astolfo = """░░░░░▒▒▒▒▒▓██▒▒▒▒█▒▒▒▒▒▒▒▒█▒▒▒▒▒▒▒▒▒▒▒▒░░░░░░░░░▒▓
+░▒▒▒▒▒▒▒▒▒██░░▒▒▒▓▓▒▒▒▒▒▒▓█▒▒▒░░░▒▒▒▒▒▒▒▒░░░░▒▒░░░
+▒▓▒▒▒▒▒▒▒▓█░░░▒▒▓▒▓▒▒▒▒▒▒▒█▓▒▒░░░░▒▒▒▒▒▒▒▒░░░░▒▓▓░
+█▓▒▒▒▒▒▒▒█▒░░░▒▒▒▒▒▒▒▒▒▒▒▒█▓▒░░░░░▒▒▒▒▒▒▒▒▒▒░░▒▒▒▒
+█▒▒▒▒▒▓▒▒▓░░▒▓▓▓▒░▒▒▒▒▒▒▒░▓█▒▒▓▓████▓▒▒▓▒▒▒▒▒░▒█▓▒
+█▒▒▒▒▒▒▒▒▒░░░░▒▒▓░░▒▒▒▒▒░░▒████████▒▒▓▒▓▓▒▒▒▒▒░▒▓▓
+▓▒▒▒▒▒▒▒▒▒▓▓▓▓▓▓▓░░░▒▒▒░░░▒▒░▓▓██▒▒▒░░▓▒▓▒▒▒▒▒▒▒▒▒
+▒▒▒▒▒▒▒▒████████░░░░░▒▒░░░░░░░▓▒░▒▓░░░▒▒▒▒▒▒▒▒▒▒▒▒
+▒▒▒▒▒▓▒▒▓░▒▓▒█▓▒░░░░░░░░░░░░░░░▒▒▒░░░░▒▒░▒▒▒▒▒▒▒▒▒
+▒▓▒▒▒▒▒▒▒░░▒▓▒▓▓░░░░▒░░░░░░░░░░░░░░░░░░▒▒▒▒▒▒▒▒▒▒▒
+▒▒▓▒▒▒▒▒▒░░░░░░░░░░░▓░░░░░░░░░░░░░░░░░░░▒▒▒▒▒▒▒▒▒▒
+▒▒█▓▒▒▒▓▒░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░▒▒▒▒▒▓▒▒▓▒
+▒▒▒█▓▓▒█▒░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░▒▒▒▒▒▒▒▒█▓
+▒▒▓▓██▒▓▓░░░░░░░░░░▒▒▓▓████████░░░░░░░░▒▒▒▒▒▒▒▓▓██
+▒▒▓▓██▒▓▓▒░░░░░░░▒█████████████░░░░░░░▒▓▒▒▒▒▒▓▓▒▓█
+▒▒▓█▓█▓▓▓▓▒░░░░░░░██▓▓▓▓▓▓▓▓▓▓▒░░░░░░▓▓▓▒▒▒▒▓█▓░░▓
+▒▒▒██▓▓▒▒▒▒▓▒░░░░░░▓▓▓▓▓▓▓▓▓▓▒░░░░░▓█▓▓▒▒▒▒▒▒█▒░░░
+▒▒▒███▓▒▒░░░▒▒▒▒░░░░░▒▓▓▓▓▓▒░░░░░▒███▓▓▒▒▒▒░▒█░░░░
+▒▒▒██▒█▓▒░░░░░░▒▒░░░░░░░░░░░░░▒▓████▓▓▒▒▒▒▒░▓░░░░░
+▒▒▒██░░█▓▒░░░░░░░░▒▓▒▒░░░░░░▒▓████▓▒▓▓▒▒▓▓▓▒▒░░░░░"""
+    success = f"{astolfo}\n\n🎉 SUCCESSFULLY AUTHENTICATED!! 🎉\n\nEntering shell...\n\nastolfo@asu:~$ cat /flag.txt\nFLAG: {flag}\nastolfo@asu:~$ exit\nlogout"
+
+    astolfo_sed = """██████████████████████████████████████████████████
+██████████████████▓██▓█▓█████▓██▓▓▓███████████████
+███████████████████████████▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓████████
+███��█████████████▓▓▓▓▓▓▓▓▓▓▓▓█▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓██
+██████████████▓▓▓▓▓▓▓▒▒░▒▓▒▒░░▒▒▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓██
+████████████▓▓▓▓▓▓▒▒░░▒▒░░░▒▒░░░░░▒▓▓▓▓▓▓▓▓▓▓▓▓▓▓█
+███████████████▓▒░░░░░░░░░░░░░░░░░░░▒▓▓▓▓▓▓▓▓▓▓▓▓▓
+██████████████▓░░░░░░░░░░░░░░░░░░░░░░░▓███▓▓▓▓▓▓▓▓
+█████████▓▓██▒░░░░░░░░░░░░░░░░░░░░░░░░░██▓▓▓▓▓▓▓▓█
+████████▓▓▓█▒░░░░░░░░▒░░░░░░░░░░░░░░░░░░██▓▓▓▓▓▓▓█
+████████▓██▒░░░░░░░░░░░▒░░░▒░░░░░░░░░░░░▒█▓▓▓▓▓▓▓█
+██████████▓░░░░░░░▒▒░░░░░░░▒░░░░░░░░░░░░░▓██▓▓▓▓▓█
+█████████▓▓░░░░░░▒▒▒▒░░░░░░░░▒▒▒░░░░░░░░░░▓▓▓▓▓▓▓█
+█████████▓░░░░░▒▓▓▒▓▓▓▒░░░░▒▓██▓▓▓░░░░▒░░░░▓▓▓▓▓▓█
+████████▓░░░░░░▒▓▒░░░░░░░░░░░░░░░▒░░░░░░░▒▒▓▓▓▓▓██
+████████░▓░░░░░▒▒▓░░░░░░░░░░░░░░▒▒░░░░░░░▓▓▓▓▓▓███
+███████▓█▓▒▒▒░░▒▒▒▒░░░░░░░░░░░░▓▓▒▒░░░░░░░▓▓▓▓████
+████████████░░▒▒██▒█▓▒░░░░░░░▓███▓▒▒░▒█▒▒▒▒███████"""
+    fail = f"{astolfo_sed}\n\n🥲 FAILED TO AUTHENTICATE 🥲"
+
+    return f"{success if ssim_value>0.9 and predicted_class == 'True' else fail}\n=====================\n\nPredicted: {predicted_class}\nProbability: {probability:.4f}\nSSIM with reference: {ssim_value:.4f}"
 
 
 iface = gr.Interface(