Update app.py

app.py

@@ -1,9 +1,9 @@
 import json
+from pprint import pprint
 
-import gradio
 import torch
 import torch.hub
-from gradio import inputs, outputs
+from gradio import Interface, inputs, outputs
 from PIL import Image
 from torchvision import transforms
 
@@ -20,6 +20,7 @@ torch.hub.load_state_dict_from_url = load_state_dict_from_url
 model = torch.hub.load(
     "RF5/danbooru-pretrained",
     "resnet50",
+    map_location="cpu",
 )
 model.eval()
 
@@ -52,6 +53,7 @@ def main(input_image: Image.Image, threshold: float):
     tag_confidences = {}
     for index in inds[0 : len(results)]:
         tag_confidences[tags[index]] = float(probs[index].cpu().numpy())
+    pprint(tag_confidences)
     return tag_confidences
 
 
@@ -62,5 +64,5 @@ threshold = inputs.Slider(
 
 labels = outputs.Label(type="confidence")
 
-interface = gradio.Interface(main, inputs=[image, threshold], outputs=[labels])
+interface = Interface(main, inputs=[image, threshold], outputs=[labels])
 interface.launch()