Added on demand model loading

app.py

@@ -3,8 +3,21 @@ import gradio as gr
 from transformers import Owlv2Processor, Owlv2ForObjectDetection
 import os
 import torchvision
+import shutil
 
 # --- Setup ---
+# Clean caches each restart (helps avoid 50GB limit)
+for cache_dir in [
+    os.path.expanduser("~/.cache/huggingface"),
+    os.path.expanduser("~/.cache/torch"),
+]:
+    shutil.rmtree(cache_dir, ignore_errors=True)
+
+# Force Hugging Face cache to /tmp (ephemeral)
+os.environ["HF_HUB_CACHE"] = "/tmp/hf_cache"
+os.makedirs(os.environ["HF_HUB_CACHE"], exist_ok=True)
+
+# Gradio temp folder
 os.environ["GRADIO_TEMP_DIR"] = "tmp"
 os.makedirs(os.environ["GRADIO_TEMP_DIR"], exist_ok=True)
 
@@ -19,22 +32,34 @@ except ImportError:
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# --- Load Models ---
-print("Loading models...")
-noctowlv2_base = Owlv2ForObjectDetection.from_pretrained(
-    "lorebianchi98/NoctOWLv2-base-patch16"
-).to(device)
-processorv2_base = Owlv2Processor.from_pretrained("google/owlv2-base-patch16")
+# --- Lazy Model Loader ---
+MODELS = {}
 
-noctowlv2_large = Owlv2ForObjectDetection.from_pretrained(
-    "lorebianchi98/NoctOWLv2-large-patch14"
-).to(device)
-processorv2_large = Owlv2Processor.from_pretrained("google/owlv2-large-patch14")
+def get_model(selected_model):
+    """Load model + processor on demand and cache in memory."""
+    if selected_model in MODELS:
+        return MODELS[selected_model]
 
-MODELS = {
-    "NoctOWLv2-Base": (noctowlv2_base, processorv2_base),
-    "NoctOWLv2-Large": (noctowlv2_large, processorv2_large),
-}
+    print(f"Loading {selected_model}...")
+
+    if selected_model == "NoctOWLv2-Base":
+        model = Owlv2ForObjectDetection.from_pretrained(
+            "lorebianchi98/NoctOWLv2-base-patch16"
+        ).to(device)
+        processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16")
+
+    elif selected_model == "NoctOWLv2-Large":
+        model = Owlv2ForObjectDetection.from_pretrained(
+            "lorebianchi98/NoctOWLv2-large-patch14"
+        ).to(device)
+        processor = Owlv2Processor.from_pretrained("google/owlv2-large-patch14")
+
+    else:
+        raise gr.Error(f"Unknown model: {selected_model}")
+
+    # Cache in memory so re-selections don't re-load from disk
+    MODELS[selected_model] = (model, processor)
+    return model, processor
 
 
 # --- Inference Function ---
@@ -47,7 +72,7 @@ def query_image(img, text_queries, score_threshold, selected_model):
     if selected_model is None or selected_model == "":
         raise gr.Error("Please select a model before running inference.")
 
-    model, processor = MODELS[selected_model]
+    model, processor = get_model(selected_model)
     model = model.to(device)
 
     # Prepare text
@@ -154,7 +179,7 @@ with gr.Blocks(title="NoctOWLv2 — Fine-Grained Zero-Shot Object Detection") as
         outputs=output_image,
     )
 
-    # --- Example Images (without predefined model) ---
+    # --- Example Images ---
     gr.Examples(
         examples=[
             ["assets/desciglio.jpg", "striped football shirt, plain red football shirt, yellow shoes, red shoes", 0.07],