Update app.py

app.py

@@ -77,6 +77,7 @@ def predict(input_img, top_k, concept, neg_concept, max_strength):
         reconstructed_features, _, full_latents = sae_model(image_features)
         
     fvu_score = calculate_fvu(image_features, reconstructed_features)
+    error = image_features - reconstructed_features
     
     # Get the top K activating SAE features for the image
     full_latents = full_latents.cpu().flatten()
@@ -137,7 +138,7 @@ def predict(input_img, top_k, concept, neg_concept, max_strength):
             modified_latents[:, concept_feature_id] = st
             
             # Decode the modified latents back into feature space
-            modified_reconstructed = sae_model.model.decode(modified_latents)
+            modified_reconstructed = sae_model.decode(modified_latents) + error
             
             # Normalize for comparison
             modified_reconstructed_norm = modified_reconstructed / modified_reconstructed.norm(dim=-1, keepdim=True)