Update app.py

app.py

@@ -1,15 +1,10 @@
 #!/usr/bin/env python3
 """
-🔥 Icon Challenge Solver - FINAL FIX HF Spaces
+🔥 Icon Challenge Solver - 100% Exact LoRA Inference Match
 """
 
 # ===================== ASYNCIO WARNING FIX =====================
 import asyncio
-import sys
-
-# Patch: Suppress "Exception ignored in BaseEventLoop.__del__" on Python 3.12
-# Ini terjadi karena Gradio/Ultralytics membuat event loop tanpa close eksplisit
-_original_event_loop_del = asyncio.BaseEventLoop.__del__
 
 def _safe_event_loop_del(self):
     try:
@@ -22,7 +17,6 @@ def _safe_event_loop_del(self):
 
 asyncio.BaseEventLoop.__del__ = _safe_event_loop_del
 
-# Suppress warnings
 import warnings
 warnings.filterwarnings("ignore", category=ResourceWarning)
 warnings.filterwarnings("ignore", message="Exception ignored in.*BaseEventLoop.*")
@@ -30,24 +24,20 @@ warnings.filterwarnings("ignore")
 
 # ===================== IMPORTS =====================
 import os
-import io
-import re
 import gc
-import base64
 import json
 import logging
 from pathlib import Path
-from typing import List, Union
+from typing import List
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.cuda.amp import autocast
+import torchvision.transforms as T
 
 import numpy as np
 from PIL import Image
 import cv2
-import torchvision.transforms as T
 
 import gradio as gr
 import spaces
@@ -56,7 +46,6 @@ from ultralytics import YOLO
 import open_clip
 from peft import PeftModel
 
-# ===================== ENV SETUP =====================
 os.environ["PYTHONASYNCIODEBUG"] = "0"
 os.environ["GRADIO_ANALYTICS_ENABLED"] = "0"
 os.environ["YOLO_CONFIG_DIR"] = "/tmp/Ultralytics"
@@ -80,7 +69,14 @@ class Config:
 
 CONFIG = Config()
 
-# ===================== MODEL CLASSES =====================
+# ===================== PREPROCESS (EXACT MATCH) =====================
+preprocess = T.Compose([
+    T.Resize((256, 256), interpolation=T.InterpolationMode.NEAREST),
+    T.ToTensor(),
+    T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+])
+
+# ===================== PROJECTION HEAD (EXACT MATCH) =====================
 class ProjectionHead(nn.Module):
     def __init__(self, in_dim=1024, hidden_dim=512, out_dim=256, dropout=0.1):
         super().__init__()
@@ -93,66 +89,41 @@ class ProjectionHead(nn.Module):
         self.skip = nn.Linear(in_dim, out_dim)
 
     def forward(self, x):
+        # EXACT: self.eval() bukan self.train(False)
         if x.shape[0] <= 1:
-            self.train(False)
+            self.eval()
         residual = self.skip(x)
-        out = self.layer1(x)
-        out = self.bn1(out)
-        out = self.act(out)
-        out = self.dropout(out)
-        out = self.layer2(out)
-        out = self.bn2(out)
-        return F.normalize(out + residual, p=2, dim=1)
-
-class IconClassifier(nn.Module):
-    def __init__(self, backbone, head_state, num_classes=21):
+        x = self.layer1(x)
+        x = self.bn1(x)
+        x = self.act(x)
+        x = self.dropout(x)
+        x = self.layer2(x)
+        x = self.bn2(x)
+        return F.normalize(x + residual, p=2, dim=1)
+
+# ===================== ICON MATCHER (EXACT MATCH REFERENSI) =====================
+class IconMatcher(nn.Module):
+    """
+    100% match dengan IconMatcherCPU dari referensi,
+    cuma ditambah .to(device) dan path resolution.
+    """
+    def __init__(self, lora_path, head_path, device='cuda'):
         super().__init__()
-        self.backbone = backbone
-        self.projection = ProjectionHead()
-        
-        # Load projection weights
-        if 'projection' in head_state:
-            proj_dict = head_state['projection']
-            cleaned = {}
-            for k, v in proj_dict.items():
-                if k.startswith('projection.'):
-                    cleaned[k[11:]] = v
-                else:
-                    cleaned[k] = v
-            self.projection.load_state_dict(cleaned, strict=False)
-        
-        # Find weight dengan berbagai kemungkinan key
-        weight = None
-        if 'arc_margin' in head_state and 'weight' in head_state['arc_margin']:
-            weight = head_state['arc_margin']['weight']
-        elif 'weight' in head_state:
-            weight = head_state['weight']
-        elif 'fc.weight' in head_state:
-            weight = head_state['fc.weight']
-        
-        if weight is None:
-            logger.warning("Head weight not found, initializing random!")
-            weight = torch.randn(num_classes, 256)
-        
-        self.weight = nn.Parameter(weight)
-        self.s = 30.0
+        backbone, _, _ = open_clip.create_model_and_transforms(
+            'ViT-L-16-SigLIP-256', pretrained='webli'
+        )
+        self.backbone = PeftModel.from_pretrained(backbone.visual, lora_path)
+        self.backbone = self.backbone.to(device)
 
-    def forward(self, x):
-        features = self.backbone(x)
-        embedding = self.projection(features)
-        logits = F.linear(F.normalize(embedding), F.normalize(self.weight))
-        return logits * self.s
+        head_state = torch.load(head_path, map_location=device)
+        self.projection = ProjectionHead(1024, 512, 256)
+        self.projection.load_state_dict(head_state['projection'])
+        self.projection = self.projection.to(device)
 
-    def get_embedding(self, x):
-        features = self.backbone(x)
-        embedding = self.projection(features)
-        return F.normalize(embedding, p=2, dim=1)
-
-custom_preprocess = T.Compose([
-    T.Resize((256, 256), interpolation=T.InterpolationMode.NEAREST),
-    T.ToTensor(),
-    T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
-])
+    def forward(self, x):
+        with torch.no_grad():
+            features = self.backbone(x)
+            return self.projection(features)
 
 # ===================== MODEL MANAGER =====================
 class ModelManager:
@@ -166,105 +137,76 @@ class ModelManager:
     def init(self):
         if self.initialized:
             return
-        
-        logger.info(f"Working directory: {os.getcwd()}")
-        logger.info(f"Base directory: {BASE_DIR}")
-        logger.info("Loading models on GPU...")
 
-        # YOLO
+        logger.info(f"Base dir: {BASE_DIR}")
+        logger.info("Loading models...")
+
+        # --- YOLO ---
         yolo_path = Path(CONFIG.yolo_model_path)
-        logger.info(f"Looking for YOLO at: {yolo_path} (exists: {yolo_path.exists()})")
-        
         if not yolo_path.exists():
-            alt_path = BASE_DIR / "best.pt"
-            if alt_path.exists():
-                yolo_path = alt_path
-            else:
-                raise FileNotFoundError(f"YOLO not found: {yolo_path}")
+            yolo_path = BASE_DIR / "best.pt"
+        if not yolo_path.exists():
+            raise FileNotFoundError(f"YOLO not found: {yolo_path}")
 
         self.yolo = YOLO(str(yolo_path)).to(CONFIG.device)
         logger.info(f"✅ YOLO loaded: {yolo_path}")
 
-        # Config
+        # --- LoRA Path ---
         lora_base = Path(CONFIG.lora_model_dir)
-        logger.info(f"Looking for LoRA at: {lora_base} (exists: {lora_base.exists()})")
-        
         if not lora_base.exists():
-            alt_lora = BASE_DIR / "icon_model_hybrid"
-            if alt_lora.exists():
-                lora_base = alt_lora
-            else:
-                raise FileNotFoundError(f"LoRA dir not found: {lora_base}")
-
-        config_path = lora_base / "config.json"
-        if not config_path.exists():
-            raise FileNotFoundError(f"config.json not found in {lora_base}")
-
-        with open(config_path) as f:
-            model_config = json.load(f)
-
-        num_classes = model_config.get('num_classes') or len(model_config.get('class_to_idx', {}))
-        logger.info(f"Num classes: {num_classes}")
-
-        # Backbone
-        logger.info(f"Loading backbone: {CONFIG.model_name}")
-        backbone, _, _ = open_clip.create_model_and_transforms(CONFIG.model_name, pretrained='webli')
-        backbone = backbone.visual.to(CONFIG.device)
-        logger.info("✅ Backbone loaded")
+            lora_base = BASE_DIR / "icon_model_hybrid"
+        if not lora_base.exists():
+            raise FileNotFoundError(f"LoRA dir not found")
 
-        # LoRA Adapter
+        # Cek adapter subdir atau langsung root
         adapter_dir = lora_base / CONFIG.adapter_subdir
-        peft_path = str(adapter_dir) if adapter_dir.exists() else str(lora_base)
-        logger.info(f"Loading LoRA from: {peft_path}")
-
-        try:
-            model_with_lora = PeftModel.from_pretrained(
-                backbone, peft_path, is_trainable=False
-            )
-            logger.info("✅ LoRA adapter loaded")
-        except Exception as e:
-            logger.error(f"LoRA load failed: {e}")
-            model_with_lora = backbone
-            logger.warning("⚠️ Using backbone without LoRA!")
+        lora_path = str(adapter_dir) if adapter_dir.exists() else str(lora_base)
+        logger.info(f"LoRA path: {lora_path}")
 
-        # Head Weights
+        # --- Head Weights ---
         head_candidates = [
             lora_base / "head_weights.pt",
             lora_base / "weights.pt",
-            lora_base / "classifier.pt",
-            lora_base / "model.pt"
         ]
-        head_state = None
         head_path = None
-        for hpath in head_candidates:
-            if hpath.exists():
-                head_state = torch.load(hpath, map_location=CONFIG.device)
-                head_path = hpath
-                logger.info(f"✅ Head weights loaded: {hpath}")
+        for h in head_candidates:
+            if h.exists():
+                head_path = str(h)
+                logger.info(f"✅ Head weights: {h}")
                 break
-        
-        if head_state is None:
-            raise FileNotFoundError(f"Head weights not found! Tried: {head_candidates}")
+        if head_path is None:
+            raise FileNotFoundError(f"Head weights not found in {lora_base}")
 
-        self.inference_model = IconClassifier(model_with_lora, head_state, num_classes).to(CONFIG.device).eval()
-        logger.info("✅ IconClassifier created")
+        # --- Load IconMatcher (EXACT MATCH REFERENSI) ---
+        self.matcher = IconMatcher(lora_path, head_path, device=CONFIG.device).eval()
+        logger.info("✅ IconMatcher loaded (100% exact with reference)")
+
+        # --- Config untuk class names ---
+        config_path = lora_base / "config.json"
+        if config_path.exists():
+            with open(config_path) as f:
+                cfg = json.load(f)
+            self.idx_to_class = {
+                int(k): v for k, v in cfg.get('idx_to_class', {}).items()
+            }
+        else:
+            self.idx_to_class = {}
 
-        self.idx_to_class = {int(k): v for k, v in model_config.get('idx_to_class', {}).items()}
-        self.class_to_idx = {v: int(k) for k, v in self.idx_to_class.items()}
         self.initialized = True
-        logger.info("✅ All models loaded successfully!")
+        logger.info("✅ All models ready!")
 
+    # --- YOLO: Detect + Crop ---
     def yolo_detect_and_crop(self, image: Image.Image) -> Image.Image:
         if image.mode != 'RGB':
             image = image.convert('RGB')
-            
+
         img_np = np.array(image)
         img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
-        
+
         results = self.yolo(img_bgr, conf=CONFIG.yolo_conf, imgsz=CONFIG.yolo_img_size, verbose=False)
-        
+
         best_box = None
-        best_conf = 0
+        best_conf = 0.0
         for result in results:
             if result.boxes is None:
                 continue
@@ -279,69 +221,82 @@ class ModelManager:
             h, w = img_np.shape[:2]
             x1, y1, x2, y2 = max(0, x1), max(0, y1), min(w, x2), min(h, y2)
             if (x2 - x1) >= CONFIG.yolo_min_crop_size and (y2 - y1) >= CONFIG.yolo_min_crop_size:
-                cropped = img_np[y1:y2, x1:x2]
-                return Image.fromarray(cropped)
-        
-        return image
+                return Image.fromarray(img_np[y1:y2, x1:x2])
 
-    def extract_embedding(self, image: Image.Image) -> torch.Tensor:
-        if image.mode != 'RGB':
-            image = image.convert('RGB')
-            
-        img_tensor = custom_preprocess(image).unsqueeze(0).to(CONFIG.device)
-        
-        with torch.no_grad(), autocast():
-            emb = self.inference_model.get_embedding(img_tensor)
-        return emb
+        return image
 
+    # --- INFERENCE: 100% EXACT MATCH REFERENSI ---
     def match_options(self, challenge_crop: Image.Image, options: List[Image.Image]):
-        challenge_emb = self.extract_embedding(challenge_crop)
-        
-        option_embs = []
-        for opt in options:
-            option_embs.append(self.extract_embedding(opt))
-            
-        option_stack = torch.cat(option_embs, dim=0)
-        similarities = (challenge_emb @ option_stack.T).squeeze(0).cpu().numpy()
-
-        best_idx = int(np.argmax(similarities))
-        
-        sim_scaled = similarities * 10
-        sim_max = np.max(sim_scaled)
-        exp_sim = np.exp(sim_scaled - sim_max)
-        probs = exp_sim / exp_sim.sum()
+        """
+        EXACT MATCH dengan:
+        imgs = [Image.open(p).convert('RGB') for p in [target_path] + option_paths]
+        batch = torch.stack([preprocess(img) for img in imgs])
+        embs = model(batch)
+        scores = torch.mm(embs[0:1], embs[1:].t()).squeeze(0)
+        """
+        # 1. Siapkan semua image (challenge + 8 options)
+        imgs = [challenge_crop] + options
+        pil_imgs = []
+        for img in imgs:
+            if img.mode != 'RGB':
+                img = img.convert('RGB')
+            pil_imgs.append(img)
+
+        # 2. Preprocess & stack (EXACT: torch.stack + preprocess)
+        tensors = [preprocess(img) for img in pil_imgs]
+        batch = torch.stack(tensors).to(CONFIG.device)
+
+        # 3. Forward sekali untuk semua (EXACT: model(batch))
+        with torch.no_grad():
+            embs = self.matcher(batch)
+
+        # 4. Similarity (EXACT: torch.mm(embs[0:1], embs[1:].t()).squeeze(0))
+        scores = torch.mm(embs[0:1], embs[1:].t()).squeeze(0)
+        best_idx = int(torch.argmax(scores).item())
+        max_score = float(scores.max().item())
+
+        # 5. Buat probability dari scores untuk display (optional, nggak di referensi tapi buat UI)
+        # Karena embedding sudah L2 normalized, scores = cosine similarity [-1, 1]
+        # Kita convert ke "probability-like" dengan softmax(scores * 10) seperti sebelumnya
+        # tapi INFERENCE UTAMA tetap pakai torch.mm exact di atas
+        probs = torch.softmax(scores * 10, dim=0).cpu().numpy()
 
         return {
             "predicted_index": best_idx,
             "confidence": float(probs[best_idx]),
-            "all_scores": {f"option_{i}": {
-                "similarity": float(similarities[i]),
-                "probability": float(probs[i])
-            } for i in range(len(options))},
+            "max_score": max_score,
+            "all_scores": {
+                f"option_{i}": {
+                    "score": float(scores[i]),
+                    "probability": float(probs[i])
+                }
+                for i in range(len(options))
+            },
             "predicted_class": self.idx_to_class.get(best_idx, "unknown")
         }
 
 manager = ModelManager()
 
-# ===================== MAIN FUNCTION =====================
+# ===================== GRADIO FUNCTION =====================
 @spaces.GPU(request_timeout=180, enable=True)
 def solve_gradio(challenge_img, *options):
     if not manager.initialized:
         try:
             manager.init()
         except Exception as e:
-            logger.error(f"Model init failed: {e}", exc_info=True)
-            return f"❌ Model initialization failed: {str(e)}", None
+            logger.error(f"Init failed: {e}", exc_info=True)
+            return f"❌ Init failed: {e}", None
 
     try:
         if challenge_img is None or any(o is None for o in options):
             return "❌ Semua gambar harus diisi!", None
 
+        # Convert ke PIL
         if isinstance(challenge_img, np.ndarray):
             challenge_img = Image.fromarray(challenge_img)
         elif isinstance(challenge_img, str):
             challenge_img = Image.open(challenge_img)
-            
+
         options_pil = []
         for opt in options:
             if isinstance(opt, np.ndarray):
@@ -352,11 +307,14 @@ def solve_gradio(challenge_img, *options):
                 options_pil.append(opt)
             else:
                 options_pil.append(None)
-        
+
         if any(o is None for o in options_pil):
-            return "❌ Format gambar option tidak valid!", None
+            return "❌ Format option tidak valid!", None
 
+        # YOLO crop
         challenge_crop = manager.yolo_detect_and_crop(challenge_img)
+
+        # Inference (EXACT MATCH REFERENSI)
         result = manager.match_options(challenge_crop, options_pil)
 
         best_idx = result["predicted_index"]
@@ -364,6 +322,7 @@ def solve_gradio(challenge_img, *options):
             "🎯 HASIL\n",
             f"✅ Best Option: {best_idx + 1}",
             f"Confidence: {result['confidence']:.4f} ({result['confidence']*100:.2f}%)",
+            f"Raw Score: {result['max_score']:.4f}",
             f"Class: {result['predicted_class']}\n",
             "📊 Skor Semua Opsi:"
         ]
@@ -371,7 +330,7 @@ def solve_gradio(challenge_img, *options):
         for i in range(len(options_pil)):
             s = result["all_scores"][f"option_{i}"]
             marker = "✅" if i == best_idx else "   "
-            lines.append(f"{marker} Option {i+1}: {s['probability']:.4f} ({s['similarity']:.4f})")
+            lines.append(f"{marker} Option {i+1}: prob={s['probability']:.4f} score={s['score']:.4f}")
 
         gc.collect()
         if torch.cuda.is_available():
@@ -384,9 +343,9 @@ def solve_gradio(challenge_img, *options):
         return f"❌ Error: {str(e)}", None
 
 
-# ===================== GRADIO INTERFACE =====================
+# ===================== GRADIO UI =====================
 with gr.Blocks(title="🔥 Icon Challenge Solver", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🔥 Icon Challenge Solver\n**YOLO + SigLIP LoRA Hybrid**")
+    gr.Markdown("# 🔥 Icon Challenge Solver\n**YOLO + SigLIP LoRA Hybrid (Exact Match Inference)**")
 
     with gr.Row():
         with gr.Column(scale=1):