Add noise

models/ScoreFunction_HDC.py

@@ -4,6 +4,42 @@ import copy
 import torch.nn.functional as F
 from torch.nn import Parameter
 
+def _fake_quant_sym(x: torch.Tensor, bits: int, eps: float = 1e-8) -> torch.Tensor:
+    bits = int(bits)
+    if bits >= 32:
+        return x
+    if bits == 16:
+        # pick fp16; if you prefer bf16: x.to(torch.bfloat16).to(torch.float32)
+        return x.to(torch.float16).to(torch.float32)
+    if bits == 1:
+        return torch.sign(x)
+
+    # signed symmetric levels: [-Qmax, Qmax]
+    Qmax = (1 << (bits - 1)) - 1
+    # per-row scale (last dim); works for both (N,d) and (...,d)
+    max_abs = x.abs().amax(dim=-1, keepdim=True).clamp(min=eps)
+    scale = max_abs / Qmax
+    q = torch.round(x / scale).clamp(-Qmax, Qmax)
+    return (q * scale).to(x.dtype)
+
+def qlinear(x: torch.Tensor, layer: torch.nn.Linear, bits: int) -> torch.Tensor:
+    """Quantize BOTH activation and weight, then do linear in float."""
+    if int(bits) >= 32:
+        return layer(x)
+    if int(bits) == 16:
+        # do true fp16 compute-ish (still uses PyTorch kernels)
+        x16 = x.to(torch.float16)
+        w16 = layer.weight.to(torch.float16)
+        b16 = None if layer.bias is None else layer.bias.to(torch.float16)
+        y16 = F.linear(x16, w16, b16)
+        return y16.to(torch.float32)
+
+    xq = _fake_quant_sym(x, bits)
+    wq = _fake_quant_sym(layer.weight, bits)
+    b = layer.bias  # keep bias float (common & stable)
+    y = F.linear(xq, wq, b)
+    return _fake_quant_sym(y, bits)
+
 class HDReason(torch.nn.Module):
     def __init__(self, d=10, D=256):
         super().__init__()
@@ -21,6 +57,7 @@ class HDReason(torch.nn.Module):
         self.activation0 = torch.nn.ReLU()
         self.activation1 = torch.nn.ReLU()
 
+    """
     def forward(self, x):
         #NOTE: build adjacency graph
         q = self.activation1(self.HDC_encoder(self.activation0(self.q_proj(x))))
@@ -34,7 +71,48 @@ class HDReason(torch.nn.Module):
         out = adj @ v
         out = out*0.3 + 0.7*self.HDC_encoder(self.activation0(self.Linear(x)))
         return out
+    """
+    def forward(self, x, quant_bits: int = 32):
+        b = int(quant_bits)
+
+        # q path
+        q = qlinear(x, self.q_proj, b)
+        q = self.activation0(q)
+        q = qlinear(q, self.HDC_encoder, b)
+        q = self.activation1(q)
+
+        # k path
+        k = qlinear(x, self.k_proj, b)
+        k = self.activation0(k)
+        k = qlinear(k, self.HDC_encoder, b)
+        k = self.activation1(k)
+
+        q = _fake_quant_sym(q * self.scale, b)
+        k = _fake_quant_sym(k, b)
+
+        # adj matmul + softmax
+        adj = _fake_quant_sym(q @ k.transpose(-2, -1), b)
+
+        # softmax is sensitive at low-bit; keep it in fp32 but quantize output
+        adj = adj.softmax(dim=-1)
+        adj = _fake_quant_sym(adj, b)
+
+        # v path
+        v = qlinear(x, self.v_proj, b)
+        v = self.activation0(v)
+        v = qlinear(v, self.HDC_encoder, b)
+        v = self.activation1(v)
+        v = _fake_quant_sym(v, b)
+
+        out = _fake_quant_sym(adj @ v, b)
 
+        # skip/mix branch
+        base = qlinear(x, self.Linear, b)
+        base = self.activation0(base)
+        base = qlinear(base, self.HDC_encoder, b)
+
+        out = _fake_quant_sym(out * 0.3 + 0.7 * base, b)
+        return out
 
 class ScoreFunctionHDC(torch.nn.Module):
     def __init__(self, N_words=20, HDV_D=512) -> None:
@@ -50,6 +128,7 @@ class ScoreFunctionHDC(torch.nn.Module):
         self.Activation2 = torch.nn.Sigmoid()
         self.register_parameter('bias',Parameter(torch.zeros(1)))
 
+    """
     def forward(self, x):
         #NOTE: input has shape NxN_word
         #NOTE: N_bbox x N_word 
@@ -62,4 +141,31 @@ class ScoreFunctionHDC(torch.nn.Module):
         output = self.Activation1(output)
         output = self.Linear4(output) + self.bias
         output = self.Activation2(output)
+        return output
+    """
+    def forward(self, x, quant_bits: int = 32):
+        b = int(quant_bits)
+
+        # input activation quant (optional but consistent)
+        if b < 32:
+            x = _fake_quant_sym(x, b)
+
+        output = self.HDReason(x, quant_bits=b)
+        output = self.norm(output)  # LayerNorm usually best left fp32
+        output = self.Activation1(output)
+        if b < 16:
+            output = _fake_quant_sym(output, b)
+
+        output = qlinear(output, self.Linear2, b)
+        output = self.Activation1(output)
+        if b < 16:
+            output = _fake_quant_sym(output, b)
+
+        output = qlinear(output, self.Linear3, b)
+        output = self.Activation1(output)
+        if b < 16:
+            output = _fake_quant_sym(output, b)
+
+        output = qlinear(output, self.Linear4, b) + self.bias
+        output = self.Activation2(output)
         return output

models/TaskCLIP.py

@@ -55,13 +55,52 @@ class TaskCLIP(torch.nn.Module):
         self.glob_adapter = torch.nn.MultiheadAttention(self.d_model, 
                                                         self.nhead, 
                                                         dropout=self.dropout)
+        self.score_function_name = model_config["score_function"]
         if model_config['score_function'] != 'HDC':
             self.ScoreFunction = ScoreFunction(N_words=self.N_words)
         else:
             self.ScoreFunction = ScoreFunctionHDC(N_words=self.N_words, HDV_D=int(model_config['HDV_D']))
         self.threshold = 0.1
     
-    def forward(self, tgt, memory, image_embedding,norm=False):
+    def _apply_hw_noise(self, score_raw: torch.Tensor, dist: str, width_0_100: int, strength_0_100: int) -> torch.Tensor:
+        dist = (dist or "none").lower()
+        w = max(0, min(100, int(width_0_100)))
+        s = max(0, min(100, int(strength_0_100)))
+
+        if dist == "none" or w == 0 or s == 0:
+            return score_raw
+
+        # Tune this constant to match your desired “device noise” magnitude.
+        # score_raw here is a dot-product similarity matrix; typical scale depends on your embeddings.
+        MAX_WIDTH = 5.0
+
+        base = (w / 100.0) * MAX_WIDTH
+        scale = (s / 100.0)
+        eps = base * scale
+
+        if dist == "gaussian":
+            noise = torch.randn_like(score_raw) * eps
+        elif dist == "uniform":
+            noise = (torch.rand_like(score_raw) * 2.0 - 1.0) * eps
+        elif dist == "laplace":
+            # Laplace(0, b): sample via inverse-CDF
+            u = torch.rand_like(score_raw) - 0.5
+            noise = -eps * torch.sign(u) * torch.log1p(-2.0 * torch.abs(u))
+        else:
+            return score_raw
+
+        return score_raw + noise
+
+    def forward(
+        self, 
+        tgt, 
+        memory, 
+        image_embedding,
+        norm=False,
+        hw_noise_dist: str = "none",
+        hw_noise_width: int = 0,
+        hw_noise_strength: int = 0,
+        hdc_bits: int = 32):
         if self.norm_before:
             tgt /= tgt.norm(dim=-1, keepdim=True)
             memory /= memory.norm(dim=-1, keepdim=True)
@@ -74,9 +113,14 @@ class TaskCLIP(torch.nn.Module):
         tgt = self.ratio_glob*self.glob_adapter(tgt, image_embedding_temp, image_embedding_temp)[0] + (1 - self.ratio_glob)*tgt
         tgt_new, memory_new = self.decoder(tgt,memory,None)
         score_raw = torch.mm(tgt_new,memory_new.T)
+        # add noise
+        score_raw = self._apply_hw_noise(score_raw, hw_noise_dist, hw_noise_width, hw_noise_strength)
         if self.norm_after:
             score_raw = self.Norm(score_raw)
-        score_res = self.ScoreFunction(score_raw)
+        if self.score_function_name == 'HDC':
+            score_res = self.ScoreFunction(score_raw, quant_bits=hdc_bits)
+        else:
+            score_res = self.ScoreFunction(score_raw)
         return tgt_new, memory_new, score_res, score_raw
 
     def Norm(self, score):

webui/app.py

@@ -1,8 +1,13 @@
 import os
 import uuid
+import io
+import traceback
 from pathlib import Path
 
-import traceback
+import numpy as np
+import torch
+from PIL import Image, ImageFilter
+
 from fastapi import FastAPI, Request, UploadFile, File, Form
 from fastapi.responses import HTMLResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
@@ -59,13 +64,105 @@ DEFAULT_OD = "xlarge"
 DEFAULT_SAM_CKPT = str(CKPT_DIR / "sam2.1_l.pt")
 DEFAULT_IMAGEBIND_CKPT = str(CKPT_DIR / "imagebind_huge.pth")  # optional but recommended
 
+
+def _clamp_int(x, lo=0, hi=100) -> int:
+    try:
+        v = int(x)
+    except Exception:
+        v = 0
+    return max(lo, min(hi, v))
+
+
+def apply_noise_pil(img: Image.Image, noise_type: str, strength_0_100: int) -> Image.Image:
+    """
+    Simple input-noise layer applied before running YOLO/TaskCLIP.
+    strength_0_100: 0..100
+    """
+    strength = _clamp_int(strength_0_100, 0, 100)
+    t = (noise_type or "none").lower()
+
+    if strength == 0 or t in ["none", "default", "off"]:
+        return img
+
+    arr = np.asarray(img).astype(np.float32)
+
+    if t == "gaussian":
+        # sigma in [0, 25] roughly
+        sigma = (strength / 100.0) * 25.0
+        noise = np.random.normal(0.0, sigma, size=arr.shape).astype(np.float32)
+        out = np.clip(arr + noise, 0, 255).astype(np.uint8)
+        return Image.fromarray(out)
+
+    if t == "linear":
+        # simple brightness/contrast-like linear shift
+        alpha = 1.0 + (strength / 100.0) * 0.6  # 1.0 -> 1.6
+        beta = (strength / 100.0) * 20.0        # 0 -> 20
+        out = np.clip(arr * alpha + beta, 0, 255).astype(np.uint8)
+        return Image.fromarray(out)
+
+    # adversarial-ish synthetic corruptions (fast, deterministic-ish)
+    if t in ["adv", "adv_rand_sign"]:
+        amp = (strength / 100.0) * 18.0
+        sign = np.random.choice([-1.0, 1.0], size=arr.shape).astype(np.float32)
+        out = np.clip(arr + sign * amp, 0, 255).astype(np.uint8)
+        return Image.fromarray(out)
+
+    if t == "adv_edge_sign":
+        # edge sign from Laplacian filter, then apply sign perturbation
+        gray = img.convert("L").filter(ImageFilter.FIND_EDGES)
+        g = np.asarray(gray).astype(np.float32) / 255.0
+        sign2d = np.where(g > 0.2, 1.0, -1.0).astype(np.float32)  # crude edge mask
+        amp = (strength / 100.0) * 18.0
+        sign = np.repeat(sign2d[..., None], 3, axis=2)
+        out = np.clip(arr + sign * amp, 0, 255).astype(np.uint8)
+        return Image.fromarray(out)
+
+    if t == "adv_patch":
+        # random square occlusion / noise patch
+        out = arr.copy()
+        w, h = img.size
+        s = int(min(w, h) * (0.10 + 0.30 * (strength / 100.0)))  # 10% -> 40%
+        x0 = np.random.randint(0, max(1, w - s))
+        y0 = np.random.randint(0, max(1, h - s))
+        patch = np.random.uniform(0, 255, size=(s, s, 3)).astype(np.float32)
+        out[y0:y0 + s, x0:x0 + s, :] = patch
+        return Image.fromarray(np.clip(out, 0, 255).astype(np.uint8))
+
+    if t == "adv_stripes":
+        out = arr.copy()
+        h, w = out.shape[0], out.shape[1]
+        period = max(4, int(40 - 30 * (strength / 100.0)))  # 40 -> 10
+        amp = (strength / 100.0) * 35.0
+        for x in range(0, w, period):
+            out[:, x:x+2, :] = np.clip(out[:, x:x+2, :] + amp, 0, 255)
+        return Image.fromarray(out.astype(np.uint8))
+
+    if t == "adv_jpeg":
+        # JPEG compression artifacts
+        quality = int(95 - (strength / 100.0) * 75)  # 95 -> 20
+        quality = max(10, min(95, quality))
+        buf = io.BytesIO()
+        img.save(buf, format="JPEG", quality=quality)
+        buf.seek(0)
+        return Image.open(buf).convert("RGB")
+
+    # fallback: no-op
+    return img
+
+
 # ---- Load runner ONCE at startup ----
+device_env = os.getenv("DEVICE", "").strip()
+if device_env:
+    device = device_env
+else:
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
 runner = ModelRunner(
     project_root=str(PROJECT_ROOT),
-    device=os.getenv("DEVICE", "cuda:0"),
+    device=device,
     yolo_ckpt=OD_VALUE_TO_CKPT[DEFAULT_OD],
     sam_ckpt=DEFAULT_SAM_CKPT,
-    imagebind_ckpt=DEFAULT_IMAGEBIND_CKPT,  # if missing, runner can fall back to pretrained=True
+    imagebind_ckpt=DEFAULT_IMAGEBIND_CKPT,
     id2task_name_file="./id2task_name.json",
     task2prompt_file="./task20.json",
     threshold=0.01,
@@ -74,6 +171,7 @@ runner = ModelRunner(
     forward_thre=0.1,
 )
 
+
 @app.get("/", response_class=HTMLResponse)
 def index(request: Request):
     return templates.TemplateResponse(
@@ -92,14 +190,18 @@ def index(request: Request):
         },
     )
 
+
 @app.post("/api/run")
 async def api_run(
+    request: Request,
     vlm_model: str = Form(DEFAULT_VLM),
     od_model: str = Form(DEFAULT_OD),
     task_id: int = Form(1),
     score_function: str = Form(DEFAULT_SCORE_FUNC),
     hdv_dim: int = Form(DEFAULT_HDV),
     viz_mode: str = Form("bbox"),
+    noise_type: str = Form("none"),
+    noise_strength: int = Form(0),
     upload: UploadFile = File(...),
 ):
     # validate + pick decoder
@@ -121,11 +223,19 @@ async def api_run(
     if not yolo_ckpt:
         return JSONResponse({"ok": False, "error": f"Unknown od_model size: {od_model}"}, status_code=400)
 
-    # save upload
-    suffix = Path(upload.filename).suffix or ".jpg"
+    # save upload (apply noise first)
     job_id = uuid.uuid4().hex
+    suffix = Path(upload.filename).suffix or ".jpg"
     upload_path = UPLOAD_DIR / f"{job_id}{suffix}"
-    upload_path.write_bytes(await upload.read())
+
+    raw = await upload.read()
+    try:
+        img = Image.open(io.BytesIO(raw)).convert("RGB")
+    except Exception:
+        return JSONResponse({"ok": False, "error": "Failed to decode image upload"}, status_code=400)
+
+    img = apply_noise_pil(img, noise_type=noise_type, strength_0_100=noise_strength)
+    img.save(upload_path, quality=95)
 
     # run
     try:
@@ -141,13 +251,9 @@ async def api_run(
             viz_mode=viz_mode,
         )
     except Exception as e:
-        # return JSONResponse({"ok": False, "error": repr(e)}, status_code=500)
         tb = traceback.format_exc()
-        print(tb)  # shows in HF logs
-        return JSONResponse(
-            {"ok": False, "error": str(e), "traceback": tb},
-            status_code=500
-        )
+        print(tb)
+        return JSONResponse({"ok": False, "error": str(e), "traceback": tb}, status_code=500)
 
     # save results
     job_dir = RESULT_DIR / job_id
@@ -161,6 +267,7 @@ async def api_run(
     out["images"]["yolo"].save(p_yolo, quality=95)
     out["images"]["selected"].save(p_sel, quality=95)
 
+    base = str(request.base_url).rstrip("/")
     return {
         "ok": True,
         "job_id": job_id,
@@ -168,8 +275,8 @@ async def api_run(
         "task_name": out["task_name"],
         "selected_indices": out["selected_indices"],
         "image_urls": {
-            "input": f"/results/{job_id}/input.jpg",
-            "yolo": f"/results/{job_id}/yolo.jpg",
-            "selected": f"/results/{job_id}/selected.jpg",
+            "input": f"{base}/results/{job_id}/input.jpg",
+            "yolo": f"{base}/results/{job_id}/yolo.jpg",
+            "selected": f"{base}/results/{job_id}/selected.jpg",
         },
-    }
+    }