reorder

tour_comfy/embeddings.py

@@ -0,0 +1,30 @@
+import torch
+import open_clip
+from PIL import Image
+import os
+
+_model = None
+_preprocess = None
+_device = None
+
+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()
+    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()
+    except Exception as e:
+        print(f"Error generating embedding for {image_path}: {e}")
+        return None