Update handler.py

handler.py

@@ -140,14 +140,13 @@ class EndpointHandler:
         self.labels: LabelData = load_labels_hf(repo_id=repo_id)
 
         print("Creating data transform...")
-        self.transform = create_transform(**resolve_data_config(model.pretrained_cfg, model=model))
+        self.transform = create_transform(**resolve_data_config(self.model.pretrained_cfg, model=self.model))
 
-        with torch.inference_mode():
-            # move model to GPU, if available
-            if torch_device.type != "cpu":
-                self.model = self.model.to(torch_device)
+        # move model to GPU, if available
+        if torch_device.type != "cpu":
+            self.model = self.model.to(torch_device)
 
-        uri = os.environ.get("MongoDB", "mongodb+srv://jamie:qJiuKQpqhXMHGb74@cluster0.i5ujz.mongodb.net/")
+        uri = os.environ.get("MongoDB", "")
         self.client = MongoClient(uri)
 
         self.db = self.client['nomorecopyright']
@@ -184,16 +183,18 @@ class EndpointHandler:
                 inputs: Tensor = self.transform(img_input).unsqueeze(0)
                 # NCHW image RGB to BGR
                 inputs = inputs[:, [2, 1, 0]]
-                inputs = inputs.to(torch_device)
-                print("Running inference...")
-                outputs = self.model.forward(inputs)
-                # apply the final activation function (timm doesn't support doing this internally)
-                outputs = F.sigmoid(outputs)
-                # move inputs, outputs, and model back to to cpu if we were on GPU
-                if torch_device.type != "cpu":
-                    inputs = inputs.to("cpu")
-                    outputs = outputs.to("cpu")
-
+                with torch.inference_mode():
+                    # move model to GPU, if available
+                    if torch_device.type != "cpu":
+                        inputs = inputs.to(torch_device)
+                    print("Running inference...")
+                    outputs = self.model.forward(inputs)
+                    # apply the final activation function (timm doesn't support doing this internally)
+                    outputs = F.sigmoid(outputs)
+                    # move inputs, outputs, and model back to to cpu if we were on GPU
+                    if torch_device.type != "cpu":
+                        inputs = inputs.to("cpu")
+                        outputs = outputs.to("cpu")
                 print("Processing results...")
                 caption, taglist, ratings, character, general = get_tags(
                     probs=outputs.squeeze(0),
@@ -203,6 +204,7 @@ class EndpointHandler:
                 )
 
                 results={**ratings, **character, **general}
+                results={key: float(value) for key, value in results.items()}
                 print(results)
 
                 saveQuery = {"_id": document.get('_id')}