Update app.py

app.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 """
-🔥 Icon Challenge Solver - FINAL FIX HF Spaces
+🔥 Icon Challenge Solver - FIXED for HF Spaces
 """
 
 import os
@@ -36,10 +36,13 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(
 logger = logging.getLogger(__name__)
 
 # ===================== CONFIG =====================
+# FIX: Gunakan absolute path berdasarkan lokasi file script ini
+BASE_DIR = Path(__file__).resolve().parent
+
 class Config:
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    yolo_model_path = "./yolo/best.pt"
-    lora_model_dir = "./icon_model_hybrid"
+    yolo_model_path = BASE_DIR / "yolo" / "best.pt"
+    lora_model_dir = BASE_DIR / "icon_model_hybrid"
     adapter_subdir = "lora_adapter"
     model_name = "ViT-L-16-SigLIP-256"
     yolo_conf = 0.25
@@ -62,24 +65,50 @@ class ProjectionHead(nn.Module):
         self.skip = nn.Linear(in_dim, out_dim)
 
     def forward(self, x):
+        # FIX: Handle batch size 1 untuk BatchNorm
         if x.shape[0] <= 1:
             self.train(False)
         residual = self.skip(x)
-        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)
+        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):
         super().__init__()
         self.backbone = backbone
         self.projection = ProjectionHead()
-        self.weight = nn.Parameter(head_state['arc_margin']['weight'])
-        self.projection.load_state_dict(head_state['projection'])
+        
+        # FIX: Load projection weights dengan strict=False untuk handle key mismatch
+        if 'projection' in head_state:
+            proj_dict = head_state['projection']
+            # Handle case kalau key-nya ada prefix '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)
+        
+        # FIX: Handle arc_margin weight dengan berbagai kemungkinan key name
+        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
 
     def forward(self, x):
@@ -109,19 +138,39 @@ class ModelManager:
         return cls._instance
 
     def init(self):
-        if self.initialized: return
+        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
+        # FIX: Check YOLO path dengan absolute path
         yolo_path = Path(CONFIG.yolo_model_path)
-        if yolo_path.exists():
-            self.yolo = YOLO(str(yolo_path)).to(CONFIG.device)
-            logger.info(f"✅ YOLO loaded: {yolo_path}")
-        else:
-            raise FileNotFoundError(f"YOLO not found: {yolo_path}")
+        logger.info(f"Looking for YOLO at: {yolo_path} (exists: {yolo_path.exists()})")
+        
+        if not yolo_path.exists():
+            # Fallback: cari di cwd juga
+            alt_path = Path("yolo") / "best.pt"
+            if alt_path.exists():
+                yolo_path = alt_path
+            else:
+                raise FileNotFoundError(f"YOLO not found: {yolo_path} or {alt_path}")
+
+        self.yolo = YOLO(str(yolo_path)).to(CONFIG.device)
+        logger.info(f"✅ YOLO loaded: {yolo_path}")
 
         # Config
         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 = Path("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}")
@@ -130,49 +179,72 @@ class ModelManager:
             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 Adapter
         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}")
 
-        model_with_lora = PeftModel.from_pretrained(
-            backbone, peft_path, is_trainable=False
-        )
+        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}")
+            # Fallback: gunakan backbone tanpa LoRA
+            model_with_lora = backbone
+            logger.warning("⚠️ Using backbone without LoRA!")
 
         # Head Weights
         head_candidates = [
             lora_base / "head_weights.pt",
-            lora_base / "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}")
                 break
+        
         if head_state is None:
-            raise FileNotFoundError("head_weights.pt or weights.pt not found!")
+            raise FileNotFoundError(f"Head weights not found! Tried: {head_candidates}")
 
         self.inference_model = IconClassifier(model_with_lora, head_state, num_classes).to(CONFIG.device).eval()
+        logger.info("✅ IconClassifier created")
 
         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!")
 
     def yolo_detect_and_crop(self, image: Image.Image) -> Image.Image:
+        # FIX: Handle PIL Image dengan berbagai mode
+        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
         for result in results:
-            if result.boxes is None: continue
+            if result.boxes is None:
+                continue
             for box in result.boxes:
                 conf = float(box.conf[0])
                 if conf > best_conf:
@@ -182,24 +254,39 @@ class ModelManager:
         if best_box is not None and best_conf >= CONFIG.yolo_min_conf:
             x1, y1, x2, y2 = best_box
             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:
-                return Image.fromarray(img_np[y1:y2, x1:x2])
+            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
 
     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():
-            return self.inference_model.get_embedding(img_tensor)
+            emb = self.inference_model.get_embedding(img_tensor)
+        return emb
 
     def match_options(self, challenge_crop: Image.Image, options: List[Image.Image]):
         challenge_emb = self.extract_embedding(challenge_crop)
-        option_embs = [self.extract_embedding(opt) for opt in options]
+        
+        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))
-        exp_sim = np.exp(similarities * 10)
+        
+        # FIX: Temperature scaling yang lebih stabil
+        sim_scaled = similarities * 10
+        sim_max = np.max(sim_scaled)
+        exp_sim = np.exp(sim_scaled - sim_max)  # Numerical stability
         probs = exp_sim / exp_sim.sum()
 
         return {
@@ -217,32 +304,59 @@ manager = ModelManager()
 # ===================== MAIN FUNCTION =====================
 @spaces.GPU(request_timeout=180, enable=True)
 def solve_gradio(challenge_img, *options):
+    # FIX: Inisialisasi dengan try-except yang lebih detail
     if not manager.initialized:
-        manager.init()
+        try:
+            manager.init()
+        except Exception as e:
+            logger.error(f"Model init failed: {e}", exc_info=True)
+            return f"❌ Model initialization failed: {str(e)}\n\nCek path model di HF Spaces!", None
 
     try:
         if challenge_img is None or any(o is None for o in options):
             return "❌ Semua gambar harus diisi!", None
 
+        # FIX: Handle berbagai input format (numpy, PIL, path)
         if isinstance(challenge_img, np.ndarray):
             challenge_img = Image.fromarray(challenge_img)
-        options_pil = [Image.fromarray(opt) if isinstance(opt, np.ndarray) else opt for opt in options]
+        elif isinstance(challenge_img, str):
+            challenge_img = Image.open(challenge_img)
+            
+        options_pil = []
+        for opt in options:
+            if isinstance(opt, np.ndarray):
+                options_pil.append(Image.fromarray(opt))
+            elif isinstance(opt, str):
+                options_pil.append(Image.open(opt))
+            elif isinstance(opt, Image.Image):
+                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
 
+        # YOLO crop
         challenge_crop = manager.yolo_detect_and_crop(challenge_img)
+        
+        # Match
         result = manager.match_options(challenge_crop, options_pil)
 
         best_idx = result["predicted_index"]
-        lines = ["🎯 HASIL\n",
-                 f"✅ Best Option: {best_idx + 1}",
-                 f"Confidence: {result['confidence']:.4f} ({result['confidence']*100:.2f}%)",
-                 f"Class: {result['predicted_class']}\n",
-                 "📊 Skor Semua Opsi:"]
+        lines = [
+            "🎯 HASIL\n",
+            f"✅ Best Option: {best_idx + 1}",
+            f"Confidence: {result['confidence']:.4f} ({result['confidence']*100:.2f}%)",
+            f"Class: {result['predicted_class']}\n",
+            "📊 Skor Semua Opsi:"
+        ]
 
-        for i in range(8):
+        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})")
 
+        # FIX: Cleanup VRAM setelah inference
         gc.collect()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
@@ -280,4 +394,6 @@ with gr.Blocks(title="🔥 Icon Challenge Solver", theme=gr.themes.Soft()) as de
 
     gr.Markdown("---\nMade with ❤️ for Icon Challenge")
 
-demo.launch()
+# FIX: Gunakan if __name__ guard untuk HF Spaces
+if __name__ == "__main__":
+    demo.launch()