Upload 4 files

app.py

@@ -394,25 +394,36 @@ class LightweightBrainDecoder(nn.Module):
 brain_decoder = None
 
 def get_brain_decoder():
-    """Get or train the brain decoder on first use"""
+    """Get pre-trained brain decoder (loads weights, no training needed)"""
     global brain_decoder
     if brain_decoder is not None:
         return brain_decoder
 
-    print("[Brain Imaging] Training lightweight decoder on Fashion-MNIST (CPU, ~30s)...")
+    decoder = LightweightBrainDecoder(latent_dim=16, num_steps=4)
+
+    # Try to load pre-trained weights (zero latency!)
+    import os
+    weights_path = os.path.join(os.path.dirname(__file__), "decoder.pth")
+    if os.path.exists(weights_path):
+        print("[Brain Imaging] Loading pre-trained decoder (instant!)...")
+        decoder.load_state_dict(torch.load(weights_path, map_location='cpu', weights_only=True))
+        decoder.eval()
+        brain_decoder = decoder
+        print("[Brain Imaging] Decoder ready (pre-trained weights loaded)")
+        return decoder
+
+    # Fallback: train from scratch if weights not found
+    print("[Brain Imaging] Pre-trained weights not found, training from scratch (~30s)...")
     t0 = time.time()
 
     from torchvision import datasets, transforms
     from torch.utils.data import DataLoader
 
-    decoder = LightweightBrainDecoder(latent_dim=16, num_steps=4)
     decoder.train()
-
     transform = transforms.Compose([transforms.ToTensor()])
     try:
         train_ds = datasets.FashionMNIST('./data', train=True, download=True, transform=transform)
     except Exception:
-        # Fallback: use random data if download fails
         print("[Brain Imaging] FashionMNIST download failed, using synthetic data")
         train_ds = torch.utils.data.TensorDataset(
             torch.rand(1000, 1, 28, 28),
@@ -922,8 +933,6 @@ with gr.Blocks(
             4. Listen to the AI's "heartbeat" — steady for normal, arrhythmic for attacks
             
             > 💡 **Try a normal prompt first, then a jailbreak prompt** to see the dramatic difference!
-            
-            > ⚠️ The decoder trains on first use (~30 seconds on CPU). Please be patient!
             """)
             
             with gr.Row():

decoder.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d27712da397e15c2921c2e4014438f2fe052cbfd30bbfac846cd5d38db77d527
+size 1701760