Update app.py

app.py

@@ -130,15 +130,11 @@ def format_derm_disclaimer(ans: str) -> str:
     )
 
 def _load_base_plus_lora(dtype: torch.dtype = torch.float16):
-    """
-    Load base model on GPU and attach LoRA adapters from ADAPTER_ID.
-    Returns the Peft-wrapped model (eval mode).
-    """
     logger.info(f"Loading BASE on GPU: {BASE_MODEL_ID}")
     base = VisionTextModelClass.from_pretrained(
         BASE_MODEL_ID,
         torch_dtype=dtype,
-        device_map="cuda",         # ZeroGPU worker device
+        device_map="cuda",
         trust_remote_code=True,
         low_cpu_mem_usage=True,
     )
@@ -148,6 +144,7 @@ def _load_base_plus_lora(dtype: torch.dtype = torch.float16):
     model.eval()
     return model
 
+
 # ---------------------------
 # Inference (ZeroGPU-safe: only here we touch CUDA)
 # ---------------------------
@@ -205,34 +202,62 @@ def create_interface() -> gr.Blocks:
 
 def _get_sanitized_adapter_dir(adapter_id: str) -> str:
     """
-    Download the adapter repo and remove unsupported keys from adapter_config.json
-    so PEFT can load it.
+    Download the adapter repo locally and sanitize adapter_config.json so PEFT LoraConfig
+    can parse it (remove any unknown / experimental blocks like 'corda_config', 'eva_config', etc.).
     """
     repo_dir = snapshot_download(adapter_id)
-    tmp_dir = tempfile.mkdtemp(prefix="peft_adapter_")
-    shutil.copytree(repo_dir, os.path.join(tmp_dir, "adapter"), dirs_exist_ok=True)
-    cfg_path = os.path.join(tmp_dir, "adapter", "adapter_config.json")
-
-    try:
-        with open(cfg_path, "r") as f:
-            cfg = json.load(f)
-    except Exception as e:
-        raise RuntimeError(f"Failed to read adapter_config.json: {e}")
-
-    # Remove keys PEFT LoraConfig doesn't recognize
-    for k in ["corda_config", "CoRDA_config"]:
-        if k in cfg:
-            cfg.pop(k)
-
-    # If DoRA wasn’t used, drop its block too
-    if str(cfg.get("use_dora", "false")).lower() in ["false", "0", "no"]:
-        if "dora_config" in cfg:
-            cfg.pop("dora_config")
+    tmp_root = tempfile.mkdtemp(prefix="peft_adapter_")
+    adapter_dir = os.path.join(tmp_root, "adapter")
+    shutil.copytree(repo_dir, adapter_dir, dirs_exist_ok=True)
+
+    cfg_path = os.path.join(adapter_dir, "adapter_config.json")
+    if not os.path.isfile(cfg_path):
+        raise RuntimeError(f"adapter_config.json not found in adapter repo: {adapter_id}")
+
+    with open(cfg_path, "r") as f:
+        cfg = json.load(f)
+
+    # Minimal, widely-supported PEFT LoRA keys:
+    allowed = {
+        "peft_type", "task_type",
+        "r", "lora_alpha", "lora_dropout",
+        "target_modules", "bias",
+        "inference_mode",
+        "base_model_name_or_path",
+        "fan_in_fan_out",
+        "modules_to_save",
+        "layers_to_transform",
+        "layers_pattern",
+        "use_rslora",
+        "rank_dropout", "module_dropout",
+        "init_lora_weights",
+        "use_dora",  # keep if your PEFT version supports DoRA; harmless otherwise if False
+    }
+
+    # If DoRA isn't actually used, nuke its block to avoid parser issues
+    if str(cfg.get("use_dora", "false")).lower() in ("false", "0", "no"):
+        cfg.pop("dora_config", None)
+
+    # Drop any unknown top-level configs (e.g., 'corda_config', 'CoRDA_config', 'eva_config', etc.)
+    to_delete = [k for k in cfg.keys() if k not in allowed]
+    for k in to_delete:
+        cfg.pop(k, None)
+
+    # Ensure required fields exist
+    cfg.setdefault("peft_type", "LORA")
+    cfg.setdefault("task_type", "CAUSAL_LM")
+    cfg.setdefault("bias", "none")
+    cfg.setdefault("inference_mode", True)
+
+    # Normalize booleans in case they were strings
+    for k in ("inference_mode", "use_rslora", "use_dora", "fan_in_fan_out"):
+        if k in cfg and isinstance(cfg[k], str):
+            cfg[k] = cfg[k].lower() in ("true", "1", "yes")
 
     with open(cfg_path, "w") as f:
         json.dump(cfg, f, indent=2)
 
-    return os.path.join(tmp_dir, "adapter")
+    return adapter_dir
 
 def main():
     demo = create_interface()