feat: dynamically initialize model weights using MODEL_REGISTRY and adjust weights for specific models

agents/ensemble_weights.py

@@ -1,5 +1,6 @@
 import logging
 import torch
+from utils.registry import MODEL_REGISTRY
 
 logger = logging.getLogger(__name__)
 
@@ -15,7 +16,7 @@ class ContextualWeightOverrideAgent:
                 "model_2": 0.7, 
                 "model_7": 1.3,
             },
-            "sunny": {
+            "SDXL": {
                 "model_3": 0.9,
                 "model_4": 1.1,
             }
@@ -36,21 +37,18 @@ class ContextualWeightOverrideAgent:
 
 class ModelWeightManager:
     def __init__(self):
-        self.base_weights = {
-            "model_1": 0.15,  # SwinV2 Based
-            "model_2": 0.15,  # ViT Based
-            "model_3": 0.15,  # SDXL Dataset
-            "model_4": 0.15,  # SDXL + FLUX
-            "model_5": 0.15,  # ViT Based
-            "model_5b": 0.10, # ViT Based, Newer Dataset
-            "model_6": 0.10,  # Swin, Midj + SDXL
-            "model_7": 0.05   # ViT
-        }
+        # Dynamically initialize base_weights for all models
+        self.base_weights = {model_id: 1.0 for model_id in MODEL_REGISTRY}
+        # Assign a higher weight to the strongest model
+        if "simple_prediction" in self.base_weights:
+            self.base_weights["simple_prediction"] = 2.5  # Tune as needed
+
+        # Situation-based multipliers (applied to all models)
         self.situation_weights = {
             "high_confidence": 1.2,    # Boost weights for high confidence predictions
             "low_confidence": 0.8,     # Reduce weights for low confidence
-            "conflict": 0.5,          # Reduce weights when models disagree
-            "consensus": 1.5          # Boost weights when models agree
+            "conflict": 0.5,           # Reduce weights when models disagree
+            "consensus": 1.5           # Boost weights when models agree
         }
         self.context_override_agent = ContextualWeightOverrideAgent()