Update utils/model_loader.py

utils/model_loader.py

@@ -1,18 +1,49 @@
-from transformers import pipeline
+from transformers import pipeline, AutoProcessor, AutoModelForCausalLM
 import torch
+from typing import Optional
 
 def load_llava_model():
+    """Load LLaVA model with 4-bit quantization for HF Spaces"""
+    model_id = "llava-hf/llava-1.5-7b-hf"
+    
     return pipeline(
         "image-to-text",
-        model="llava-hf/llava-1.5-7b-hf",
-        torch_dtype=torch.float16,
+        model=model_id,
         device_map="auto",
-        max_new_tokens=200
+        model_kwargs={
+            "torch_dtype": torch.float16,
+            "load_in_4bit": True,
+            "quantization_config": {
+                "load_in_4bit": True,
+                "bnb_4bit_compute_dtype": torch.float16,
+                "bnb_4bit_use_double_quant": True,
+                "bnb_4bit_quant_type": "nf4"
+            }
+        }
     )
 
-def load_clip_model():
+def load_caption_model():
+    """BLIP-2 with efficient loading"""
     return pipeline(
-        "feature-extraction",
-        model="openai/clip-vit-base-patch32",
-        device_map="auto"
-    )
+        "image-to-text",
+        model="Salesforce/blip2-opt-2.7b",
+        device_map="auto",
+        torch_dtype=torch.float16,
+        model_kwargs={"cache_dir": "/tmp/models"}
+    )
+
+def load_retrieval_models():
+    """Load encoders with shared weights"""
+    models = {}
+    models['text_encoder'] = SentenceTransformer(
+        'sentence-transformers/all-MiniLM-L6-v2',
+        device="cuda" if torch.cuda.is_available() else "cpu"
+    )
+    
+    models['image_encoder'] = AutoModel.from_pretrained(
+        "openai/clip-vit-base-patch32",
+        device_map="auto",
+        torch_dtype=torch.float16
+    )
+    
+    return models