Update app.py

app.py

@@ -57,7 +57,7 @@ def calculate_fvu(original_input, reconstruction):
     fvu_val = (recon_error_variance / (variance + 1e-8)).mean()
     return fvu_val.item()
 
-def predict(input_img, top_k, concept, neg_concept, max_strength):
+def predict(input_img, top_k, concept, neg_concept, max_strength, add_error):
     """
     Main function to process an image, identify top concepts, and visualize concept manipulation.
     """
@@ -138,7 +138,9 @@ 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.decode(modified_latents) + error
+            modified_reconstructed = sae_model.decode(modified_latents) 
+            if add_error:
+                modified_reconstructed = modified_reconstructed + error
             
             # Normalize for comparison
             modified_reconstructed_norm = modified_reconstructed / modified_reconstructed.norm(dim=-1, keepdim=True)
@@ -197,6 +199,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Matryoshka Sparse Autoencoder (MSA
             concept_input = gr.Textbox(label="Concept to Manipulate", value="hair", placeholder="e.g., hair")
             neg_concept_input = gr.Textbox(label="Negative Concept (Optional)", placeholder="e.g., a frown")
             max_strength_slider = gr.Slider(minimum=1.0, maximum=20.0, value=10.0, step=0.5, label="Max Concept Strength")
+            add_error_checkbox = gr.Checkbox(value=True, label="Add error term to reconstruction")
             submit_btn = gr.Button("Analyze and Interpret", variant="primary")
         
         with gr.Column(scale=2):
@@ -210,7 +213,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Matryoshka Sparse Autoencoder (MSA
             ["bird.jpg", 10, "birds", "", 10.0],
             ["statue.jpg", 10, "statue", "humans", 10.0],
         ],
-        inputs=[image_input, top_k_slider, concept_input, neg_concept_input, max_strength_slider],
+        inputs=[image_input, top_k_slider, concept_input, neg_concept_input, max_strength_slider, add_error_checkbox],
         outputs=[output_image, output_bar_plot, output_line_plot],
         fn=predict,
         cache_examples=True