Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 os.system("gdown https://drive.google.com/uc?id=14pXWwB4Zm82rsDdvbGguLfx9F8aM7ovT")
+os.system("gdown https://drive.google.com/uc?id=1IdaBtMSvtyzF0ByVaBHtvM0JYSXRExRX")
 import clip
 import os
 from torch import nn
@@ -32,7 +33,6 @@ TA = Union[T, ARRAY]
 D = torch.device
 CPU = torch.device('cpu')
 
-model_path = 'conceptual_weights.pt'
 
 def get_device(device_id: int) -> D:
     if not torch.cuda.is_available():
@@ -234,14 +234,16 @@ prefix_length = 10
 
 model = ClipCaptionModel(prefix_length)
 
-model.load_state_dict(torch.load(model_path, map_location=CPU)) 
-
-model = model.eval() 
-device = CUDA(0) if is_gpu else "cpu"
 model = model.to(device)
 
 
-def inference(img):
+def inference(img,model):
+  if model == "COCO":
+      model_path = 'coco_weights.pt'
+  else:
+      model_path = 'conceptual_weights.pt'
+  model.load_state_dict(torch.load(model_path, map_location=CPU)) 
+
   use_beam_search = False 
   image = io.imread(img.name)
   pil_image = PIL.Image.fromarray(image)  
@@ -262,7 +264,7 @@ article = "<p style='text-align: center'><a href='https://github.com/rmokady/CLI
 examples=[['water.jpeg']]
 gr.Interface(
     inference, 
-    gr.inputs.Image(type="file", label="Input"), 
+    [gr.inputs.Image(type="file", label="Input"),gr.inputs.Radio(choices["COCO","Conceptual captions"], type="value", default="COCO", label="Model")], 
     gr.outputs.Textbox(label="Output"),
     title=title,
     description=description,