Update app.py

app.py

@@ -81,7 +81,6 @@ class Predictor:
         self.text_emb_model = AutoModel.from_pretrained(model_path, trust_remote_code=True)
         self.text_emb_model = self.text_emb_model.eval()
 
-    @torch.inference_mode()
     def embed_text(self, input_strings):
         with torch.no_grad():
             # Tokenize the input texts
@@ -118,7 +117,7 @@ class Predictor:
         image_features = self.cls_model[0].forward_features(image.unsqueeze(0))
         outputs = self.cls_model[0].head(image_features, q = query).sigmoid().float()
 
-        general_tag_list = list(zip(self.tag_names, outputs[0].tolist()))
+        general_tag_list = list(zip(tag_names, outputs[0].tolist()))
         general_tag_list.sort(key=lambda y: y[1], reverse=True)
         general_tag_preds_dict = {}
         for tag, prob in general_tag_list[:50]:
@@ -138,7 +137,7 @@ class Predictor:
         image,
         description,
     ):
-        return self.predict(image, self.embed_text(description), ["embedding"])["embedding"]
+        return self.predict(image, self.embed_text([description]), ["embedding"])["embedding"]
 
 
 def main():