update inference script for comma separated tags and using cuda if available

inference_gradio.py

@@ -1,18 +1,15 @@
 import json
 
-from PIL import Image
 import gradio as gr
+from PIL import Image
+import safetensors.torch
+import timm
+from timm.models import VisionTransformer
 import torch
 from torchvision.transforms import transforms
 from torchvision.transforms import InterpolationMode
 import torchvision.transforms.functional as TF
 
-import timm
-from timm.models import VisionTransformer
-import safetensors.torch
-
-
-torch.jit.script = lambda f: f
 torch.set_grad_enabled(False)
 
 class Fit(torch.nn.Module):
@@ -123,13 +120,26 @@ model = timm.create_model(
 safetensors.torch.load_model(model, "JTP_PILOT/JTP_PILOT-e4-vit_so400m_patch14_siglip_384.safetensors")
 model.eval()
 
+if torch.cuda.is_available():
+    model.cuda()
+    if torch.cuda.get_device_capability()[0] >= 7: # tensor cores
+        model.half()
+
 with open("JTP_PILOT/tags.json", "r") as file:
     tags = json.load(file) # type: dict
 allowed_tags = list(tags.keys())
 
+for idx, tag in enumerate(allowed_tags):
+    allowed_tags[idx] = tag.replace("_", " ")
+
 def create_tags(image, threshold):
     img = image.convert('RGB')
-    tensor = transform(img).unsqueeze(0)
+    tensor = transform(img).unsqueeze(0) # type: torch.Tensor
+
+    if torch.cuda.is_available():
+        tensor.cuda()
+        if torch.cuda.get_device_capability()[0] >= 7:
+            tensor.half()
 
     with torch.no_grad():
         logits = model(tensor)