Add attention visualization code

JTP_PILOT-e4-vit_so400m_patch14_siglip_384.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:2bac5c99e38e946b09b8813e28598783b2aabbea24ecafd04261142343185f69
-size 1754826116

JTP_PILOT2-2-e3-vit_so400m_patch14_siglip_384.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:565cbd6d3f453940c12d73aa2496bab102caf9f1c9a2a85433533c768df03555
-size 1796716928

JTP_PILOT2-e3-vit_so400m_patch14_siglip_384.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:0ac0e46bc773cfb486a83a79de9497566d91359e7962d225afdb7822dffc603d
-size 1796716928

app.py

@@ -1,11 +1,12 @@
 
-import json
+import msgspec
 import os
 import zipfile
 from io import BytesIO
 from tempfile import NamedTemporaryFile
 import tempfile
-
+import numpy as np
+import matplotlib.cm as cm
 import gradio as gr
 import pandas as pd
 from PIL import Image
@@ -23,8 +24,7 @@ from typing import Callable
 from functools import partial
 import spaces.config
 from spaces.zero.decorator import P, R
-
-torch.set_grad_enabled(False)
+from huggingface_hub import hf_hub_download
 
 def _dynGPU(
     fn: Callable[P, R] | None, duration: Callable[P, int], min=10, max=300, step=5
@@ -189,47 +189,130 @@ class GatedHead(torch.nn.Module):
 
 model.head = GatedHead(min(model.head.weight.shape), 9083)
 
-safetensors.torch.load_model(model, "JTP_PILOT2-2-e3-vit_so400m_patch14_siglip_384.safetensors")
+cached_model = hf_hub_download(
+    repo_id="RedRocket/JointTaggerProject",
+    subfolder="JTP_PILOT2",
+    filename="JTP_PILOT2-e3-vit_so400m_patch14_siglip_384.safetensors"
+)
+
+safetensors.torch.load_model(model, cached_model)
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
+if torch.cuda.is_available():
+    model.to(device='cuda', dtype=torch.float16, memory_format=torch.channels_last)
 model.eval()
 
-with open("tagger_tags.json", "r") as file:
-    tags = json.load(file) # type: dict
-allowed_tags = list(tags.keys())
+with open("tagger_tags.json", "rb") as file:
+    tags = msgspec.json.decode(file.read(), type=dict[str, int])
 
-for idx, tag in enumerate(allowed_tags):
-    allowed_tags[idx] = tag.replace("_", " ")
+for tag in tags.keys():
+    tags[tag.replace("_", " ")] = tags.pop(tag)
 
-sorted_tag_score = {}
+allowed_tags = list(tags.keys())
 
 @spaces.GPU(duration=6)
-def run_classifier(image, threshold):
-    global sorted_tag_score
+def run_classifier(image: Image.Image, threshold):
     img = image.convert('RGBA')
-    tensor = transform(img).unsqueeze(0).to(device)
+    tensor = transform(img).unsqueeze(0)
 
     with torch.no_grad():
-        probits = model(tensor)[0]
-        values, indices = probits.topk(250)
+        probits = model(tensor)[0] # type: torch.Tensor
+        values, indices = probits.cpu().topk(250)
+
+    tag_score = {allowed_tags[idx.item()]: val.item() for idx, val in zip(indices, values)}
 
-    tag_score = dict()
-    for i in range(indices.size(0)):
-        tag_score[allowed_tags[indices[i]]] = values[i].item()
     sorted_tag_score = dict(sorted(tag_score.items(), key=lambda item: item[1], reverse=True))
 
-    return create_tags(threshold)
+    return *create_tags(threshold, sorted_tag_score), img, sorted_tag_score
 
-def create_tags(threshold):
-    global sorted_tag_score
+def create_tags(threshold, sorted_tag_score: dict):
     filtered_tag_score = {key: value for key, value in sorted_tag_score.items() if value > threshold}
     text_no_impl = ", ".join(filtered_tag_score.keys())
     return text_no_impl, filtered_tag_score
 
 def clear_image():
-    global sorted_tag_score
-    sorted_tag_score = {}
-    return "", {}
+    return "", {}, None, {}, None
+
+@spaces.GPU(duration=5)
+def cam_inference(img, threshold, alpha, evt: gr.SelectData):
+    target_tag_index = tags[evt.value]
+    tensor = transform(img).unsqueeze(0)
+
+    gradients = {}
+    activations = {}
+
+    def hook_forward(module, input, output):
+        activations['value'] = output
+
+    def hook_backward(module, grad_in, grad_out):
+        gradients['value'] = grad_out[0]
+
+    handle_forward = model.norm.register_forward_hook(hook_forward)
+    handle_backward = model.norm.register_full_backward_hook(hook_backward)
+
+    probits = model(tensor)[0]
+ 
+    model.zero_grad()
+    probits[target_tag_index].backward(retain_graph=True)
+
+    with torch.no_grad():
+        patch_grads = gradients.get('value')
+        patch_acts = activations.get('value')
+    
+        weights = torch.mean(patch_grads, dim=1).squeeze(0)
+    
+        cam_1d = torch.einsum('pe,e->p', patch_acts.squeeze(0), weights)
+        cam_1d = torch.relu(cam_1d)
+    
+        cam = cam_1d.reshape(27, 27).detach().cpu().numpy()
+
+    handle_forward.remove()
+    handle_backward.remove()
+
+    return create_cam_visualization_pil(img, cam, alpha=alpha, vis_threshold=threshold), cam
+
+def create_cam_visualization_pil(image_pil, cam, alpha=0.6, vis_threshold=0.2):
+    """
+    Overlays CAM on image and returns a PIL image.
+
+    Args:
+        image_pil: PIL Image (RGB)
+        cam: 2D numpy array (activation map)
+        alpha: float, blending factor
+        vis_threshold: float, minimum normalized CAM value to show color
+
+    Returns:
+        PIL.Image.Image with overlay
+    """
+    if cam is None:
+        return image_pil
+    w, h = image_pil.size
+    size = max(w, h)
+
+    # Normalize CAM to [0, 1]
+    cam -= cam.min()
+    cam /= cam.max()
+
+    # Create heatmap using matplotlib colormap
+    colormap = cm.get_cmap('inferno')
+    cam_rgb = colormap(cam)[:, :, :3]  # RGB
+
+    # Create alpha channel
+    cam_alpha = (cam >= vis_threshold).astype(np.float32) * alpha  # Alpha mask
+    cam_rgba = np.dstack((cam_rgb, cam_alpha))  # Shape: (H, W, 4)
+    
+    # Coarse upscale for CAM output -- keeps "blocky" effect that is truer to what is measured
+    cam_pil = Image.fromarray((cam_rgba * 255).astype(np.uint8), mode="RGBA")
+    cam_pil = cam_pil.resize((216,216), resample=Image.Resampling.NEAREST)
+
+    # Model uses padded image as input, this matches attention map to input image aspect ratio
+    cam_pil = cam_pil.resize((size, size), resample=Image.Resampling.BICUBIC)
+    cam_pil = transforms.CenterCrop((h, w))(cam_pil)
+
+    # Composite over original
+    composite = Image.alpha_composite(image_pil, cam_pil)
+
+    return composite
 
 class ImageDataset(Dataset):
     def __init__(self, image_files, transform):
@@ -311,35 +394,84 @@ def process_zip(zip_file, threshold):
 
     return temp_file.name, df
 
-with gr.Blocks(css=".output-class { display: none; }") as demo:
+custom_css = """
+.output-class { display: none; }
+.inferno-slider input[type=range] {
+    background: linear-gradient(to right,
+        #000004, #1b0c41, #4a0c6b, #781c6d,
+        #a52c60, #cf4446, #ed6925, #fb9b06,
+        #f7d13d, #fcffa4
+    ) !important;
+    background-size: 100% 100% !important;
+}
+#image_container-image {
+    width: 100%;
+    aspect-ratio: 1 / 1;
+    max-height: 100%;
+}
+#image_container img {
+    object-fit: contain !important;
+}
+"""
+
+with gr.Blocks(css=custom_css) as demo:
     gr.Markdown("""
     ## Joint Tagger Project: JTP-PILOT² Demo **BETA**
-    This tagger is designed for use on furry images (though may very well work on out-of-distribution images, potentially with funny results).  A threshold of 0.2 is recommended.  Lower thresholds often turn up more valid tags, but can also result in some amount of hallucinated tags.
-
-    This tagger is the result of joint efforts between members of the RedRocket team, with distinctions given to Thessalo for creating the foundation for this project with his efforts, RedHotTensors for redesigning the process into a second-order method that models information expectation, and drhead for dataset prep, creation of training code and supervision of training runs.
-
-    Special thanks to Minotoro at frosting.ai for providing the compute power for this project.
     """)
     with gr.Tabs():
         with gr.TabItem("Single Image"):
+            original_image_state = gr.State() # stash a copy of the input image
+            sorted_tag_score_state = gr.State(value={}) # stash a copy of the input image
+            cam_state = gr.State()
             with gr.Row():
                 with gr.Column():
-                    image_input = gr.Image(label="Source", sources=['upload'], type='pil', height=512, show_label=False)
-                    threshold_slider = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold")
+                    image = gr.Image(label="Source", sources=['upload', 'clipboard'], type='pil', show_label=False, elem_id="image_container")
+                    cam_slider = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.40, label="CAM Threshold", elem_classes="inferno-slider")
+                    alpha_slider = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.60, label="CAM Alpha")
                 with gr.Column():
                     tag_string = gr.Textbox(label="Tag String")
+                    threshold_slider = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Tag Threshold")
                     label_box = gr.Label(label="Tag Predictions", num_top_classes=250, show_label=False)
 
-            image_input.upload(
+            image.upload(
                 fn=run_classifier,
-                inputs=[image_input, threshold_slider],
-                outputs=[tag_string, label_box]
+                inputs=[image, threshold_slider],
+                outputs=[tag_string, label_box, original_image_state, sorted_tag_score_state],
+                show_progress='minimal'
+            )
+
+            image.clear(
+                fn=clear_image,
+                inputs=[],
+                outputs=[tag_string, label_box, original_image_state, sorted_tag_score_state, cam_state]
             )
 
             threshold_slider.input(
                 fn=create_tags,
-                inputs=[threshold_slider],
-                outputs=[tag_string, label_box]
+                inputs=[threshold_slider, sorted_tag_score_state],
+                outputs=[tag_string, label_box],
+                show_progress='hidden'
+            )
+
+            label_box.select(
+                fn=cam_inference,
+                inputs=[original_image_state, cam_slider, alpha_slider],
+                outputs=[image, cam_state],
+                show_progress='minimal'
+            )
+
+            cam_slider.input(
+                fn=create_cam_visualization_pil,
+                inputs=[original_image_state, cam_state, alpha_slider, cam_slider],
+                outputs=[image],
+                show_progress='hidden'
+            )
+
+            alpha_slider.input(
+                fn=create_cam_visualization_pil,
+                inputs=[original_image_state, cam_state, alpha_slider, cam_slider],
+                outputs=[image],
+                show_progress='hidden'
             )
         
         with gr.TabItem("Multiple Images"):
@@ -357,6 +489,11 @@ with gr.Blocks(css=".output-class { display: none; }") as demo:
                 inputs=[zip_input, multi_threshold_slider],
                 outputs=[zip_output, dataframe_output]
             )
+    gr.Markdown("""
+    This tagger is designed for use on furry images (though may very well work on out-of-distribution images, potentially with funny results).  A threshold of 0.2 is recommended.  Lower thresholds often turn up more valid tags, but can also result in some amount of hallucinated tags.
+    This tagger is the result of joint efforts between members of the RedRocket team, with distinctions given to Thessalo for creating the foundation for this project with his efforts, RedHotTensors for redesigning the process into a second-order method that models information expectation, and drhead for dataset prep, creation of training code and supervision of training runs.
+    Special thanks to Minotoro at frosting.ai for providing the compute power for this project.
+    """)
 
 if __name__ == "__main__":
     demo.launch()

requirements.txt

@@ -3,4 +3,6 @@ torchvision
 timm
 pillow
 safetensors
-rarfile
+rarfile
+pydantic==2.10.6
+msgspec