test

app.py

@@ -160,6 +160,16 @@ def inpaint_image(image, prompt, subject, editor_value):
     target_text_prompt=prompt
     prompt_final=f'A two side-by-side image of same {subject_name}. LEFT: a photo of the {subject_name}; RIGHT: a photo of the {subject_name} {target_text_prompt}.'
 
+    # Convert attention mask to PIL image format
+    # Take first head's mask after prompt tokens (shape is now H*W x H*W)
+    attn_vis = full_attn_scale_mask[0, 0]
+    attn_vis[attn_vis <= 1.0] = 0
+    attn_vis[attn_vis > 1.0] = 255
+    attn_vis = attn_vis.cpu().float().numpy().astype(np.uint8)
+    # # Convert to PIL Image 
+    attn_vis_img = Image.fromarray(attn_vis)
+    attn_vis_img.save('attention_mask_vis.png')
+
     # Inpaint
     result = pipe(
         prompt=prompt_final,
@@ -175,7 +185,7 @@ def inpaint_image(image, prompt, subject, editor_value):
         true_guidance_scale=1.0,
         attn_scale_mask=full_attn_scale_mask,
     ).images[0]
-    return result, diptych_image
+    return result, attn_vis_img
 
 # Create Gradio interface with structured layout
 with gr.Blocks() as iface:
@@ -192,12 +202,12 @@ with gr.Blocks() as iface:
         with gr.Column():
             editor_value = gr.ImageEditor(type="pil", label="Image with Mask", sources="upload", visible=False)
             inpainted_image = gr.Image(type="pil", label="Inpainted Image")
-            diptych_image = gr.Image(type="pil", label="Diptych Image")
+            attn_vis_img = gr.Image(type="pil", label="Attn Vis Image")
     with gr.Row():
         
         inpaint_button = gr.Button("Inpaint")
 
-    inpaint_button.click(fn=inpaint_image, inputs=[input_image, prompt, subject, editor_value], outputs=[inpainted_image, diptych_image])
+    inpaint_button.click(fn=inpaint_image, inputs=[input_image, prompt, subject, editor_value], outputs=[inpainted_image, attn_vis_img])
 
 # Launch the app
 iface.launch()