Update README.md

README.md

@@ -43,15 +43,20 @@ import importlib.util
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load model class
+# Load the model class definition dynamically
 class_path = hf_hub_download(repo_id="PerceptCLIP/PerceptCLIP_Emotions", filename="modeling.py")
-spec = importlib.util.spec_from_file_location("clip_lora_model", class_path)
-clip_lora_model = importlib.util.module_from_spec(spec)
-spec.loader.exec_module(clip_lora_model)
+spec = importlib.util.spec_from_file_location("modeling", class_path)
+modeling = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(modeling)
+
+# initialize a model
+ModelClass = modeling.clip_lora_model 
+model = ModelClass().to(device)
 
 # Load pretrained model
 model_path = hf_hub_download(repo_id="PerceptCLIP/PerceptCLIP_Emotions", filename="perceptCLIP_Emotions.pth")
-model = torch.load(model_path).to(device).eval()
+model.load_state_dict(torch.load(model_path, map_location=device))
+model.eval()
 
 # Emotion label mapping
 idx2label = {
@@ -82,7 +87,7 @@ image = emo_preprocess()(image).unsqueeze(0).to(device)
 # Run inference
 with torch.no_grad():
     outputs = model(image)
-    _, predicted = outputs.max(1)  # Get the class index
+    _, predicted = outputs.max(1)
 
 # Get emotion label
 predicted_emotion = idx2label[predicted.item()]