Update app.py

app.py

@@ -326,7 +326,7 @@ def create_intervene_part(pipe: HookedStableDiffusionXLPipeline, saes_dict, mean
             return images.images[0]
 
     @spaces.GPU
-    def image_mod(prompt, block_str, brush_index, strength, num_steps, input_image, guidance_scale=None, start_index=None, end_index=None):
+    def image_mod(prompt, block_str, brush_index, strength, num_steps, input_image, guidance_scale=None, start_index=None, end_index=None, block_start=None, block_end=None):
         block = block_str.split(" ")[0]
 
         mask = (input_image["layers"][0] > 0)[:, :, -1].astype(float)
@@ -373,7 +373,7 @@ def create_intervene_part(pipe: HookedStableDiffusionXLPipeline, saes_dict, mean
                     seed=42,
                     num_inference_steps=int(num_steps),
                     edit_fn= lambda input, output: hook(None, input, output),
-                    layers_for_edit_fn=[i for i in range(18, 57)],
+                    layers_for_edit_fn=[i for i in range(block_start, block_end)],
                     stream="image").images[0]
             else:
 
@@ -464,6 +464,9 @@ def create_intervene_part(pipe: HookedStableDiffusionXLPipeline, saes_dict, mean
                     with gr.Column(scale=3):
                         with gr.Row():
                             strength = gr.Number(value=10, label="Strength", minimum=-40, maximum=40, elem_id="strength", precision=2)
+                        with gr.Row():
+                            block_start = gr.Number(value=18, label="Block start", minimum=0, maximum=57, elem_id="block_start")
+                            block_finish = gr.Number(value=57, label="Block finish", minimum=0, maximum=57, elem_id="block_start")
                         with gr.Row():
                             button = gr.Button('Apply', elem_id="apply_button")
 
@@ -493,7 +496,7 @@ def create_intervene_part(pipe: HookedStableDiffusionXLPipeline, saes_dict, mean
         num_steps.change(update_index_maxes, [num_steps], [start_index, end_index])
         
         button.click(image_mod, 
-                    inputs=[prompt_field, block_select, brush_index, strength, num_steps, i_image, start_index, end_index], 
+                    inputs=[prompt_field, block_select, brush_index, strength, num_steps, i_image, start_index, end_index, block_start, block_finish], 
                     outputs=o_image)
         # button_icon.click(feature_icon, inputs=[block_select, brush_index], outputs=o_image)
         demo.load(image_gen, [prompt_field, num_steps], outputs=[image_state])