update app

app.py

@@ -3,6 +3,8 @@ import numpy as np
 import random
 import torch
 import spaces
+import base64
+from io import BytesIO
 from typing import Iterable
 from PIL import Image, ImageDraw
 from diffusers import FlowMatchEulerDiscreteScheduler
@@ -114,6 +116,19 @@ loaded = False
 DEFAULT_PROMPT = "Remove the red highlighted object from the scene"
 
 
+def b64_to_pil(b64_str: str) -> Image.Image | None:
+    """Helper to decode base64 string from JS into a PIL Image"""
+    if not b64_str or not b64_str.startswith("data:image"):
+        return None
+    try:
+        _, data = b64_str.split(',', 1)
+        image_data = base64.b64decode(data)
+        return Image.open(BytesIO(image_data)).convert("RGB")
+    except Exception as e:
+        print(f"Error decoding image: {e}")
+        return None
+
+
 def burn_boxes_onto_image(pil_image: Image.Image, boxes_json_str: str) -> Image.Image:
     """Burn red outline-only rectangles onto the image (no fill)."""
     import json
@@ -146,7 +161,7 @@ def burn_boxes_onto_image(pil_image: Image.Image, boxes_json_str: str) -> Image.
 
 @spaces.GPU
 def infer_object_removal(
-    source_image: Image.Image,
+    b64_str: str,
     boxes_json: str,
     prompt: str,
     seed: int = 0,
@@ -175,8 +190,9 @@ def infer_object_removal(
     print(f"Prompt: {prompt}")
     print(f"Boxes JSON received: '{boxes_json}'")
 
+    source_image = b64_to_pil(b64_str)
     if source_image is None:
-        raise gr.Error("Please upload an image first.")
+        raise gr.Error("Please upload an image first using the Bbox editor area.")
 
     import json
     try:
@@ -211,7 +227,8 @@ def infer_object_removal(
     return result, seed, marked
 
 
-def update_dimensions_on_upload(image):
+def update_dimensions_on_upload(b64_str: str):
+    image = b64_to_pil(b64_str)
     if image is None:
         return 1024, 1024
     original_width, original_height = image.size
@@ -259,12 +276,57 @@ label{font-weight:600!important;color:#333!important}
 ::-webkit-scrollbar-thumb{background:linear-gradient(135deg,#A855F7,#C084FC);border-radius:4px}
 ::-webkit-scrollbar-thumb:hover{background:linear-gradient(135deg,#9333EA,#A855F7)}
 
-#bbox-draw-wrap{position:relative;border:2px solid #C084FC;border-radius:12px;overflow:hidden;background:#1a1a1a;min-height:420px}
+#bbox-draw-wrap{position:relative;border:2px dashed #C084FC;border-radius:12px;overflow:hidden;background:#1a1a1a;min-height:420px;transition: border-color 0.2s ease;}
 #bbox-draw-wrap:hover{border-color:#A855F7}
 #bbox-draw-canvas{cursor:crosshair;display:block;margin:0 auto}
 .bbox-hint{background:rgba(168,85,247,.08);border:1px solid #E9D5FF;border-radius:8px;padding:10px 16px;margin:8px 0;font-size:.9rem;color:#6B21A8}
 .dark .bbox-hint{background:rgba(168,85,247,.15);border-color:rgba(168,85,247,.3);color:#C084FC}
 
+/* Custom Upload Prompt overlay */
+#upload-prompt {
+    position: absolute;
+    top: 50%; left: 50%;
+    transform: translate(-50%, -50%);
+    text-align: center;
+    z-index: 20;
+    background: rgba(255, 255, 255, 0.95);
+    padding: 40px;
+    border-radius: 12px;
+    box-shadow: 0 8px 32px rgba(168,85,247,0.15);
+    backdrop-filter: blur(8px);
+    width: 80%;
+    max-width: 400px;
+}
+.dark #upload-prompt {
+    background: rgba(30, 30, 30, 0.95);
+}
+#upload-prompt p {
+    margin: 0 0 20px 0;
+    font-family: 'Outfit', sans-serif;
+    color: #6B21A8;
+    font-weight: 500;
+    font-size: 1.1rem;
+}
+.dark #upload-prompt p { color: #DAB2FF; }
+
+#custom-upload-btn {
+    background: linear-gradient(90deg, #A855F7, #9333EA);
+    color: white;
+    border: none;
+    padding: 12px 28px;
+    font-size: 16px;
+    font-weight: 600;
+    border-radius: 8px;
+    cursor: pointer;
+    box-shadow: 0 4px 12px rgba(168,85,247,0.3);
+    transition: transform 0.2s, box-shadow 0.2s;
+    font-family: 'Outfit', sans-serif;
+}
+#custom-upload-btn:hover {
+    transform: translateY(-2px);
+    box-shadow: 0 6px 16px rgba(168,85,247,0.4);
+}
+
 .bbox-toolbar-section{
     display:flex;
     gap:8px;
@@ -324,6 +386,8 @@ label{font-weight:600!important;color:#333!important}
 .bbox-tb-undo:hover{background:#9333EA}
 .bbox-tb-clear{background:#be123c}
 .bbox-tb-clear:hover{background:#e11d48}
+.bbox-tb-change{background:#4b5563}
+.bbox-tb-change:hover{background:#6b7280}
 
 #bbox-status{position:absolute;top:10px;left:10px;background:rgba(0,0,0,.75);color:#00ff88;padding:5px 10px;border-radius:6px;font-family:'IBM Plex Mono',monospace;font-size:11px;z-index:10;display:none;pointer-events:none}
 #bbox-count{position:absolute;top:10px;right:10px;background:rgba(147,51,234,.85);color:#fff;padding:4px 10px;border-radius:6px;font-family:'IBM Plex Mono',monospace;font-size:11px;z-index:10;display:none}
@@ -341,15 +405,8 @@ label{font-weight:600!important;color:#333!important}
 }
 .dark #bbox-debug-count{color:#C084FC;background:rgba(168,85,247,.12)}
 
-#boxes-json-input{
-    max-height:0!important;
-    overflow:hidden!important;
-    margin:0!important;
-    padding:0!important;
-    opacity:0!important;
-    pointer-events:none!important;
-    position:absolute!important;
-    z-index:-1!important;
+.hidden-input {
+    display: none !important;
 }
 """
 
@@ -364,13 +421,18 @@ bbox_drawer_js = r"""
         const badge      = document.getElementById('bbox-count');
         const debugCount = document.getElementById('bbox-debug-count');
 
-        const btnDraw   = document.getElementById('tb-draw');
-        const btnSelect = document.getElementById('tb-select');
-        const btnDel    = document.getElementById('tb-del');
-        const btnUndo   = document.getElementById('tb-undo');
-        const btnClear  = document.getElementById('tb-clear');
-
-        if (!canvas || !wrap || !debugCount || !btnDraw) {
+        const btnDraw    = document.getElementById('tb-draw');
+        const btnSelect  = document.getElementById('tb-select');
+        const btnDel     = document.getElementById('tb-del');
+        const btnUndo    = document.getElementById('tb-undo');
+        const btnClear   = document.getElementById('tb-clear');
+        const btnChange  = document.getElementById('tb-change-img');
+        
+        const uploadPrompt = document.getElementById('upload-prompt');
+        const fileInput    = document.getElementById('custom-file-input');
+        const uploadBtn    = document.getElementById('custom-upload-btn');
+
+        if (!canvas || !wrap || !debugCount || !btnDraw || !fileInput) {
             console.log('[BBox] waiting for DOM...');
             setTimeout(initCanvasBbox, 250);
             return;
@@ -397,9 +459,7 @@ bbox_drawer_js = r"""
         const RED_STROKE_WIDTH = 3;
         const SEL_STROKE = 'rgba(0,120,255,0.95)';
 
-        function n2px(b) {
-            return {x1:b.x1*dispW, y1:b.y1*dispH, x2:b.x2*dispW, y2:b.y2*dispH};
-        }
+        function n2px(b) { return {x1:b.x1*dispW, y1:b.y1*dispH, x2:b.x2*dispW, y2:b.y2*dispH}; }
         function px2n(x1,y1,x2,y2) {
             return {
                 x1: Math.min(x1,x2)/dispW, y1: Math.min(y1,y2)/dispH,
@@ -436,43 +496,46 @@ bbox_drawer_js = r"""
             }
 
             const container = document.getElementById('boxes-json-input');
-            if (!container) {
-                console.warn('[BBox] #boxes-json-input not in DOM');
-                return;
-            }
+            if (!container) return;
             const targets = [
                 ...container.querySelectorAll('textarea'),
-                ...container.querySelectorAll('input[type="text"]'),
-                ...container.querySelectorAll('input:not([type])')
+                ...container.querySelectorAll('input:not([type="file"])')
             ];
             targets.forEach(el => {
-                const proto = el.tagName === 'TEXTAREA'
-                    ? HTMLTextAreaElement.prototype
-                    : HTMLInputElement.prototype;
+                const proto = el.tagName === 'TEXTAREA' ? HTMLTextAreaElement.prototype : HTMLInputElement.prototype;
                 const ns = Object.getOwnPropertyDescriptor(proto, 'value');
                 if (ns && ns.set) {
                     ns.set.call(el, jsonStr);
                     el.dispatchEvent(new Event('input',  {bubbles:true, composed:true}));
                     el.dispatchEvent(new Event('change', {bubbles:true, composed:true}));
-                    el.dispatchEvent(new Event('blur',   {bubbles:true, composed:true}));
                 }
             });
         }
 
-        function placeholder() {
-            ctx.fillStyle='#2a2a2a'; ctx.fillRect(0,0,dispW,dispH);
-            ctx.strokeStyle='#444'; ctx.lineWidth=2; ctx.setLineDash([8,4]);
-            ctx.strokeRect(20,20,dispW-40,dispH-40); ctx.setLineDash([]);
-            ctx.fillStyle='#888'; ctx.font='16px Outfit,sans-serif';
-            ctx.textAlign='center'; ctx.textBaseline='middle';
-            ctx.fillText('Upload an image above first',dispW/2,dispH/2-10);
-            ctx.font='13px Outfit'; ctx.fillStyle='#666';
-            ctx.fillText('Then draw red boxes on objects to remove',dispW/2,dispH/2+14);
+        function syncImageToGradio(dataUrl) {
+            const container = document.getElementById('hidden-image-b64');
+            if (!container) return;
+            const targets = [
+                ...container.querySelectorAll('textarea'),
+                ...container.querySelectorAll('input')
+            ];
+            targets.forEach(el => {
+                const proto = el.tagName === 'TEXTAREA' ? HTMLTextAreaElement.prototype : HTMLInputElement.prototype;
+                const ns = Object.getOwnPropertyDescriptor(proto, 'value');
+                if (ns && ns.set) {
+                    ns.set.call(el, dataUrl);
+                    el.dispatchEvent(new Event('input', {bubbles:true, composed:true}));
+                    el.dispatchEvent(new Event('change', {bubbles:true, composed:true}));
+                }
+            });
         }
 
         function redraw(tempRect) {
             ctx.clearRect(0,0,dispW,dispH);
-            if (!baseImg) { placeholder(); updateBadge(); return; }
+            if (!baseImg) { 
+                ctx.fillStyle='#1a1a1a'; ctx.fillRect(0,0,dispW,dispH);
+                updateBadge(); return; 
+            }
             ctx.drawImage(baseImg, 0, 0, dispW, dispH);
 
             boxes.forEach((b,i) => {
@@ -709,6 +772,57 @@ bbox_drawer_js = r"""
         canvas.addEventListener('touchend',   onUp,   {passive:false});
         canvas.addEventListener('touchcancel',(e)=>{e.preventDefault();dragging=false;redraw();},{passive:false});
 
+        // --- File Upload Logic ---
+        function processFile(file) {
+            if (!file || !file.type.startsWith('image/')) return;
+            const reader = new FileReader();
+            reader.onload = (event) => {
+                const dataUrl = event.target.result;
+                const img = new window.Image();
+                img.crossOrigin = 'anonymous';
+                img.onload = () => {
+                    baseImg = img;
+                    boxes.length = 0;
+                    window.__bboxBoxes = boxes;
+                    selectedIdx = -1;
+                    fitSize(img.naturalWidth, img.naturalHeight);
+                    syncToGradio(); redraw(); hideStatus();
+                    uploadPrompt.style.display = 'none';
+                    syncImageToGradio(dataUrl);
+                };
+                img.src = dataUrl;
+            };
+            reader.readAsDataURL(file);
+        }
+
+        uploadBtn.addEventListener('click', () => fileInput.click());
+        btnChange.addEventListener('click', () => fileInput.click());
+        
+        fileInput.addEventListener('change', (e) => {
+            processFile(e.target.files[0]);
+            e.target.value = ''; // Reset input to allow re-upload of same file
+        });
+
+        wrap.addEventListener('dragover', (e) => {
+            e.preventDefault();
+            wrap.style.borderColor = '#A855F7';
+            wrap.style.boxShadow = '0 0 15px rgba(168,85,247,0.3)';
+        });
+        wrap.addEventListener('dragleave', (e) => {
+            e.preventDefault();
+            wrap.style.borderColor = '';
+            wrap.style.boxShadow = '';
+        });
+        wrap.addEventListener('drop', (e) => {
+            e.preventDefault();
+            wrap.style.borderColor = '';
+            wrap.style.boxShadow = '';
+            if (e.dataTransfer.files.length) {
+                processFile(e.dataTransfer.files[0]);
+            }
+        });
+
+        // --- Toolbar Logic ---
         btnDraw.addEventListener('click',   ()=>setMode('draw'));
         btnSelect.addEventListener('click', ()=>setMode('select'));
 
@@ -742,55 +856,6 @@ bbox_drawer_js = r"""
             syncToGradio(); redraw(); hideStatus();
         });
 
-        let lastSrc = null;
-        setInterval(() => {
-            const imgs = document.querySelectorAll('#source-image-component img');
-            let el = null;
-            for (const img of imgs) {
-                if (img.src && img.src.length > 30 &&
-                    (img.src.startsWith('data:') ||
-                     img.src.startsWith('blob:') ||
-                     img.src.includes('/file=') ||
-                     img.src.includes('/upload') ||
-                     img.src.includes('localhost') ||
-                     img.src.includes('127.0.0.1') ||
-                     img.src.startsWith('http'))) {
-                    el = img;
-                    break;
-                }
-            }
-
-            if (el && el.src && el.src !== lastSrc) {
-                lastSrc = el.src;
-                const img = new window.Image();
-                img.crossOrigin = 'anonymous';
-                img.onload = () => {
-                    baseImg = img;
-                    boxes.length = 0;
-                    window.__bboxBoxes = boxes;
-                    selectedIdx = -1;
-                    fitSize(img.naturalWidth, img.naturalHeight);
-                    syncToGradio(); redraw(); hideStatus();
-                    console.log('[BBox] loaded image', img.naturalWidth, 'x', img.naturalHeight);
-                };
-                img.onerror = () => {
-                    console.warn('[BBox] image load failed for', el.src.substring(0,60));
-                };
-                img.src = el.src;
-            } else if (!el || !el.src) {
-                if (baseImg) {
-                    baseImg = null;
-                    boxes.length = 0;
-                    window.__bboxBoxes = boxes;
-                    selectedIdx = -1;
-                    fitSize(512,400); syncToGradio(); redraw(); hideStatus();
-                }
-                lastSrc = null;
-            }
-        }, 300);
-
-        setInterval(() => { syncToGradio(); }, 500);
-
         new ResizeObserver(() => {
             if (baseImg) { fitSize(baseImg.naturalWidth, baseImg.naturalHeight); redraw(); }
         }).observe(wrap);
@@ -809,21 +874,22 @@ with gr.Blocks(css=css, theme=purple_theme) as demo:
     gr.Markdown("# **QIE-Object-Remover-Bbox**", elem_id="main-title")
     gr.Markdown(
         "Perform diverse image edits using a specialized [LoRA](https://huggingface.co/prithivMLmods/QIE-2511-Object-Remover-v2). "
-        "Upload an image, draw red bounding boxes over the objects you want to remove, and click Remove Object. "
+        "Upload an image directly into the bounding box editor area below, draw red bounding boxes over the objects you want to remove, and click Remove Object. "
         "Multiple boxes supported. Select, move, resize or delete individual boxes.",
         elem_id="subtitle",
     )
 
     with gr.Row():
         with gr.Column(scale=1):
-            source_image = gr.Image(
-                label="Upload Image",
-                type="pil",
-                height=350,
-                elem_id="source-image-component",
+            
+            # --- We replace the visible gr.Image upload with a hidden textbox for the Base64 data ---
+            hidden_image_b64 = gr.Textbox(
+                elem_id="hidden-image-b64", 
+                elem_classes="hidden-input", 
+                container=False
             )
 
-            gr.Markdown("# **Bbox Edit Controller**")
+            gr.Markdown("### **Bbox Edit Controller**")
             gr.HTML(
                 '<div class="bbox-hint">'
                 "<b>Draw mode:</b> Click & drag to create red rectangles. "
@@ -832,9 +898,15 @@ with gr.Blocks(css=css, theme=purple_theme) as demo:
                 "</div>"
             )
 
+            # Insert custom upload markup right over the canvas
             gr.HTML(
                 """
                 <div id="bbox-draw-wrap">
+                    <div id="upload-prompt">
+                        <p>Drag and drop an image here</p>
+                        <button id="custom-upload-btn">Browse Files</button>
+                        <input type="file" id="custom-file-input" accept="image/*" style="display:none;" />
+                    </div>
                     <canvas id="bbox-draw-canvas" width="512" height="400"></canvas>
                     <div id="bbox-status"></div>
                     <div id="bbox-count"></div>
@@ -850,9 +922,11 @@ with gr.Blocks(css=css, theme=purple_theme) as demo:
                     <button id="tb-select" class="bbox-tb-select"       title="Select / move / resize">🔲 Select</button>
                     <div class="toolbar-divider"></div>
                     <span class="toolbar-label">Actions:</span>
-                    <button id="tb-del"    class="bbox-tb-del"          title="Delete selected box">✕ Delete Selected</button>
+                    <button id="tb-del"    class="bbox-tb-del"          title="Delete selected box">✕ Delete</button>
                     <button id="tb-undo"   class="bbox-tb-undo"         title="Remove last box">↩ Undo</button>
                     <button id="tb-clear"  class="bbox-tb-clear"        title="Remove all boxes">🗑 Clear All</button>
+                    <div class="toolbar-divider"></div>
+                    <button id="tb-change-img" class="bbox-tb-change"   title="Upload a different image">📸 Change Image</button>
                 </div>
                 """
             )
@@ -861,10 +935,9 @@ with gr.Blocks(css=css, theme=purple_theme) as demo:
 
             boxes_json = gr.Textbox(
                 value="[]",
-                visible=True,
-                interactive=True,
                 elem_id="boxes-json-input",
-                label="boxes-json",
+                elem_classes="hidden-input",
+                container=False
             )
 
             prompt = gr.Textbox(
@@ -904,24 +977,23 @@ with gr.Blocks(css=css, theme=purple_theme) as demo:
 
     run_btn.click(
         fn=infer_object_removal,
-        inputs=[source_image, boxes_json, prompt, seed, randomize_seed,
+        inputs=[hidden_image_b64, boxes_json, prompt, seed, randomize_seed,
                 guidance_scale, num_inference_steps, height_slider, width_slider],
         outputs=[result, seed, preview],
-        js="""(src, bj, p, s, rs, gs, nis, h, w) => {
+        js="""(b64, bj, p, s, rs, gs, nis, h, w) => {
             const boxes = window.__bboxBoxes || [];
             const json  = JSON.stringify(boxes);
             console.log('[BBox] submitting', boxes.length, 'boxes:', json);
-            return [src, json, p, s, rs, gs, nis, h, w];
+            return [b64, json, p, s, rs, gs, nis, h, w];
         }""",
     )
 
-    source_image.upload(
+    hidden_image_b64.change(
         fn=update_dimensions_on_upload,
-        inputs=[source_image],
+        inputs=[hidden_image_b64],
         outputs=[width_slider, height_slider],
     )
 
-
 if __name__ == "__main__":
     demo.launch(
         css=css, theme=purple_theme,