prompt, from image working

app.py

@@ -33,50 +33,21 @@ embedding_base64s = [None for i in range(max_tabs)]
 
 
 def image_to_embedding(input_im):
-    # approch A: 
-    tform = transforms.Compose([
-    transforms.ToTensor(),
-    transforms.Resize(
-        (336, 336),
-        interpolation=transforms.InterpolationMode.BICUBIC,
-        antialias=False,
-        ),
-        transforms.Normalize(
-          [0.48145466, 0.4578275, 0.40821073],
-          [0.26862954, 0.26130258, 0.27577711]),
-    ])
-    input = tform(input_im).to(device)
-
-    # approch B: convert input_im to torch
-    # inp = torch.from_numpy(np.array(input_im)).to(device)
-    # inp = torch.from_numpy(np.array(input_im)).permute(2, 0, 1).to(device)
-
-    # dtype = torch.float32
-    # input = input.to(device=device, dtype=dtype)
-    input = input.unsqueeze(0)
-    # image_embeddings = pipe.image_encoder(image).image_embeds
-    # image_embeddings = image_embeddings[0]
-
+    input_im = Image.fromarray(input_im)
+    prepro = preprocess(input_im).unsqueeze(0).to(device)
     with torch.no_grad():
-        # image_embeddings_np = model.get_text_features(prompt_tokens.to(device))
-        image_embeddings = model.get_image_features(input)
-
+        image_embeddings = model.encode_image(prepro)
     # image_embeddings /= image_embeddings.norm(dim=-1, keepdim=True)
-    image_embeddings_np = image_embeddings.cpu().detach().numpy()
+    image_embeddings_np = image_embeddings.cpu().to(torch.float32).detach().numpy()
     return image_embeddings_np
 
 def prompt_to_embedding(prompt):
-    # inputs = processor(prompt, images=imgs, return_tensors="pt", padding=True)
-    inputs = processor(prompt, return_tensors="pt", padding='max_length', max_length=77)
-    # labels = torch.tensor(labels)
-    # prompt_tokens = inputs.input_ids[0]
-    prompt_tokens = inputs.input_ids
-    # image = inputs.pixel_values
+    text = tokenizer([prompt]).to(device)
     with torch.no_grad():
-        prompt_embededdings = model.get_text_features(prompt_tokens.to(device))
+        prompt_embededdings = model.encode_text(text)
     # prompt_embededdings /= prompt_embededdings.norm(dim=-1, keepdim=True)
-    prompt_embededdings = prompt_embededdings[0].cpu().detach().numpy()
-    return prompt_embededdings
+    prompt_embededdings_np = prompt_embededdings.cpu().to(torch.float32).detach().numpy()    
+    return prompt_embededdings_np
 
 def embedding_to_image(embeddings):
     size = math.ceil(math.sqrt(embeddings.shape[0]))
@@ -87,15 +58,15 @@ def embedding_to_image(embeddings):
 
 def embedding_to_base64(embeddings):
     import base64
-    # ensure float16
-    embeddings = embeddings.astype(np.float16)
+    # ensure float32
+    embeddings = embeddings.astype(np.float32)
     embeddings_b64 = base64.urlsafe_b64encode(embeddings).decode()
     return embeddings_b64
 
 def base64_to_embedding(embeddings_b64):
     import base64
     embeddings = base64.urlsafe_b64decode(embeddings_b64)
-    embeddings = np.frombuffer(embeddings, dtype=np.float16)
+    embeddings = np.frombuffer(embeddings, dtype=np.float32)
     # embeddings = torch.tensor(embeddings)
     return embeddings
 
@@ -177,6 +148,9 @@ def update_average_embeddings(embedding_base64s_state, embedding_powers):
     # TODO toggle this to support average or sum
     final_embedding = final_embedding / num_embeddings
     
+    # normalize embeddings in numpy
+    final_embedding /= np.linalg.norm(final_embedding)
+    
     embeddings_b64 = embedding_to_base64(final_embedding)
     return embeddings_b64
 
@@ -229,35 +203,12 @@ def on_example_image_click_set_image(input_image, image_url):
 
 # device = torch.device("mps" if torch.backends.mps.is_available() else "cuda:0" if torch.cuda.is_available() else "cpu")
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
-torch_size = torch.float16 if device == ('cuda') else torch.float32
-# torch_size = torch.float32
-# pipe = StableDiffusionPipeline.from_pretrained(
-#     model_id,
-#     custom_pipeline="pipeline.py",
-#     torch_dtype=torch_size,
-#     # , revision="fp16",
-#     requires_safety_checker = False, safety_checker=None,
-#     text_encoder = CLIPTextModel,
-#     tokenizer = CLIPTokenizer,
-#     )
-# pipe = pipe.to(device)
-
-from transformers import AutoProcessor, AutoModel
-# processor = AutoProcessor.from_pretrained(clip_model_id)
-# model = AutoModel.from_pretrained(clip_model_id)
-# model = model.to(device)
 
 from clip_retrieval.load_clip import load_clip, get_tokenizer 
 # model, preprocess = load_clip(clip_model, use_jit=True, device=device)
 model, preprocess = load_clip(clip_model, use_jit=True, device=device)
 tokenizer = get_tokenizer(clip_model)
 
-test_url = "https://placekitten.com/400/600"
-test_caption = "an image of a cat"
-test_image_1 = "tests/test_clip_inference/test_images/123_456.jpg"
-test_image_2 = "tests/test_clip_inference/test_images/416_264.jpg"
-
-# clip_retrieval_service_url = "https://knn.laion.ai/knn-service"
 clip_retrieval_client = ClipClient(
     url=clip_retrieval_service_url, 
     indice_name=clip_model_id,