|
import gradio as gr |
|
from Models import VisionModel |
|
import huggingface_hub |
|
from PIL import Image |
|
import torch.amp.autocast_mode |
|
from pathlib import Path |
|
|
|
|
|
MODEL_REPO = "fancyfeast/joytag" |
|
|
|
|
|
@torch.no_grad() |
|
def predict(image: Image.Image): |
|
with torch.amp.autocast_mode.autocast('cuda', enabled=True): |
|
preds = model(image) |
|
tag_preds = preds['tags'].sigmoid().cpu() |
|
|
|
return {top_tags[i]: tag_preds[i] for i in range(len(top_tags))} |
|
|
|
|
|
print("Downloading model...") |
|
path = huggingface_hub.snapshot_download(MODEL_REPO) |
|
print("Loading model...") |
|
model = VisionModel.load_model(path) |
|
model.eval() |
|
|
|
with open(Path(path) / 'top_tags.txt', 'r') as f: |
|
top_tags = [line.strip() for line in f.readlines() if line.strip()] |
|
|
|
print("Starting server...") |
|
|
|
gradio_app = gr.Interface( |
|
predict, |
|
inputs=gr.Image(label="Source", sources=['upload', 'webcam'], type='pil'), |
|
outputs=[gr.Label(label="Result", num_top_classes=5)], |
|
title="JoyTag", |
|
) |
|
|
|
|
|
if __name__ == '__main__': |
|
gradio_app.launch() |
|
|
