Simplify model loading - use direct CNN instead of PreTrainedModel wrapper

app.py

@@ -6,15 +6,16 @@ Deploy to Hugging Face Spaces
 
 import gradio as gr
 import torch
+import torch.nn as nn
 import numpy as np
 import plotly.graph_objects as go
-from transformers import AutoModel, AutoConfig
+from huggingface_hub import hf_hub_download
 import tempfile
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Constants
-MODEL_ID = "Tumo505/SSL-ECG-Classificcation"
+REPO_ID = "Tumo505/SSL-ECG-Classificcation"
 CLASS_LABELS = ["NORM", "MI", "STTC", "HYP", "CD"]
 CLASS_COLORS = {
     "NORM": "#90EE90",
@@ -24,16 +25,55 @@ CLASS_COLORS = {
     "CD": "#A29BFE"
 }
 
+# Define model architecture (1D CNN)
+class ECGClassifier(nn.Module):
+    def __init__(self, num_classes=5, num_leads=12, output_size=128):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            nn.Conv1d(num_leads, 32, kernel_size=7, padding=3),
+            nn.BatchNorm1d(32),
+            nn.ReLU(),
+            nn.MaxPool1d(2),
+            nn.Conv1d(32, 64, kernel_size=5, padding=2),
+            nn.BatchNorm1d(64),
+            nn.ReLU(),
+            nn.MaxPool1d(2),
+            nn.Conv1d(64, 128, kernel_size=3, padding=1),
+            nn.BatchNorm1d(128),
+            nn.ReLU(),
+            nn.AdaptiveAvgPool1d(1),
+            nn.Flatten(),
+            nn.Linear(128, output_size),
+        )
+        self.classifier = nn.Linear(output_size, num_classes)
+    
+    def forward(self, x):
+        embeddings = self.encoder(x)
+        logits = self.classifier(embeddings)
+        return logits
+
 # Load model
 model = None
 try:
     print("Loading model from Hub...")
-    model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True)
+    model = ECGClassifier(num_classes=len(CLASS_LABELS), num_leads=12, output_size=128)
+    
+    # Download weights from Hub
+    weights_path = hf_hub_download(repo_id=REPO_ID, filename="model.safetensors")
+    
+    # Load safetensors
+    from safetensors.torch import load_file
+    state_dict = load_file(weights_path)
+    
+    # Load weights into model
+    model.load_state_dict(state_dict, strict=False)
     model.to(device)
     model.eval()
-    print("Model loaded successfully")
+    print("✅ Model loaded successfully")
 except Exception as e:
-    print(f"Error loading model: {e}")
+    print(f"❌ Error loading model: {e}")
+    import traceback
+    traceback.print_exc()
 
 def predict_ecg(file_obj):
     """Main prediction function"""
@@ -95,8 +135,7 @@ def predict_ecg(file_obj):
         
         # Predict
         with torch.no_grad():
-            output = model(x)
-            logits = output["logits"][0].cpu().numpy()
+            logits = model(x)[0].cpu().numpy()
             probs = torch.softmax(torch.tensor(logits), dim=0).numpy()
         
         # Get prediction