feat: initialize local git repository and update project configurations

app.py

@@ -3,29 +3,20 @@ import os
 import subprocess
 import time
 
-# --- Dynamic Repository Clone ---
-# If the MambaEye source code isn't deployed directly alongside this app.py,
-# we clone it from GitHub before trying to import it.
 mamba_dir = os.path.join(os.path.dirname(__file__), "MambaEye")
 if not os.path.exists(mamba_dir) or not os.path.exists(os.path.join(mamba_dir, "mambaeye")):
     print("Cloning MambaEye repository from GitHub...", flush=True)
-    # Ensure any empty/partial directory is removed before cloning
     if os.path.exists(mamba_dir):
         import shutil
         shutil.rmtree(mamba_dir)
     subprocess.check_call(["git", "clone", "https://github.com/usingcolor/MambaEye.git", mamba_dir])
 
-# --- Dynamic Dependency Injection for HuggingFace Spaces ---
-# HuggingFace ZeroGPU builder environments lack `nvcc`. 
-# We intercept the import and softly compile mamba-ssm using CPU-fallback PyTorch natives 
-# so we pass the build requirements perfectly.
 try:
     import mamba_ssm
     import causal_conv1d
 except ImportError:
     print("Installing mamba_ssm and causal_conv1d in backend...", flush=True)
     env = os.environ.copy()
-    # Bypass CUDA extensions because we don't have nvcc locally or in standard Hub build container
     env["MAMBA_SKIP_CUDA_BUILD"] = "TRUE"
     env["CAUSAL_CONV1D_SKIP_CUDA_BUILD"] = "TRUE"
     subprocess.check_call(
@@ -33,7 +24,6 @@ except ImportError:
         env=env
     )
 
-# Add the cloned MambaEye repository to the Python path
 sys.path.append(os.path.join(os.path.dirname(__file__), "MambaEye"))
 
 import gradio as gr
@@ -46,13 +36,11 @@ from torchvision.models import ResNet50_Weights
 from huggingface_hub import hf_hub_download
 import spaces
 
-# MambaEye Imports
 from mambaeye.model import MambaEye
 from mambaeye.scan import generate_scan_positions
 from mambaeye.positional_encoding import sinusoidal_position_encoding_2d
 from mamba_ssm.utils.generation import InferenceParams
 
-# Global Configuration
 TARGET_CANVAS_SIZE = 512
 PATCH_SIZE = 16
 CATEGORIES = ResNet50_Weights.IMAGENET1K_V1.meta["categories"]
@@ -70,10 +58,105 @@ MODEL_CONFIG = {
 
 MODEL_REPO = "usingcolor/MambaEye-base"
 MODEL_FILENAME = "mambaeye_base_ft.pt"
-
-# Global Model Cache
 _GLOBAL_MODEL = None
 
+
+# --- HOVER SCRIPT INJECTION ---
+JS_HOVER_SCRIPT = """
+function() {
+    let overlay = document.getElementById('mamba-hover-overlay');
+    if (!overlay) {
+        overlay = document.createElement('div');
+        overlay.id = 'mamba-hover-overlay';
+        overlay.style.position = 'fixed';
+        overlay.style.pointerEvents = 'none';
+        overlay.style.border = '2px solid rgba(0, 102, 255, 0.8)';
+        overlay.style.backgroundColor = 'rgba(0, 102, 255, 0.2)';
+        overlay.style.zIndex = '99999';
+        overlay.style.display = 'none';
+        document.body.appendChild(overlay);
+    }
+
+    document.addEventListener('mousemove', (e) => {
+        let imgs = document.querySelectorAll('img');
+        let targetImg = null;
+        for (let img of imgs) {
+            if (img.closest('.gradio-image-hook')) {
+                if (img.src && !img.src.includes('data:image/svg')) {
+                    targetImg = img;
+                }
+            }
+        }
+        if (!targetImg) { overlay.style.display = 'none'; return; }
+        
+        let rect = targetImg.getBoundingClientRect();
+        if (e.clientX >= rect.left && e.clientX <= rect.right && e.clientY >= rect.top && e.clientY <= rect.bottom) {
+            let nw = targetImg.naturalWidth;
+            let nh = targetImg.naturalHeight;
+            if (nw === 0 || nh === 0) return;
+            
+            let cw = rect.width;
+            let ch = rect.height;
+            let imgRatio = nw / nh;
+            let containerRatio = cw / ch;
+            
+            let renderW, renderH, renderX, renderY;
+            if (imgRatio > containerRatio) {
+                renderW = cw;
+                renderH = cw / imgRatio;
+                renderX = 0;
+                renderY = (ch - renderH) / 2;
+            } else {
+                renderH = ch;
+                renderW = ch * imgRatio;
+                renderY = 0;
+                renderX = (cw - renderW) / 2;
+            }
+            
+            let relX = e.clientX - rect.left - renderX;
+            let relY = e.clientY - rect.top - renderY;
+            
+            if (relX >= 0 && relX <= renderW && relY >= 0 && relY <= renderH) {
+                let scale = renderW / nw;
+                let TARGET_CANVAS_SIZE = 512;
+                let ratio = Math.min(TARGET_CANVAS_SIZE / nw, TARGET_CANVAS_SIZE / nh);
+                
+                let origX = relX / scale;
+                let origY = relY / scale;
+                
+                let y_offset = (TARGET_CANVAS_SIZE - nw * ratio) / 2;
+                let x_offset = (TARGET_CANVAS_SIZE - nh * ratio) / 2;
+                
+                let canvas_y = origX * ratio + y_offset;
+                let canvas_x = origY * ratio + x_offset;
+                
+                let px = Math.floor(canvas_x / 16) * 16;
+                let py = Math.floor(canvas_y / 16) * 16;
+                
+                let start_orig_y = (py - y_offset) / ratio;
+                let start_orig_x = (px - x_offset) / ratio;
+                
+                let render_box_x = rect.left + renderX + start_orig_y * scale;
+                let render_box_y = rect.top + renderY + start_orig_x * scale;
+                
+                let size_scale = (16 / ratio) * scale;
+                
+                overlay.style.left = render_box_x + "px";
+                overlay.style.top = render_box_y + "px";
+                overlay.style.width = size_scale + "px";
+                overlay.style.height = size_scale + "px";
+                overlay.style.display = 'block';
+            } else {
+                overlay.style.display = 'none';
+            }
+        } else {
+            overlay.style.display = 'none';
+        }
+    });
+}
+"""
+# -----------------------------
+
 def get_model():
     global _GLOBAL_MODEL
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -82,9 +165,7 @@ def get_model():
         try:
             checkpoint_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILENAME)
             model = MambaEye(**MODEL_CONFIG)
-            
-            # Since this runs inside ZeroGPU worker, load directly to device
-            model.load_state_dict(torch.load(checkpoint_path, map_location=device))
+            model.load_state_dict(torch.load(checkpoint_path, map_location=device, weights_only=True))
             model.to(device)
             model.eval()
             _GLOBAL_MODEL = model
@@ -95,7 +176,6 @@ def get_model():
     return _GLOBAL_MODEL, device
 
 def transfer_inference_params(params, device):
-    """Recursively moves the KV cache state of MambaEye InferenceParams to CPU or CUDA."""
     if params is None or getattr(params, "key_value_memory_dict", None) is None:
         return params
         
@@ -106,7 +186,7 @@ def transfer_inference_params(params, device):
             params.key_value_memory_dict[k] = tuple(x.to(device) if isinstance(x, torch.Tensor) else x for x in v)
         elif isinstance(v, list):
             params.key_value_memory_dict[k] = [x.to(device) if isinstance(x, torch.Tensor) else x for x in v]
-        elif isinstance(v, dict): # E.g., layers map
+        elif isinstance(v, dict): 
             for k2, v2 in v.items():
                 if hasattr(v2, "to"): 
                     params.key_value_memory_dict[k][k2] = v2.to(device)
@@ -155,10 +235,18 @@ def extract_patch(canvas_tensor, px, py):
     return patch.flatten()
 
 def draw_patches_on_image(image_arr, positions, x_offset, y_offset, h, w):
-    img = Image.fromarray(image_arr).convert("RGB")
-    draw = ImageDraw.Draw(img)
+    img = np.array(image_arr)
+    
+    # Create the greyed-out ambient background
+    grey_base = Image.fromarray(img).convert("L").convert("RGB")
+    grey_np = np.array(grey_base) * 0.4 + np.full_like(grey_np, 160) # Note: broadcasting handles full_like internally safely via float math
+    grey_base_np = (np.array(grey_base).astype(float) * 0.4 + 160).clip(0, 255).astype(np.uint8)
+    
+    temp_img = Image.fromarray(grey_base_np)
+    orig_pil = Image.fromarray(img)
+    draw = ImageDraw.Draw(temp_img)
     
-    orig_w, orig_h = img.size
+    orig_w, orig_h = orig_pil.size
     ratio = min(TARGET_CANVAS_SIZE / orig_w, TARGET_CANVAS_SIZE / orig_h)
     
     for i, (px, py) in enumerate(positions):
@@ -166,8 +254,14 @@ def draw_patches_on_image(image_arr, positions, x_offset, y_offset, h, w):
         orig_x = (px - x_offset) / ratio
         orig_px_size = PATCH_SIZE / ratio
         
+        box = (int(orig_y), int(orig_x), int(orig_y + orig_px_size), int(orig_x + orig_px_size))
+        
+        # Paste original color into the highlighted region
+        patch_crop = orig_pil.crop(box)
+        temp_img.paste(patch_crop, box)
+        
         color = "red" if i == len(positions) - 1 else "blue"
-        draw.rectangle([orig_y, orig_x, orig_y + orig_px_size, orig_x + orig_px_size], outline=color, width=2)
+        draw.rectangle(box, outline=color, width=2)
         
         if i > 0:
             prev_py, prev_px = positions[i-1]
@@ -178,7 +272,7 @@ def draw_patches_on_image(image_arr, positions, x_offset, y_offset, h, w):
             center_curr = (orig_y + orig_px_size / 2, orig_x + orig_px_size / 2)
             draw.line([center_prev, center_curr], fill="blue", width=2)
             
-    return np.array(img), positions
+    return np.array(temp_img), positions
 
 def init_state_for_image(image):
     canvas_tensor, x_offset, y_offset, h, w = preprocess_image(image)
@@ -244,7 +338,6 @@ def run_auto_scan(image, scan_pattern, sequence_length):
     state['drawn_positions'] = positions
     state['sequence_length'] = sequence_length
     
-    # On ZeroGPU spaces securely move Tensors back to CPU State
     state['canvas_tensor'] = state['canvas_tensor'].cpu()
     state['inference_params'] = transfer_inference_params(inference_params, torch.device('cpu'))
     
@@ -256,7 +349,7 @@ def run_auto_scan(image, scan_pattern, sequence_length):
     return img_display, format_predictions(final_probs), state, f"Auto Scan Complete. Extracted {sequence_length} patches. Click to add more!"
 
 @spaces.GPU
-def on_click(evt: gr.SelectData, original_image, state):
+def process_click_inference(x_orig, y_orig, original_image, state):
     if original_image is None:
         return None, {"Upload Image": 1.0}, state, "Upload Image"
         
@@ -266,10 +359,8 @@ def on_click(evt: gr.SelectData, original_image, state):
         state = init_state_for_image(original_image)
         state['inference_params'] = InferenceParams(max_seqlen=4000, max_batch_size=1)
     
-    # Move InferenceParams back to the functional device correctly!
     state['inference_params'] = transfer_inference_params(state['inference_params'], device)
     
-    x_orig, y_orig = evt.index
     orig_h, orig_w = state['original_image'].shape[:2]
     ratio = min(TARGET_CANVAS_SIZE / orig_w, TARGET_CANVAS_SIZE / orig_h)
     
@@ -285,8 +376,8 @@ def on_click(evt: gr.SelectData, original_image, state):
     
     patch = extract_patch(state['canvas_tensor'], px, py).to(device)
     
-    img_seq = patch.unsqueeze(0).unsqueeze(0) # (1, 1, 768)
-    move_seq = move_emb.unsqueeze(0) # (1, 1, 512)
+    img_seq = patch.unsqueeze(0).unsqueeze(0)
+    move_seq = move_emb.unsqueeze(0)
     
     with torch.no_grad():
         out = model(img_seq, move_seq, inference_params=state['inference_params'])
@@ -297,7 +388,6 @@ def on_click(evt: gr.SelectData, original_image, state):
     state['drawn_positions'].append((px, py))
     state['sequence_length'] += 1
     
-    # Strip back to CPU for Gradio Session Memory
     state['inference_params'] = transfer_inference_params(state['inference_params'], torch.device('cpu'))
     
     img_display, _ = draw_patches_on_image(
@@ -305,27 +395,39 @@ def on_click(evt: gr.SelectData, original_image, state):
         state['x_offset'], state['y_offset'], state['h'], state['w']
     )
     
-    return img_display, format_predictions(final_probs), state, f"Added patch {state['sequence_length']} (Total {state['inference_params'].seqlen_offset} inference steps)."
+    return img_display, format_predictions(final_probs), state, f"Added patch {state['sequence_length']} (Total {state['inference_params'].seqlen_offset} steps)."
+
+def on_click(evt: gr.SelectData, original_image, state):
+    x_orig, y_orig = evt.index
+    return process_click_inference(x_orig, y_orig, original_image, state)
 
 def on_upload(image):
     if image is None:
-        return None, {"Waiting...": 1.0}, None, "Upload Image"
-    # Delay model load until auto-scan triggers, saving memory overhead in preloads
-    return image, {"Click Auto Scan or click the image": 1.0}, None, "Ready. You can Auto Scan or click."
+        return None, None, {"Waiting...": 1.0}, None, "Upload Image"
+    
+    # Pre-render the grey background immediately on upload
+    grey_base = Image.fromarray(image).convert("L").convert("RGB")
+    grey_base_np = (np.array(grey_base).astype(float) * 0.4 + 160).clip(0, 255).astype(np.uint8)
+    
+    return grey_base_np, image, {"Click Auto Scan or click the image": 1.0}, None, "Ready. You can Auto Scan or click."
 
 def on_clear(original_image):
     if original_image is None:
         return None, {"Cleared": 1.0}, None, "Cleared"
-    return original_image, {"Cleared": 1.0}, init_state_for_image(original_image), "Selections cleared. Ready for new patch sequence."
+        
+    grey_base = Image.fromarray(original_image).convert("L").convert("RGB")
+    grey_base_np = (np.array(grey_base).astype(float) * 0.4 + 160).clip(0, 255).astype(np.uint8)
+    
+    return grey_base_np, {"Cleared": 1.0}, init_state_for_image(original_image), "Selections cleared. Ready for new patch sequence."
 
-# Build the Gradio App Blocks
-with gr.Blocks(title="MambaEye Interactive Demo", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# MambaEye Interactive inference Demo")
-    gr.Markdown("This interface incorporates the full **MambaEye-base** model inference natively. Using **ZeroGPU** inference via PyTorch equivalents.")
+with gr.Blocks(title="MambaEye Interactive Demo") as demo:
+    gr.Markdown("# MambaEye Interactive Inference Demo")
+    gr.Markdown("This interface incorporates the full **MambaEye-base** model natively.")
     
     with gr.Row():
         with gr.Column(scale=2):
-            input_image = gr.Image(type="numpy", label="Upload and Select Patches", interactive=True)
+            # elem_classes targets the JS overlay script correctly
+            input_image = gr.Image(type="numpy", label="Upload and Select Patches", interactive=True, elem_classes="gradio-image-hook")
             
             with gr.Row():
                 scan_pattern = gr.Dropdown(
@@ -343,17 +445,13 @@ with gr.Blocks(title="MambaEye Interactive Demo", theme=gr.themes.Soft()) as dem
             model_output_label = gr.Label(label="MambaEye Output Predictions", num_top_classes=5)
             status_text = gr.Markdown("Status: Waiting for image upload...")
             
-            # Application State
             state = gr.State(None)
             original_image_state = gr.State(None)
             
-    # Event wiring
     input_image.upload(
         fn=on_upload,
         inputs=[input_image],
-        outputs=[input_image, model_output_label, state, status_text]
-    ).then(
-        fn=lambda img: img, inputs=[input_image], outputs=[original_image_state]
+        outputs=[input_image, original_image_state, model_output_label, state, status_text]
     )
     
     auto_btn.click(
@@ -373,6 +471,8 @@ with gr.Blocks(title="MambaEye Interactive Demo", theme=gr.themes.Soft()) as dem
         inputs=[original_image_state],
         outputs=[input_image, model_output_label, state, status_text]
     )
+    
+    demo.load(js=JS_HOVER_SCRIPT)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(theme=gr.themes.Soft(), ssr_mode=False)