updated steering

app.py

@@ -8,24 +8,43 @@ UNICORN_NO_HORN_PATH = "unicorn_without_horn.png"
 
 # --- Model / head info ---
 NUM_HEADS = 64  # SDXL UNet attention heads per layer
+LAYER_CHOICES = [
+    'unet.down_blocks.0.attentions.0.transformer_blocks.0.attn2',
+    'unet.down_blocks.0.attentions.1.transformer_blocks.0.attn2',
+    'unet.down_blocks.1.attentions.0.transformer_blocks.0.attn2',
+    'unet.down_blocks.1.attentions.1.transformer_blocks.0.attn2',
+    'unet.down_blocks.2.attentions.0.transformer_blocks.0.attn2',
+    'unet.down_blocks.2.attentions.1.transformer_blocks.0.attn2',
+    'unet.mid_block.attentions.0.transformer_blocks.0.attn2',
+    'unet.up_blocks.1.attentions.0.transformer_blocks.0.attn2',
+    'unet.up_blocks.1.attentions.1.transformer_blocks.0.attn2',
+    'unet.up_blocks.1.attentions.2.transformer_blocks.0.attn2',
+    'unet.up_blocks.2.attentions.0.transformer_blocks.0.attn2',
+    'unet.up_blocks.2.attentions.1.transformer_blocks.0.attn2',
+    'unet.up_blocks.2.attentions.2.transformer_blocks.0.attn2',
+    'unet.up_blocks.3.attentions.0.transformer_blocks.0.attn2',
+    'unet.up_blocks.3.attentions.1.transformer_blocks.0.attn2',
+    'unet.up_blocks.3.attentions.2.transformer_blocks.0.attn2'
+]
 HEAD_CHOICES = [f"head_{i}" for i in range(NUM_HEADS)]
 
 
 # --- Callbacks ---
 
-def steer_spectacles(strength: float):
+def steer_spectacles(strength: int):
     """
     Simple placeholder:
     - For strength ~0, show no glasses on both sides.
     - For strength > 0, show original on the left, glasses on the right.
     """
-    if strength <= 0.01:
+    if strength == 0:
         return CAT_WITHOUT_GLASSES_PATH, CAT_WITHOUT_GLASSES_PATH
     else:
+        CAT_WITH_GLASSES_PATH = f"./cat_steering/Cat_step_{strength}.png.png"
         return CAT_WITHOUT_GLASSES_PATH, CAT_WITH_GLASSES_PATH
 
 
-def run_unicorn_ablation(selected_heads):
+def run_unicorn_ablation(selected_layer, selected_heads):
     """
     Given selected head labels (e.g., ["head_3", "head_17"]), return:
     - Unicorn with horn (original)
@@ -33,10 +52,21 @@ def run_unicorn_ablation(selected_heads):
 
     Enforce max of 3 heads.
     """
+    
+    layer = selected_layer.split('blocks.')[1].split('.attentions')[0]  
+       
     if selected_heads is None:
         selected_heads = []
     if len(selected_heads) > 3:
         selected_heads = selected_heads[:3]
+    if len(selected_heads) ==1:
+        path = f"./single_heads/layer{layer}__{selected_layer}_h{selected_heads[0].replace('head_', '')}.png" 
+    elif len(selected_heads) ==2:
+        path = f"./head_pairs/layer{layer}__{selected_layer}_h{selected_heads[0].replace('head_', '')}_{selected_heads[1].replace('head_', '')}.png" 
+    elif len(selected_heads) ==3:
+        path = f"./head_triples/layer{layer}__{selected_layer}_h{selected_heads[0].replace('head_', '')}_{selected_heads[1].replace('head_', '')}_{selected_heads[2].replace('head_', '')}.png"    
+    
+    UNICORN_NO_HORN_PATH = path
     # In a real experiment you'd use `selected_heads` to ablate SDXL heads.
     return UNICORN_WITH_HORN_PATH, UNICORN_NO_HORN_PATH
 
@@ -154,10 +184,10 @@ with gr.Blocks() as demo:
             )
 
         steer_slider = gr.Slider(
-            minimum=0.0,
-            maximum=1.0,
-            value=0.0,
-            step=0.05,
+            minimum=0,
+            maximum=35,
+            value=0,
+            step=5,
             label="Steer 😼 (CAA strength towards glasses)",
             info="Connect this to your actual CAA steering pipeline.",
             elem_id="cat_steer_slider",
@@ -206,6 +236,14 @@ with gr.Blocks() as demo:
                 elem_classes=["img-card"],
             )
 
+        layer_selector = gr.Dropdown(
+            choices=LAYER_CHOICES,
+            multiselect=False,
+            value=["unet.down_blocks.0.attentions.0.transformer_blocks.0.attn2"],
+            label="Attention layer to ablate",
+            info="Select an attention layer to ablate.",
+        )
+        
         head_selector = gr.Dropdown(
             choices=HEAD_CHOICES,
             multiselect=True,
@@ -213,10 +251,10 @@ with gr.Blocks() as demo:
             label="Attention heads to ablate (max 3)",
             info="Select up to three head indices (0–63). In this demo, images are fixed placeholders.",
         )
-
+        
         head_selector.change(
             fn=run_unicorn_ablation,
-            inputs=head_selector,
+            inputs=[layer_selector, head_selector],
             outputs=[unicorn_original, unicorn_ablated]
         )