sync from GitHub: fix: single model instance, OOM retry with batch reduction, better logging

demo/qdrant_utils.py

@@ -43,7 +43,9 @@ def init_embedder(model_name: str):
     try:
         from visual_rag import VisualEmbedder
 
-        return VisualEmbedder(model_name=model_name), None
+        embedder = VisualEmbedder(model_name=model_name)
+        embedder._load_model()
+        return embedder, None
     except Exception as e:
         return None, f"{e}\n\n{traceback.format_exc()}"
 

demo/ui/playground.py

@@ -47,13 +47,14 @@ def render_playground_tab():
         if not st.session_state.get("model_loaded"):
             with st.spinner(f"Loading {model_short}..."):
                 try:
-                    _ = MultiVectorRetriever(
-                        collection_name=active_collection, model_name=model_name,
-                        qdrant_url=url, qdrant_api_key=api_key,
-                    )
-                    st.session_state["model_loaded"] = True
-                    st.session_state["loaded_model_key"] = cache_key
-                    st.session_state["loaded_model_name"] = model_name
+                    from demo.qdrant_utils import init_embedder
+                    embedder, err = init_embedder(model_name)
+                    if err:
+                        st.warning(f"Failed: {model_short}")
+                    else:
+                        st.session_state["model_loaded"] = True
+                        st.session_state["loaded_model_key"] = cache_key
+                        st.session_state["loaded_model_name"] = model_name
                 except Exception:
                     st.warning(f"Failed: {model_short}")
 

demo/ui/upload.py

@@ -269,15 +269,12 @@ def process_pdfs(uploaded_files, config):
             model_status.info(f"Loading `{model_short}`...")
 
             output_dtype = np.float16 if vector_dtype == "float16" else np.float32
-            embedder_key = f"{model_name}::{vector_dtype}"
-            embedder = None
-            if st.session_state.get("upload_embedder_key") == embedder_key:
-                embedder = st.session_state.get("upload_embedder")
-            if embedder is None:
-                embedder = VisualEmbedder(model_name=model_name, output_dtype=output_dtype)
-                embedder._load_model()
-                st.session_state["upload_embedder_key"] = embedder_key
-                st.session_state["upload_embedder"] = embedder
+            from demo.qdrant_utils import init_embedder
+            embedder, err = init_embedder(model_name)
+            if err:
+                model_status.error(f"❌ Failed to load model: {err[:100]}")
+                return
+            embedder.output_dtype = output_dtype
             model_status.success(f"✅ Model `{model_short}` loaded ({vector_dtype})")
 
         with phase2:

visual_rag/embedding/visual_embedder.py

@@ -618,13 +618,55 @@ class VisualEmbedder:
 
         for i in iterator:
             batch = images[i : i + batch_size]
+            current_batch_size = len(batch)
 
-            with torch.no_grad():
-                processed = self.processor.process_images(batch).to(self.model.device)
+            while current_batch_size > 0:
+                try:
+                    sub_batch = batch[:current_batch_size]
+                    self._embed_image_batch(
+                        sub_batch, embeddings, token_infos, return_token_info
+                    )
+                    if current_batch_size < len(batch):
+                        remaining = batch[current_batch_size:]
+                        for single in remaining:
+                            self._embed_image_batch(
+                                [single], embeddings, token_infos, return_token_info
+                            )
+                    break
+                except (torch.cuda.OutOfMemoryError, RuntimeError) as e:
+                    if "out of memory" not in str(e).lower() and "CUDA" not in str(e):
+                        raise
+                    gc.collect()
+                    if torch.cuda.is_available():
+                        torch.cuda.empty_cache()
+                    new_size = max(1, current_batch_size // 2)
+                    logger.warning(
+                        f"⚠️ CUDA OOM with batch_size={current_batch_size}, "
+                        f"retrying with batch_size={new_size}"
+                    )
+                    print(
+                        f"[OOM] Reducing batch size from {current_batch_size} to {new_size}"
+                    )
+                    if new_size == current_batch_size:
+                        raise
+                    current_batch_size = new_size
+
+        if return_token_info:
+            return embeddings, token_infos
+        return embeddings
 
-                # Extract token info before model forward
-                if return_token_info:
-                    input_ids = processed["input_ids"]
+    def _embed_image_batch(
+        self,
+        batch: List[Image.Image],
+        embeddings: list,
+        token_infos: Optional[list],
+        return_token_info: bool,
+    ):
+        with torch.no_grad():
+            processed = self.processor.process_images(batch).to(self.model.device)
+
+            if return_token_info:
+                input_ids = processed["input_ids"]
                     batch_n_rows = processed.get("n_rows")
                     batch_n_cols = processed.get("n_cols")
                     # Qwen2/2.5-VL style grid information (T, H, W)
@@ -681,27 +723,20 @@ class VisualEmbedder:
                             }
                         )
 
-                # Generate embeddings
-                batch_embeddings = self.model(**processed)
-
-            # Extract per-image embeddings
-            if isinstance(batch_embeddings, torch.Tensor) and batch_embeddings.dim() == 3:
-                for j in range(batch_embeddings.shape[0]):
-                    embeddings.append(batch_embeddings[j].cpu())
-            else:
-                embeddings.extend([e.cpu() for e in batch_embeddings])
-
-            # Memory cleanup
-            del processed, batch_embeddings
-            gc.collect()
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-            elif torch.backends.mps.is_available():
-                torch.mps.empty_cache()
+            batch_embeddings = self.model(**processed)
 
-        if return_token_info:
-            return embeddings, token_infos
-        return embeddings
+        if isinstance(batch_embeddings, torch.Tensor) and batch_embeddings.dim() == 3:
+            for j in range(batch_embeddings.shape[0]):
+                embeddings.append(batch_embeddings[j].cpu())
+        else:
+            embeddings.extend([e.cpu() for e in batch_embeddings])
+
+        del processed, batch_embeddings
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        elif torch.backends.mps.is_available():
+            torch.mps.empty_cache()
 
     def extract_visual_embedding(
         self,