Update app.py

app.py

@@ -18,6 +18,7 @@ from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
 from transformers import (
     Qwen3_5ForConditionalGeneration,
     AutoProcessor,
+    AutoModelForImageTextToText,
     TextIteratorStreamer,
 )
 
@@ -31,18 +32,35 @@ DTYPE = (
     else torch.float16
 )
 
-MODEL_NAME = "Qwen/Qwen3.5-2B"
+QWEN_MODEL_NAME = "Qwen/Qwen3.5-2B"
+LFM_MODEL_NAME = "LiquidAI/LFM2.5-VL-450M"
 CATEGORIES = ["Query", "Caption", "Point", "Detect"]
 
-print(f"Loading model: {MODEL_NAME} on {DEVICE}...")
+print(f"Loading Qwen model: {QWEN_MODEL_NAME} on {DEVICE}...")
 try:
     qwen_model = Qwen3_5ForConditionalGeneration.from_pretrained(
-        MODEL_NAME, torch_dtype=DTYPE, device_map=DEVICE,
+        QWEN_MODEL_NAME, torch_dtype=DTYPE, device_map=DEVICE,
     ).eval()
-    qwen_processor = AutoProcessor.from_pretrained(MODEL_NAME)
-    print("Model loaded successfully.")
+    qwen_processor = AutoProcessor.from_pretrained(QWEN_MODEL_NAME)
+    print("Qwen model loaded successfully.")
 except Exception as e:
-    print(f"Warning: Model loading failed (Check if weights exist). Error: {e}")
+    print(f"Warning: Qwen model loading failed. Error: {e}")
+    qwen_model = None
+    qwen_processor = None
+
+print(f"Loading LFM model: {LFM_MODEL_NAME} on {DEVICE}...")
+try:
+    lfm_model = AutoModelForImageTextToText.from_pretrained(
+        LFM_MODEL_NAME,
+        device_map="auto",
+        torch_dtype=torch.bfloat16,
+    ).eval()
+    lfm_processor = AutoProcessor.from_pretrained(LFM_MODEL_NAME)
+    print("LFM model loaded successfully.")
+except Exception as e:
+    print(f"Warning: LFM model loading failed. Error: {e}")
+    lfm_model = None
+    lfm_processor = None
 
 
 # --- Utility Functions ---
@@ -63,7 +81,7 @@ def safe_parse_json(text: str):
 
 # --- Inference Generator (Streaming) ---
 @spaces.GPU(duration=120)
-def generate_inference_stream(image: Image.Image, category: str, prompt: str):
+def generate_inference_stream(image: Image.Image, category: str, prompt: str, model_id: str = "qwen"):
     if category == "Query":
         full_prompt = prompt
     elif category == "Caption":
@@ -75,51 +93,102 @@ def generate_inference_stream(image: Image.Image, category: str, prompt: str):
     else:
         full_prompt = prompt
 
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image", "image": image},
-                {"type": "text", "text": full_prompt},
-            ],
-        }
-    ]
-    
-    text_input = qwen_processor.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-    
-    inputs = qwen_processor(
-        text=[text_input], images=[image], return_tensors="pt", padding=True
-    ).to(qwen_model.device)
-
-    streamer = TextIteratorStreamer(
-        qwen_processor.tokenizer,
-        skip_prompt=True,
-        skip_special_tokens=True,
-        timeout=120,
-    )
-    
-    thread = threading.Thread(
-        target=qwen_model.generate,
-        kwargs=dict(
-            **inputs,
-            streamer=streamer,
-            max_new_tokens=1024,
-            use_cache=True,
-            temperature=1.5,
-            min_p=0.1,
-        ),
-    )
-    thread.start()
-
-    # Stream out tokens as SSE (Server-Sent Events)
-    for tok in streamer:
-        if tok:
-            # Yield as JSON chunks encoded for SSE
-            yield f"data: {json.dumps({'chunk': tok})}\n\n"
-
-    thread.join()
+    if model_id == "lfm":
+        # LFM2.5-VL inference
+        if lfm_model is None or lfm_processor is None:
+            yield f"data: {json.dumps({'chunk': '[Error] LFM model not loaded.'})}\n\n"
+            yield "data: [DONE]\n\n"
+            return
+
+        conversation = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image", "image": image},
+                    {"type": "text", "text": full_prompt},
+                ],
+            }
+        ]
+
+        inputs = lfm_processor.apply_chat_template(
+            conversation,
+            add_generation_prompt=True,
+            return_tensors="pt",
+            return_dict=True,
+            tokenize=True,
+        ).to(lfm_model.device)
+
+        streamer = TextIteratorStreamer(
+            lfm_processor.tokenizer,
+            skip_prompt=True,
+            skip_special_tokens=True,
+            timeout=120,
+        )
+
+        thread = threading.Thread(
+            target=lfm_model.generate,
+            kwargs=dict(
+                **inputs,
+                streamer=streamer,
+                max_new_tokens=1024,
+                use_cache=True,
+            ),
+        )
+        thread.start()
+        for tok in streamer:
+            if tok:
+                yield f"data: {json.dumps({'chunk': tok})}\n\n"
+        thread.join()
+
+    else:
+        # Qwen3.5 inference
+        if qwen_model is None or qwen_processor is None:
+            yield f"data: {json.dumps({'chunk': '[Error] Qwen model not loaded.'})}\n\n"
+            yield "data: [DONE]\n\n"
+            return
+
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image", "image": image},
+                    {"type": "text", "text": full_prompt},
+                ],
+            }
+        ]
+
+        text_input = qwen_processor.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+
+        inputs = qwen_processor(
+            text=[text_input], images=[image], return_tensors="pt", padding=True
+        ).to(qwen_model.device)
+
+        streamer = TextIteratorStreamer(
+            qwen_processor.tokenizer,
+            skip_prompt=True,
+            skip_special_tokens=True,
+            timeout=120,
+        )
+
+        thread = threading.Thread(
+            target=qwen_model.generate,
+            kwargs=dict(
+                **inputs,
+                streamer=streamer,
+                max_new_tokens=1024,
+                use_cache=True,
+                temperature=1.5,
+                min_p=0.1,
+            ),
+        )
+        thread.start()
+        for tok in streamer:
+            if tok:
+                yield f"data: {json.dumps({'chunk': tok})}\n\n"
+        thread.join()
+
     yield "data: [DONE]\n\n"
 
 
@@ -129,393 +198,667 @@ async def run_inference(
     image: UploadFile = File(...),
     category: str = Form(...),
     prompt: str = Form(...),
+    model_id: str = Form("qwen"),
 ):
     try:
-        # Load and resize image
         img_bytes = await image.read()
         img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
         img.thumbnail((512, 512))
 
-        # Return streaming response wrapping our inference generator
         return StreamingResponse(
-            generate_inference_stream(img, category, prompt), 
+            generate_inference_stream(img, category, prompt, model_id),
             media_type="text/event-stream"
         )
     except Exception as e:
         return JSONResponse({"error": str(e)}, status_code=500)
 
 
-# --- Frontend UI (HTML/JS/CSS) ---
+# --- Frontend UI ---
 @app.get("/", response_class=HTMLResponse)
 async def homepage(request: Request):
-    return f"""
+    return """
 <!DOCTYPE html>
 <html lang="en">
 <head>
     <meta charset="UTF-8">
     <meta name="viewport" content="width=device-width, initial-scale=1.0">
     <title>Multimodal-Edge-Comparator</title>
-    <link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap" rel="stylesheet">
+    <link href="https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;600;700&family=Space+Mono:wght@400;700&display=swap" rel="stylesheet">
     <style>
-        :root {{
-            --bg-color: #121212;
-            --grid-color: #242424;
-            --node-bg: #1e1e1e;
-            --node-header: #2a2a2a;
-            --node-border: #3a3a3a;
-            --accent: #5e81ac;
-            --accent-hover: #81a1c1;
-            --text-main: #eceff4;
-            --text-muted: #8fbcbb;
-            --port-color: #88c0d0;
-            --wire-color: #4c566a;
-        }}
-        
-        * {{ box-sizing: border-box; font-family: 'Inter', sans-serif; }}
-        
-        body {{
+        :root {
+            --bg: #0d0d0f;
+            --grid: #1a1a1f;
+            --node-bg: #13131a;
+            --node-header: #1c1c26;
+            --node-border: #2a2a3a;
+            --accent: #7c6af7;
+            --accent2: #4ecdc4;
+            --accent3: #ff6b6b;
+            --text: #e8e8f0;
+            --muted: #6b6b8a;
+            --port: #4ecdc4;
+            --wire: #2a2a4a;
+            --wire-active: #7c6af7;
+        }
+
+        * { box-sizing: border-box; }
+        body {
             margin: 0; padding: 0; overflow: hidden;
-            background-color: var(--bg-color);
-            background-image: 
-                linear-gradient(var(--grid-color) 1px, transparent 1px),
-                linear-gradient(90deg, var(--grid-color) 1px, transparent 1px);
-            background-size: 20px 20px;
-            color: var(--text-main);
+            background: var(--bg);
+            background-image:
+                radial-gradient(circle at 20% 50%, rgba(124,106,247,0.04) 0%, transparent 50%),
+                radial-gradient(circle at 80% 20%, rgba(78,205,196,0.04) 0%, transparent 50%),
+                linear-gradient(var(--grid) 1px, transparent 1px),
+                linear-gradient(90deg, var(--grid) 1px, transparent 1px);
+            background-size: 100% 100%, 100% 100%, 24px 24px, 24px 24px;
+            color: var(--text);
+            font-family: 'JetBrains Mono', monospace;
             user-select: none;
-        }}
+        }
 
-        .top-bar {{
-            position: absolute; top: 0; left: 0; right: 0;
-            height: 50px; background: rgba(18,18,18,0.8);
+        .top-bar {
+            position: fixed; top: 0; left: 0; right: 0;
+            height: 48px; background: rgba(13,13,15,0.9);
             border-bottom: 1px solid var(--node-border);
             display: flex; align-items: center; padding: 0 20px;
-            font-weight: 700; font-size: 18px; z-index: 100;
-            backdrop-filter: blur(4px);
-            color: var(--port-color);
-            pointer-events: none;
-        }}
-
-        /* Canvas & SVG Wires */
-        #canvas {{ position: relative; width: 100vw; height: 100vh; }}
-        svg {{ position: absolute; top: 0; left: 0; width: 100%; height: 100%; pointer-events: none; z-index: 1; }}
-        path.wire {{ fill: none; stroke: var(--wire-color); stroke-width: 3; stroke-linecap: round; }}
-        path.wire.active {{ stroke: var(--accent); stroke-width: 4; animation: flow 1s linear infinite; }}
-        
-        @keyframes flow {{
-            0% {{ stroke-dashoffset: 20; }}
-            100% {{ stroke-dashoffset: 0; }}
-        }}
-
-        /* Nodes */
-        .node {{
-            position: absolute; width: 340px;
-            background: var(--node-bg); border: 1px solid var(--node-border);
-            border-radius: 8px; box-shadow: 0 8px 24px rgba(0,0,0,0.5);
-            z-index: 10; display: flex; flex-direction: column;
-        }}
-        
-        .node-header {{
-            background: var(--node-header); padding: 12px 15px;
+            gap: 12px; z-index: 1000;
+            backdrop-filter: blur(12px);
+        }
+        .top-bar .logo { font-size: 13px; font-weight: 700; color: var(--accent); letter-spacing: 0.05em; }
+        .top-bar .sep { color: var(--node-border); }
+        .top-bar .sub { font-size: 11px; color: var(--muted); }
+        .top-bar .badge {
+            margin-left: auto; background: rgba(124,106,247,0.15);
+            border: 1px solid rgba(124,106,247,0.3); padding: 3px 10px;
+            border-radius: 20px; font-size: 10px; color: var(--accent);
+        }
+
+        #canvas {
+            position: relative; width: 100vw; height: 100vh;
+            padding-top: 48px;
+            overflow: hidden;
+        }
+
+        svg.wires {
+            position: absolute; top: 0; left: 0;
+            width: 100%; height: 100%;
+            pointer-events: none; z-index: 2;
+        }
+
+        path.wire {
+            fill: none; stroke: var(--wire); stroke-width: 2.5;
+            stroke-linecap: round;
+        }
+        path.wire.active {
+            stroke: var(--wire-active); stroke-width: 3;
+            stroke-dasharray: 8 4;
+            animation: flow 0.6s linear infinite;
+        }
+        @keyframes flow { to { stroke-dashoffset: -24; } }
+
+        /* ─── Nodes ─── */
+        .node {
+            position: absolute;
+            width: 300px;
+            background: var(--node-bg);
+            border: 1px solid var(--node-border);
+            border-radius: 10px;
+            box-shadow: 0 12px 40px rgba(0,0,0,0.6), 0 0 0 0px rgba(124,106,247,0);
+            z-index: 10;
+            display: flex;
+            flex-direction: column;
+            transition: box-shadow 0.2s;
+        }
+        .node:hover {
+            box-shadow: 0 12px 40px rgba(0,0,0,0.6), 0 0 0 1px rgba(124,106,247,0.3);
+        }
+
+        /* Fixed height nodes */
+        .node.fixed-height {
+            height: 420px;
+        }
+
+        .node-header {
+            background: var(--node-header);
+            padding: 10px 14px;
             border-bottom: 1px solid var(--node-border);
-            border-radius: 8px 8px 0 0; font-weight: 600; font-size: 14px;
-            cursor: grab; display: flex; justify-content: space-between;
-        }}
-        .node-header:active {{ cursor: grabbing; }}
-        
-        .node-body {{ padding: 15px; display: flex; flex-direction: column; gap: 12px; }}
-
-        /* Ports */
-        .port {{
-            position: absolute; width: 14px; height: 14px;
-            background: var(--node-bg); border: 2px solid var(--port-color);
-            border-radius: 50%; z-index: 20; top: 50%; transform: translateY(-50%);
-        }}
-        .port.out {{ right: -8px; }}
-        .port.in {{ left: -8px; }}
-
-        /* Form Elements */
-        label {{ font-size: 12px; color: var(--text-muted); font-weight: 500; display: block; margin-bottom: 4px; }}
-        
-        input[type="file"] {{ display: none; }}
-        .file-upload {{
-            border: 2px dashed var(--node-border); border-radius: 6px;
-            padding: 20px; text-align: center; cursor: pointer;
-            font-size: 13px; color: var(--text-muted); transition: border 0.2s;
-        }}
-        .file-upload:hover {{ border-color: var(--accent); }}
-        
-        .img-preview {{ width: 100%; height: auto; max-height: 280px; object-fit: contain; border-radius: 4px; display: none; margin-top: 10px; background: #000; }}
-
-        select, textarea {{
-            width: 100%; background: #121212; border: 1px solid var(--node-border);
-            color: var(--text-main); padding: 10px; border-radius: 6px;
-            outline: none; font-size: 13px; resize: none;
-        }}
-        select:focus, textarea:focus {{ border-color: var(--accent); }}
-        
-        button {{
-            background: #ff69b4; color: #fff; border: none;
-            padding: 10px; border-radius: 6px; font-weight: 600; font-size: 14px;
-            cursor: pointer; transition: background 0.2s; display: flex; justify-content: center; gap: 8px; align-items: center;
-        }}
-        button:hover {{ background: #ff1493; }}
-        button:disabled {{ background: var(--node-border); cursor: not-allowed; color: #888; }}
-
-        .output-box {{
-            background: #121212; border: 1px solid var(--node-border);
-            border-radius: 6px; padding: 12px; min-height: 150px; max-height: 300px;
-            overflow-y: auto; font-size: 13px; line-height: 1.5; color: #d8dee9;
-            white-space: pre-wrap; user-select: text;
-        }}
-
-        .loader {{ width: 14px; height: 14px; border: 2px solid #fff; border-bottom-color: transparent; border-radius: 50%; display: inline-block; animation: rotation 1s linear infinite; display: none; }}
-        @keyframes rotation {{ 0% {{ transform: rotate(0deg); }} 100% {{ transform: rotate(360deg); }} }}
+            border-radius: 10px 10px 0 0;
+            font-size: 11px; font-weight: 700;
+            cursor: grab; display: flex; justify-content: space-between; align-items: center;
+            flex-shrink: 0;
+        }
+        .node-header:active { cursor: grabbing; }
+        .node-header .id {
+            font-size: 10px; color: var(--muted);
+            background: rgba(255,255,255,0.04);
+            padding: 2px 8px; border-radius: 4px;
+        }
+
+        .node-body {
+            padding: 14px;
+            display: flex; flex-direction: column; gap: 12px;
+            flex: 1; overflow: hidden;
+        }
+
+        /* ─── Ports ─── */
+        .port {
+            position: absolute;
+            width: 12px; height: 12px;
+            background: var(--node-bg);
+            border: 2px solid var(--port);
+            border-radius: 50%;
+            z-index: 30;
+        }
+        .port.out { right: -7px; }
+        .port.in  { left: -7px; }
+
+        /* ─── Form Elements ─── */
+        label {
+            font-size: 10px; color: var(--muted);
+            font-weight: 600; display: block; margin-bottom: 4px;
+            letter-spacing: 0.08em; text-transform: uppercase;
+        }
+
+        input[type="file"] { display: none; }
+
+        .file-upload {
+            border: 1.5px dashed var(--node-border);
+            border-radius: 8px; padding: 16px 12px;
+            text-align: center; cursor: pointer;
+            font-size: 11px; color: var(--muted);
+            transition: border-color 0.2s, background 0.2s;
+            background: rgba(255,255,255,0.01);
+        }
+        .file-upload:hover { border-color: var(--accent); background: rgba(124,106,247,0.04); }
+        .file-upload .icon { font-size: 22px; margin-bottom: 6px; }
+
+        .img-preview {
+            width: 100%; height: 190px;
+            object-fit: contain;
+            border-radius: 6px;
+            display: none;
+            background: #000;
+            border: 1px solid var(--node-border);
+        }
+
+        select, textarea {
+            width: 100%;
+            background: rgba(0,0,0,0.3);
+            border: 1px solid var(--node-border);
+            color: var(--text); padding: 9px 11px;
+            border-radius: 6px; outline: none;
+            font-size: 12px; font-family: 'JetBrains Mono', monospace;
+            resize: none; transition: border-color 0.2s;
+        }
+        select:focus, textarea:focus { border-color: var(--accent); }
+        select option { background: #1c1c26; }
+
+        button.run-btn {
+            background: linear-gradient(135deg, var(--accent), #9b59b6);
+            color: #fff; border: none;
+            padding: 10px; border-radius: 7px;
+            font-weight: 700; font-size: 12px;
+            font-family: 'JetBrains Mono', monospace;
+            cursor: pointer;
+            transition: opacity 0.2s, transform 0.1s;
+            display: flex; justify-content: center; align-items: center; gap: 8px;
+            letter-spacing: 0.04em;
+        }
+        button.run-btn:hover { opacity: 0.9; }
+        button.run-btn:active { transform: scale(0.98); }
+        button.run-btn:disabled { background: var(--node-border); cursor: not-allowed; color: #555; }
+
+        .output-box {
+            background: rgba(0,0,0,0.4);
+            border: 1px solid var(--node-border);
+            border-radius: 6px; padding: 12px;
+            flex: 1;
+            overflow-y: auto;
+            font-size: 12px; line-height: 1.6;
+            color: #c8c8e0; white-space: pre-wrap;
+            user-select: text;
+            font-family: 'JetBrains Mono', monospace;
+        }
+
+        /* Grounding canvas */
+        .ground-canvas-wrap {
+            position: relative; flex: 1;
+            border: 1px solid var(--node-border);
+            border-radius: 6px; overflow: hidden;
+            background: #000;
+            min-height: 0;
+        }
+        .ground-canvas-wrap canvas {
+            width: 100%; height: 100%;
+            object-fit: contain;
+        }
+        .ground-placeholder {
+            position: absolute; inset: 0;
+            display: flex; align-items: center; justify-content: center;
+            font-size: 11px; color: var(--muted); text-align: center;
+            padding: 12px;
+        }
+
+        .loader {
+            width: 12px; height: 12px;
+            border: 2px solid rgba(255,255,255,0.3);
+            border-top-color: #fff;
+            border-radius: 50%;
+            animation: spin 0.7s linear infinite;
+            display: none;
+        }
+        @keyframes spin { to { transform: rotate(360deg); } }
+
+        /* Status dot */
+        .status-dot {
+            width: 7px; height: 7px; border-radius: 50%;
+            background: var(--muted); display: inline-block;
+            margin-right: 6px;
+        }
+        .status-dot.active { background: var(--accent2); box-shadow: 0 0 6px var(--accent2); }
+
+        /* Model badge */
+        .model-badge {
+            display: inline-block; padding: 2px 8px;
+            border-radius: 4px; font-size: 9px; font-weight: 700;
+            letter-spacing: 0.06em; text-transform: uppercase;
+        }
+        .model-badge.qwen { background: rgba(124,106,247,0.2); color: var(--accent); border: 1px solid rgba(124,106,247,0.3); }
+        .model-badge.lfm  { background: rgba(78,205,196,0.15); color: var(--accent2); border: 1px solid rgba(78,205,196,0.3); }
     </style>
 </head>
 <body>
 
-    <div class="top-bar">
-        Multimodal Edge Comparator [Qwen3.5-2B]
+<div class="top-bar">
+    <span class="logo">MULTIMODAL EDGE</span>
+    <span class="sep">|</span>
+    <span class="sub">Node-Based Inference Canvas</span>
+    <span class="badge">v2.0 — DUAL MODEL</span>
+</div>
+
+<div id="canvas">
+    <svg class="wires">
+        <!-- Left col → Task node -->
+        <path id="wire-img-task"   class="wire" />
+        <path id="wire-model-task" class="wire" />
+        <!-- Task → Output -->
+        <path id="wire-task-out"   class="wire" />
+        <!-- Task → Ground -->
+        <path id="wire-task-gnd"   class="wire" />
+    </svg>
+
+    <!-- ─── ID 01 : Image Input (left col, top) ─── -->
+    <div class="node fixed-height" id="node-img" style="left:60px; top:68px;">
+        <div class="node-header">
+            <span><span class="status-dot" id="dot-img"></span>Input Image</span>
+            <span class="id">ID: 01</span>
+        </div>
+        <div class="node-body">
+            <div>
+                <label>Upload Image</label>
+                <div class="file-upload" id="dropZone">
+                    <div class="icon">🖼️</div>
+                    Click or drop image here
+                    <input type="file" id="fileInput" accept="image/*">
+                </div>
+                <img id="imgPreview" class="img-preview" />
+            </div>
+        </div>
+        <!-- OUT port: vertically centered on right edge -->
+        <div class="port out" id="port-img-out" style="top:50%; transform:translateY(-50%);"></div>
     </div>
 
-    <div id="canvas">
-        <svg>
-            <path id="wire1" class="wire" />
-            <path id="wire2" class="wire" />
-        </svg>
-
-        <!-- Node 1: Image Input -->
-        <div class="node" id="node-img" style="left: 100px; top: 150px;">
-            <div class="node-header"><span>Input Node</span><span>ID: 01</span></div>
-            <div class="node-body">
-                <div>
-                    <label>Upload Image</label>
-                    <div class="file-upload" id="dropZone">
-                        Click or Drop Image Here
-                        <input type="file" id="fileInput" accept="image/*">
-                    </div>
-                    <img id="imgPreview" class="img-preview" />
-                </div>
+    <!-- ─── ID 02 : Model Selector (left col, bottom) ─── -->
+    <div class="node fixed-height" id="node-model" style="left:60px; top:508px;">
+        <div class="node-header">
+            <span><span class="status-dot" id="dot-model"></span>Model Selector</span>
+            <span class="id">ID: 02</span>
+        </div>
+        <div class="node-body">
+            <div>
+                <label>Active Model</label>
+                <select id="modelSelect">
+                    <option value="qwen">Qwen3.5-2B (Vision-Language)</option>
+                    <option value="lfm">LFM2.5-VL-450M (LiquidAI)</option>
+                </select>
+            </div>
+            <div id="modelInfoBox" style="background:rgba(124,106,247,0.07); border:1px solid rgba(124,106,247,0.2); border-radius:6px; padding:10px; font-size:10px; color:var(--muted); line-height:1.6;">
+                <span class="model-badge qwen">QWEN</span>
+                <br><br>
+                Qwen3.5 2B parameter multimodal model by Alibaba Cloud. Supports Query, Caption, Point & Detect tasks with streaming output.
             </div>
-            <div class="port out" id="port-img-out"></div>
+            <div style="flex:1;"></div>
         </div>
+        <!-- OUT port: vertically centered on right edge -->
+        <div class="port out" id="port-model-out" style="top:50%; transform:translateY(-50%);"></div>
+    </div>
 
-        <!-- Node 2: Task Execution -->
-        <div class="node" id="node-task" style="left: 550px; top: 150px;">
-            <div class="port in" id="port-task-in" style="top: 30%;"></div>
-            <div class="node-header"><span>Task Node</span><span>ID: 02</span></div>
-            <div class="node-body">
-                <div>
-                    <label>Task Category</label>
-                    <select id="categorySelect">
-                        <option value="Query" selected>Query</option>
-                        <option value="Caption">Caption</option>
-                        <option value="Point">Point</option>
-                        <option value="Detect">Detect</option>
-                    </select>
-                </div>
-                <div>
-                    <label>Prompt Directive</label>
-                    <textarea id="promptInput" rows="4" placeholder="e.g., Count the total number of boats and describe the environment."></textarea>
-                </div>
-                <button id="runBtn">
-                    <span>Quick Execution</span>
-                    <span class="loader" id="btnLoader"></span>
-                </button>
+    <!-- ─── ID 03 : Task Node (right col, row 1) ─── -->
+    <div class="node fixed-height" id="node-task" style="left:460px; top:68px;">
+        <!-- IN port aligned to accept wires from ID01 and ID02 -->
+        <div class="port in" id="port-task-in" style="top:50%; transform:translateY(-50%);"></div>
+        <div class="node-header">
+            <span><span class="status-dot" id="dot-task"></span>Task Config</span>
+            <span class="id">ID: 03</span>
+        </div>
+        <div class="node-body">
+            <div>
+                <label>Task Category</label>
+                <select id="categorySelect">
+                    <option value="Query">Query</option>
+                    <option value="Caption">Caption</option>
+                    <option value="Point">Point</option>
+                    <option value="Detect">Detect</option>
+                </select>
+            </div>
+            <div>
+                <label>Prompt Directive</label>
+                <textarea id="promptInput" rows="5" placeholder="e.g., Count the total number of boats and describe the environment."></textarea>
             </div>
-            <div class="port out" id="port-task-out" style="top: 70%;"></div>
+            <button class="run-btn" id="runBtn">
+                <span>Execute</span>
+                <span class="loader" id="btnLoader"></span>
+            </button>
         </div>
+        <!-- OUT port -->
+        <div class="port out" id="port-task-out" style="top:50%; transform:translateY(-50%);"></div>
+    </div>
 
-        <!-- Node 3: Output Display -->
-        <div class="node" id="node-out" style="left: 1000px; top: 150px; width: 380px;">
-            <div class="port in" id="port-out-in"></div>
-            <div class="node-header"><span>Output Node</span><span>ID: 03</span></div>
-            <div class="node-body">
-                <label>Streamed Result</label>
-                <div class="output-box" id="outputBox">Results will stream here...</div>
+    <!-- ─── ID 04 : Output Node (right col, row 1, further right) ─── -->
+    <div class="node fixed-height" id="node-out" style="left:860px; top:68px;">
+        <div class="port in" id="port-out-in" style="top:50%; transform:translateY(-50%);"></div>
+        <div class="node-header">
+            <span><span class="status-dot" id="dot-out"></span>Output Stream</span>
+            <span class="id">ID: 04</span>
+        </div>
+        <div class="node-body">
+            <label>Streamed Result</label>
+            <div class="output-box" id="outputBox">Results will stream here...</div>
+        </div>
+    </div>
+
+    <!-- ─── ID 05 : Grounding Visualiser (right col, row 2) ─── -->
+    <div class="node fixed-height" id="node-gnd" style="left:860px; top:508px;">
+        <div class="port in" id="port-gnd-in" style="top:50%; transform:translateY(-50%);"></div>
+        <div class="node-header">
+            <span><span class="status-dot" id="dot-gnd"></span>View Grounding</span>
+            <span class="id">ID: 05</span>
+        </div>
+        <div class="node-body">
+            <label>Point / Detect Overlay</label>
+            <div class="ground-canvas-wrap">
+                <canvas id="groundCanvas"></canvas>
+                <div class="ground-placeholder" id="groundPlaceholder">
+                    Active for Point / Detect tasks.<br>Run inference to visualise.
+                </div>
             </div>
         </div>
     </div>
+</div>
+
+<script>
+// ══════════════════════════════════════════════
+//  WIRE DRAWING
+// ══════════════════════════════════════════════
+function portCenter(id) {
+    const el = document.getElementById(id);
+    if (!el) return {x:0,y:0};
+    const r = el.getBoundingClientRect();
+    return { x: r.left + r.width/2, y: r.top + r.height/2 };
+}
+
+function bezier(p1, p2) {
+    const dx = Math.abs(p2.x - p1.x) * 0.55;
+    return `M ${p1.x} ${p1.y} C ${p1.x+dx} ${p1.y}, ${p2.x-dx} ${p2.y}, ${p2.x} ${p2.y}`;
+}
+
+function updateWires() {
+    const wires = [
+        ['wire-img-task',   'port-img-out',   'port-task-in'],
+        ['wire-model-task', 'port-model-out',  'port-task-in'],
+        ['wire-task-out',   'port-task-out',   'port-out-in'],
+        ['wire-task-gnd',   'port-task-out',   'port-gnd-in'],
+    ];
+    for (const [id, from, to] of wires) {
+        const el = document.getElementById(id);
+        if (el) el.setAttribute('d', bezier(portCenter(from), portCenter(to)));
+    }
+}
+
+// ══════════════════════════════════════════════
+//  DRAGGING
+// ══════════════════════════════════════════════
+document.querySelectorAll('.node').forEach(node => {
+    const header = node.querySelector('.node-header');
+    let drag = false, sx, sy, il, it;
+
+    header.addEventListener('mousedown', e => {
+        drag = true; sx = e.clientX; sy = e.clientY;
+        il = parseInt(node.style.left)||0;
+        it = parseInt(node.style.top)||0;
+        node.style.zIndex = 100;
+        e.preventDefault();
+    });
+    document.addEventListener('mousemove', e => {
+        if (!drag) return;
+        node.style.left = `${il + e.clientX - sx}px`;
+        node.style.top  = `${it + e.clientY - sy}px`;
+        updateWires();
+    });
+    document.addEventListener('mouseup', () => {
+        if (drag) { drag = false; node.style.zIndex = 10; }
+    });
+});
+
+window.addEventListener('resize', updateWires);
+// Delay first draw so layout is complete
+requestAnimationFrame(() => { updateWires(); });
+
+// ══════════════════════════════════════════════
+//  FILE UPLOAD
+// ══════════════════════════════════════════════
+let currentFile = null;
+const dropZone   = document.getElementById('dropZone');
+const fileInput  = document.getElementById('fileInput');
+const imgPreview = document.getElementById('imgPreview');
+const dotImg     = document.getElementById('dot-img');
+
+function handleFile(file) {
+    if (!file || !file.type.startsWith('image/')) return;
+    currentFile = file;
+    imgPreview.src = URL.createObjectURL(file);
+    imgPreview.style.display = 'block';
+    dropZone.style.display   = 'none';
+    dotImg.classList.add('active');
+    // Wait for layout then redraw wires
+    requestAnimationFrame(updateWires);
+}
+
+dropZone.onclick  = () => fileInput.click();
+fileInput.onchange = e => handleFile(e.target.files[0]);
+dropZone.ondragover  = e => { e.preventDefault(); dropZone.style.borderColor = 'var(--accent)'; };
+dropZone.ondragleave = ()=> { dropZone.style.borderColor = ''; };
+dropZone.ondrop      = e => {
+    e.preventDefault(); dropZone.style.borderColor = '';
+    if (e.dataTransfer.files.length) handleFile(e.dataTransfer.files[0]);
+};
+
+// ══════════════════════════════════════════════
+//  MODEL SELECTOR INFO
+// ══════════════════════════════════════════════
+const modelSelect  = document.getElementById('modelSelect');
+const modelInfoBox = document.getElementById('modelInfoBox');
+const dotModel     = document.getElementById('dot-model');
+dotModel.classList.add('active');
+
+const MODEL_INFO = {
+    qwen: `<span class="model-badge qwen">QWEN</span><br><br>Qwen3.5 2B parameter multimodal model by Alibaba Cloud. Supports Query, Caption, Point & Detect tasks with streaming output.`,
+    lfm:  `<span class="model-badge lfm">LFM</span><br><br>LFM2.5-VL 450M parameter vision-language model by LiquidAI. Ultra-lightweight edge model with strong grounding capabilities.`,
+};
+
+modelSelect.onchange = () => {
+    modelInfoBox.innerHTML = MODEL_INFO[modelSelect.value] || '';
+};
+
+// ══════════════════════════════════════════════
+//  CATEGORY → PLACEHOLDER
+// ══════════════════════════════════════════════
+const categorySelect = document.getElementById('categorySelect');
+const promptInput    = document.getElementById('promptInput');
+const PLACEHOLDERS   = {
+    Query:   'e.g., Count the total number of boats and describe the environment.',
+    Caption: 'e.g., short | normal | detailed',
+    Point:   'e.g., The gun held by the person.',
+    Detect:  'e.g., The headlight of the car.',
+};
+categorySelect.onchange = e => {
+    promptInput.placeholder = PLACEHOLDERS[e.target.value] || '';
+};
+
+// ══════════════════════════════════════════════
+//  GROUNDING VISUALIZER
+// ══════════════════════════════════════════════
+const groundCanvas      = document.getElementById('groundCanvas');
+const groundPlaceholder = document.getElementById('groundPlaceholder');
+const gCtx              = groundCanvas.getContext('2d');
+
+function drawGrounding(imgSrc, jsonText) {
+    const parsed = safeParseJSON(jsonText);
+    if (!parsed || (Array.isArray(parsed) && parsed.length === 0)) return;
+
+    const img = new Image();
+    img.onload = () => {
+        groundCanvas.width  = img.naturalWidth;
+        groundCanvas.height = img.naturalHeight;
+        gCtx.drawImage(img, 0, 0);
+        groundPlaceholder.style.display = 'none';
+
+        const W = img.naturalWidth, H = img.naturalHeight;
+        gCtx.strokeStyle = '#4ecdc4';
+        gCtx.lineWidth   = Math.max(2, W/200);
+        gCtx.fillStyle   = 'rgba(78,205,196,0.25)';
+        gCtx.font        = `bold ${Math.max(12, W/40)}px JetBrains Mono, monospace`;
+
+        const items = Array.isArray(parsed) ? parsed : [parsed];
+        items.forEach((item, i) => {
+            // Point format: [x, y] or {x, y}
+            if (Array.isArray(item) && item.length === 2 && typeof item[0] === 'number') {
+                const [x, y] = item;
+                const px = x * W, py = y * H;
+                gCtx.beginPath();
+                gCtx.arc(px, py, Math.max(8, W/60), 0, Math.PI*2);
+                gCtx.fill();
+                gCtx.stroke();
+            }
+            // BBox format: [x1,y1,x2,y2] or {x1,y1,x2,y2}
+            if (Array.isArray(item) && item.length === 4) {
+                const [x1,y1,x2,y2] = item;
+                const bx = x1*W, by = y1*H, bw = (x2-x1)*W, bh = (y2-y1)*H;
+                gCtx.fillRect(bx, by, bw, bh);
+                gCtx.strokeRect(bx, by, bw, bh);
+                gCtx.fillStyle = '#4ecdc4';
+                gCtx.fillText(`${i+1}`, bx+4, by+gCtx.font.match(/\d+/)[0]*1.1);
+                gCtx.fillStyle = 'rgba(78,205,196,0.25)';
+            }
+        });
+    };
+    img.src = imgSrc;
+}
+
+function safeParseJSON(text) {
+    text = text.trim().replace(/^```(json)?/, '').replace(/```$/, '').trim();
+    try { return JSON.parse(text); } catch(_){}
+    try { return eval('(' + text + ')'); } catch(_) { return null; }
+}
+
+// ══════════════════════════════════════════════
+//  RUN INFERENCE
+// ══════════════════════════════════════════════
+const runBtn    = document.getElementById('runBtn');
+const btnLoader = document.getElementById('btnLoader');
+const outputBox = document.getElementById('outputBox');
+const allWires  = ['wire-img-task','wire-model-task','wire-task-out','wire-task-gnd'];
+const dotTask   = document.getElementById('dot-task');
+const dotOut    = document.getElementById('dot-out');
+const dotGnd    = document.getElementById('dot-gnd');
+
+runBtn.onclick = async () => {
+    if (!currentFile) { alert('Please upload an image into the Input Node.'); return; }
+    const promptStr = promptInput.value.trim();
+    if (!promptStr) { alert('Please enter a prompt directive.'); return; }
+
+    // UI: running state
+    runBtn.disabled = true;
+    btnLoader.style.display = 'inline-block';
+    outputBox.innerText = '';
+    outputBox.style.color = '';
+    groundPlaceholder.style.display = 'flex';
+    gCtx.clearRect(0, 0, groundCanvas.width, groundCanvas.height);
+    dotTask.classList.add('active');
+    dotOut.classList.remove('active');
+    dotGnd.classList.remove('active');
+    allWires.forEach(id => document.getElementById(id)?.classList.add('active'));
+
+    const formData = new FormData();
+    formData.append('image',    currentFile);
+    formData.append('category', categorySelect.value);
+    formData.append('prompt',   promptStr);
+    formData.append('model_id', modelSelect.value);
+
+    let fullText = '';
+
+    try {
+        const response = await fetch('/api/run', { method: 'POST', body: formData });
+        if (!response.ok) {
+            const err = await response.json();
+            throw new Error(err.error || 'Execution failed.');
+        }
+
+        const reader  = response.body.getReader();
+        const decoder = new TextDecoder('utf-8');
+        let   buffer  = '';
+
+        while (true) {
+            const { value, done } = await reader.read();
+            if (done) break;
+            buffer += decoder.decode(value, { stream: true });
+            const lines = buffer.split('\\n\\n');
+            buffer = lines.pop();
+            for (const line of lines) {
+                if (!line.startsWith('data: ')) continue;
+                const payload = line.replace('data: ', '');
+                if (payload === '[DONE]') break;
+                try {
+                    const data = JSON.parse(payload);
+                    if (data.chunk) {
+                        fullText += data.chunk;
+                        outputBox.innerText = fullText;
+                        outputBox.scrollTop = outputBox.scrollHeight;
+                    }
+                } catch(_) {}
+            }
+        }
+
+        dotOut.classList.add('active');
+
+        // Visualise grounding if applicable
+        const cat = categorySelect.value;
+        if ((cat === 'Point' || cat === 'Detect') && fullText.trim()) {
+            dotGnd.classList.add('active');
+            drawGrounding(URL.createObjectURL(currentFile), fullText);
+        }
 
-    <script>
-        // --- 1. Node Canvas & Dragging Logic ---
-        function getCenter(elementId) {{
-            const el = document.getElementById(elementId);
-            const rect = el.getBoundingClientRect();
-            return {{ x: rect.left + rect.width / 2, y: rect.top + rect.height / 2 }};
-        }}
-
-        function drawCurve(p1, p2) {{
-            const dx = Math.abs(p2.x - p1.x) * 0.5;
-            return `M ${{p1.x}} ${{p1.y}} C ${{p1.x + dx}} ${{p1.y}}, ${{p2.x - dx}} ${{p2.y}}, ${{p2.x}} ${{p2.y}}`;
-        }}
-
-        function updateWires() {{
-            const pImgOut = getCenter('port-img-out');
-            const pTaskIn = getCenter('port-task-in');
-            const pTaskOut = getCenter('port-task-out');
-            const pOutIn = getCenter('port-out-in');
-
-            document.getElementById('wire1').setAttribute('d', drawCurve(pImgOut, pTaskIn));
-            document.getElementById('wire2').setAttribute('d', drawCurve(pTaskOut, pOutIn));
-        }}
-
-        // Initialize node dragging
-        document.querySelectorAll('.node').forEach(node => {{
-            const header = node.querySelector('.node-header');
-            let isDragging = false, startX, startY, initialLeft, initialTop;
-
-            header.addEventListener('mousedown', (e) => {{
-                isDragging = true;
-                startX = e.clientX; startY = e.clientY;
-                initialLeft = parseInt(node.style.left || 0);
-                initialTop = parseInt(node.style.top || 0);
-                node.style.zIndex = 100; // bring to front
-            }});
-
-            document.addEventListener('mousemove', (e) => {{
-                if (!isDragging) return;
-                const dx = e.clientX - startX;
-                const dy = e.clientY - startY;
-                node.style.left = `${{initialLeft + dx}}px`;
-                node.style.top = `${{initialTop + dy}}px`;
-                updateWires();
-            }});
-
-            document.addEventListener('mouseup', () => {{
-                if (isDragging) {{
-                    isDragging = false;
-                    node.style.zIndex = 10;
-                }}
-            }});
-        }});
-
-        window.addEventListener('resize', updateWires);
-        updateWires(); // initial draw
-
-
-        // --- 2. Application Logic ---
-        let currentFile = null;
-
-        // File Upload Handlers
-        const dropZone = document.getElementById('dropZone');
-        const fileInput = document.getElementById('fileInput');
-        const imgPreview = document.getElementById('imgPreview');
-
-        function handleFile(file) {{
-            if (file && file.type.startsWith('image/')) {{
-                currentFile = file;
-                imgPreview.src = URL.createObjectURL(file);
-                imgPreview.style.display = 'block';
-                dropZone.style.display = 'none';
-                updateWires(); // redraw wires because node size changed
-            }}
-        }}
-
-        dropZone.onclick = () => fileInput.click();
-        fileInput.onchange = e => handleFile(e.target.files[0]);
-        dropZone.ondragover = e => {{ e.preventDefault(); dropZone.style.borderColor = "var(--accent)"; }};
-        dropZone.ondragleave = e => {{ dropZone.style.borderColor = "var(--node-border)"; }};
-        dropZone.ondrop = e => {{
-            e.preventDefault();
-            dropZone.style.borderColor = "var(--node-border)";
-            if(e.dataTransfer.files.length) handleFile(e.dataTransfer.files[0]);
-        }};
-
-        // Category Selection Updates
-        const categorySelect = document.getElementById('categorySelect');
-        const promptInput = document.getElementById('promptInput');
-        const placeholders = {{
-            "Query": "e.g., Count the total number of boats and describe the environment.",
-            "Caption": "e.g., short, normal, detailed",
-            "Point": "e.g., The gun held by the person.",
-            "Detect": "e.g., The headlight of the car."
-        }};
-
-        categorySelect.onchange = (e) => {{
-            const val = e.target.value;
-            promptInput.placeholder = placeholders[val] || "Enter your prompt here.";
-        }};
-
-        // Execution Logic
-        const runBtn = document.getElementById('runBtn');
-        const btnLoader = document.getElementById('btnLoader');
-        const outputBox = document.getElementById('outputBox');
-        const wire1 = document.getElementById('wire1');
-        const wire2 = document.getElementById('wire2');
-
-        runBtn.onclick = async () => {{
-            if (!currentFile) return alert("Please upload an image into the Input Node.");
-            const promptStr = promptInput.value.trim();
-            if (!promptStr) return alert("Please enter a prompt directive in the Task Node.");
-
-            // UI State updates
-            runBtn.disabled = true;
-            btnLoader.style.display = 'inline-block';
-            outputBox.innerText = ""; // clear
-            wire1.classList.add('active');
-            wire2.classList.add('active');
-
-            const formData = new FormData();
-            formData.append("image", currentFile);
-            formData.append("category", categorySelect.value);
-            formData.append("prompt", promptStr);
-
-            try {{
-                const response = await fetch('/api/run', {{
-                    method: 'POST',
-                    body: formData
-                }});
-
-                if (!response.ok) {{
-                    const err = await response.json();
-                    throw new Error(err.error || "Execution failed.");
-                }}
-
-                // Streaming Reader setup
-                const reader = response.body.getReader();
-                const decoder = new TextDecoder('utf-8');
-                let buffer = '';
-
-                while (true) {{
-                    const {{ value, done }} = await reader.read();
-                    if (done) break;
-                    
-                    buffer += decoder.decode(value, {{ stream: true }});
-                    
-                    // Parse SSE chunks
-                    const lines = buffer.split('\\n\\n');
-                    buffer = lines.pop(); // keep last incomplete chunk
-                    
-                    for (const line of lines) {{
-                        if (line.startsWith('data: ')) {{
-                            const payload = line.replace('data: ', '');
-                            if (payload === '[DONE]') break;
-                            
-                            try {{
-                                const data = JSON.parse(payload);
-                                if (data.chunk) {{
-                                    outputBox.innerText += data.chunk;
-                                    outputBox.scrollTop = outputBox.scrollHeight; // auto-scroll
-                                }}
-                            }} catch(e) {{
-                                console.error("Chunk parsing error:", e);
-                            }}
-                        }}
-                    }}
-                }}
-            }} catch (error) {{
-                outputBox.innerText = `[Error] ${{error.message}}`;
-                outputBox.style.color = "#bf616a"; // Red error tint
-            }} finally {{
-                // Revert UI State
-                runBtn.disabled = false;
-                btnLoader.style.display = 'none';
-                wire1.classList.remove('active');
-                wire2.classList.remove('active');
-            }}
-        }};
-    </script>
+    } catch (err) {
+        outputBox.innerText = `[Error] ${err.message}`;
+        outputBox.style.color = '#ff6b6b';
+    } finally {
+        runBtn.disabled = false;
+        btnLoader.style.display = 'none';
+        dotTask.classList.remove('active');
+        allWires.forEach(id => document.getElementById(id)?.classList.remove('active'));
+    }
+};
+</script>
 </body>
 </html>
 """