Update app.py

app.py

@@ -86,6 +86,9 @@ def generate():
                 for layer in data["past_key_values"]
             )
 
+        # Ensure model config has caching enabled
+        ee_model.config.use_cache = True
+
         with torch.no_grad():
             out = ee_model(
                 inputs_embeds=inputs_embeds,
@@ -95,13 +98,16 @@ def generate():
                 output_hidden_states=True,
             )
 
-        # Return final hidden state in sigma-space — client decrypts + runs lm_head
+        # Final hidden state (sigma-space) — client decrypts + runs lm_head
         last_hidden = out.hidden_states[-1]  # (1, seq_len, hidden)
 
-        new_past = [
-            [t.cpu().tolist() for t in layer]
-            for layer in out.past_key_values
-        ]
+        # Serialize KV cache — guard against None (some models/configs don't return it)
+        new_past = None
+        if out.past_key_values is not None:
+            new_past = [
+                [t.cpu().tolist() for t in layer]
+                for layer in out.past_key_values
+            ]
 
         return jsonify({
             "last_hidden": last_hidden.cpu().tolist(),