Update app.py

app.py

@@ -5,20 +5,23 @@ import torch
 import kelip
 import gradio as gr
 
+
 def load_model():
-    model, preprocess_img, tokenizer = kelip.build_model('ViT-B/32')
+    model, preprocess_img, tokenizer = kelip.build_model("ViT-B/32")
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model = model.to(device)
     model.eval()
 
-    model_dict = {'model': model,
-                  'preprocess_img': preprocess_img,
-                  'tokenizer': tokenizer
-                  }
+    model_dict = {
+        "model": model,
+        "preprocess_img": preprocess_img,
+        "tokenizer": tokenizer,
+    }
     return model_dict
 
+
 def classify(img, user_text):
-    preprocess_img = model_dict['preprocess_img']
+    preprocess_img = model_dict["preprocess_img"]
 
     input_img = preprocess_img(img).unsqueeze(0)
 
@@ -27,17 +30,17 @@ def classify(img, user_text):
 
     # extract image features
     with torch.no_grad():
-        image_features = model_dict['model'].encode_image(input_img)
+        image_features = model_dict["model"].encode_image(input_img)
 
         # extract text features
-        user_texts = user_text.split(',')
-        if user_text == '' or user_text.isspace():
+        user_texts = user_text.split(",")
+        if user_text == "" or user_text.isspace():
             user_texts = []
 
-        input_texts = model_dict['tokenizer'].encode(user_texts)
+        input_texts = model_dict["tokenizer"].encode(user_texts)
         if torch.cuda.is_available():
             input_texts = input_texts.cuda()
-        text_features = model_dict['model'].encode_text(input_texts)
+        text_features = model_dict["model"].encode_text(input_texts)
 
     # l2 normalize
     image_features /= image_features.norm(dim=-1, keepdim=True)
@@ -50,28 +53,30 @@ def classify(img, user_text):
 
     return result
 
-if __name__ == '__main__':
+
+if __name__ == "__main__":
     global model_dict
 
     model_dict = load_model()
 
-    inputs = [gr.inputs.Image(type="pil", label="Image"),
-              gr.inputs.Textbox(lines=5, label="Caption"),
-              ]
+    inputs = [
+        gr.inputs.Image(type="pil", label="Image"),
+        gr.inputs.Textbox(lines=5, label="Caption"),
+    ]
 
-    outputs = ['label']
+    outputs = ["label"]
 
     title = "KELIP"
     description = "Zero-shot classification with KELIP -- Korean and English bilingual contrastive Language-Image Pre-training model that is trained with collected 1.1 billion image-text pairs (708 million Korean and 476 million English).<br> <br><a href='https://arxiv.org/abs/2203.14463' target='_blank'>Arxiv</a> | <a href='https://github.com/navervision/KELIP' target='_blank'>Github</a>"
-    
+
     article = ""
 
-    iface=gr.Interface(
+    iface = gr.Interface(
         fn=classify,
         inputs=inputs,
         outputs=outputs,
         title=title,
         description=description,
-        article=article
+        article=article,
     )
-    iface.launch()
+    iface.launch()