Update app.py

app.py

@@ -16,12 +16,12 @@ spec.loader.exec_module(modeling)
 
 # Initialize the model
 from modeling import clip_lora_model
-model = clip_lora_model().to(device)
+emotion_model = clip_lora_model().to(device)
 
 # Load pretrained weights
-model_path = hf_hub_download(repo_id="PerceptCLIP/PerceptCLIP_Emotions", filename="perceptCLIP_Emotions.pth")
-model.load_state_dict(torch.load(model_path, map_location=device))
-model.eval()
+emotion_model_path = hf_hub_download(repo_id="PerceptCLIP/PerceptCLIP_Emotions", filename="perceptCLIP_Emotions.pth")
+emotion_model.load_state_dict(torch.load(emotion_model_path, map_location=device))
+emotion_model.eval()
 
 # Emotion label mapping
 idx2label = {
@@ -68,12 +68,18 @@ def predict_emotion(image):
     img = emo_preprocess(img)
 
     with torch.no_grad():
-        outputs = model(img)
+        mem_score = mem_model(img).item()
+        outputs = emotion_model(img)
         predicted = outputs.argmax(1).item()
 
     emotion = idx2label[predicted]
     emoji = emotion_emoji.get(emotion, "❓")  
-    return f"{emotion} {emoji}"
+    return f"{emotion} {emoji}", f"{mem_score:.4f}"
+    
+mem_model = clip_lora_model(output_dim=1).to(device)
+mem_model_path = hf_hub_download(repo_id="PerceptCLIP/PerceptCLIP_Memorability", filename="perceptCLIP_Memorability.pth")
+mem_model.load_state_dict(torch.load(mem_model_path, map_location=device))
+mem_model.eval()
 
 # Example images 
 example_images = [
@@ -87,7 +93,7 @@ example_images = [
 iface = gr.Interface(
     fn=predict_emotion,
     inputs=gr.Image(type="pil", label="Upload an Image"),
-    outputs=gr.Textbox(label="Emotion"),
+    outputs=[gr.Textbox(label="Emotion"), gr.Textbox(label="Memorability Score")],
     title="PerceptCLIP-Emotions",
     description="This model predicts the emotion evoked by an image.",
     examples=example_images