Fix bugs

handler.py

@@ -24,10 +24,10 @@ def preprocess_text(text):
     """Tokenize text query."""
     return tokenize([str(text)])[0].unsqueeze(0).to(device)
 
-def get_fused_embedding(sketch_base64, text, model):
+def get_fused_embedding(sketch_base64, text, model, transformer):
     """Fuse sketch and text features into a single embedding."""
     with torch.no_grad():
-        sketch_tensor = preprocess_image(sketch_base64)
+        sketch_tensor = preprocess_image(sketch_base64, transformer)
         text_tensor = preprocess_text(text)
 
         sketch_feature = model.encode_sketch(sketch_tensor)
@@ -39,9 +39,9 @@ def get_fused_embedding(sketch_base64, text, model):
         fused_embedding = model.feature_fuse(sketch_feature, text_feature)
     return fused_embedding.cpu().numpy().tolist()
 
-def get_image_embedding(image_base64, model):
+def get_image_embedding(image_base64, model, transformer):
     """Convert base64 encoded image to tensor."""
-    image_tensor = preprocess_image(image_base64)
+    image_tensor = preprocess_image(image_base64, transformer)
     with torch.no_grad():
         image_feature = model.encode_image(image_tensor)
         image_feature = image_feature / image_feature.norm(dim=-1, keepdim=True)
@@ -89,13 +89,13 @@ class EndpointHandler:
                 return {"error": "Both 'sketch' (base64) and 'text' are required inputs."}
 
             # Generate Fused Embedding
-            fused_embedding = get_fused_embedding(sketch_base64, text_query)
+            fused_embedding = get_fused_embedding(sketch_base64, text_query, self.model, self.transform)
             return {"embedding": fused_embedding}
         elif "image" in inputs:
             image_base64 = inputs.get("image", "")
             if not image_base64:
                 return {"error": "Image 'image' (base64) is required input."}
-            embedding = get_image_embedding(image_base64)
+            embedding = get_image_embedding(image_base64, self.model, self.transform)
             return {"embedding": embedding}
         else:
             return {"error": "Input 'sketch' or 'image' is required."}