model load fixed but image upload persists

app.py

@@ -5,12 +5,27 @@ from PIL import Image
 import requests
 from huggingface_hub import hf_hub_download
 
+from PIL import Image
+import numpy as np
+
 # Load the model checkpoint from Hugging Face
 checkpoint_path = hf_hub_download(repo_id="ttoosi/resnet50_robust_face", filename="100_checkpoint.pt")
 
 # Initialize the model
 model = models.resnet50()
-model.load_state_dict(torch.load(checkpoint_path, map_location=torch.device('cpu'))['model'])  # Force model to load on CPU
+# change the num_classes to 500
+model.fc = torch.nn.Linear(model.fc.in_features, 500)
+checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))['model']
+# remove the prefix 'module.' from the keys
+# remove the prefix 'model.' from the keys that have it
+new_state_dict = {k.replace('module.', ''): v for k, v in checkpoint.items()}
+new_state_dict = {k.replace('model.', ''): v for k, v in new_state_dict.items()}
+new_state_dict = {k.replace('attacker.', ''): v for k, v in new_state_dict.items()}
+
+
+print(new_state_dict.keys())
+print('********************')
+model.load_state_dict(new_state_dict, strict=False)  # ignore Unexpected key(s) in state_dict: "normalizer.new_mean", "normalizer.new_std", "normalize.new_mean", "normalize.new_std". 
 model.eval()
 
 # Image preprocessing
@@ -18,11 +33,13 @@ preprocess = transforms.Compose([
     transforms.Resize(256),
     transforms.CenterCrop(224),
     transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), # vggface2
 ])
 
 # Function to make predictions
 def predict(image):
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)  # Convert to PIL Image if i
     image = preprocess(image).unsqueeze(0)  # Add batch dimension
     with torch.no_grad():
         output = model(image)  # Perform inference on CPU
@@ -30,7 +47,7 @@ def predict(image):
     return f"Predicted class: {predicted_class.item()}"
 
 # Create the Gradio interface
-iface = gr.Interface(fn=predict, inputs=gr.inputs.Image(type="pil"), outputs="text")
+iface = gr.Interface(fn=predict, inputs=gr.Image(type="numpy"), outputs="text")  # Updated from gr.inputs.Image to gr.Image
 
 # Launch the interface
 iface.launch()