reorder

tour-comfy/embeddings.py

@@ -2,29 +2,39 @@ import torch
 import open_clip
 from PIL import Image
 import os
+import threading
 
 _model = None
 _preprocess = None
 _device = None
+_lock = threading.Lock()
 
 def get_model():
     global _model, _preprocess, _device
     if _model is None:
-        _device = "cuda" if torch.cuda.is_available() else "cpu"
-        # ViT-H-14 is 1024-dim
-        _model, _, _preprocess = open_clip.create_model_and_transforms('ViT-H-14', pretrained='laion2b_s32b_b79k')
-        _model = _model.to(_device)
-        _model.eval()
+        with _lock:
+            if _model is None:
+                dev = "cuda" if torch.cuda.is_available() else "cpu"
+                # ViT-H-14 is 1024-dim
+                model, _, preprocess = open_clip.create_model_and_transforms('ViT-H-14', pretrained='laion2b_s32b_b79k')
+                model = model.to(dev)
+                model.eval()
+                
+                # Set globals only when fully ready to avoid race conditions
+                _preprocess = preprocess
+                _device = dev
+                _model = model
     return _model, _preprocess, _device
 
 def generate_embedding(image_path):
     model, preprocess, device = get_model()
     try:
-        image = preprocess(Image.open(image_path)).unsqueeze(0).to(device)
-        with torch.no_grad():
-            image_features = model.encode_image(image)
-            image_features /= image_features.norm(dim=-1, keepdim=True)
-            return image_features.cpu().numpy()[0].tolist()
+        with Image.open(image_path) as img:
+            image = preprocess(img.convert("RGB")).unsqueeze(0).to(device)
+            with torch.no_grad():
+                image_features = model.encode_image(image)
+                image_features /= image_features.norm(dim=-1, keepdim=True)
+                return image_features.cpu().numpy()[0].tolist()
     except Exception as e:
         print(f"Error generating embedding for {image_path}: {e}")
         return None