Update app.py

app.py

@@ -224,25 +224,24 @@ def generate2(
     return generated_list[0]
     
 
+device = CUDA(0) if is_gpu else "cpu"
+clip_model, preprocess = clip.load("ViT-B/32", device=device, jit=False)
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 
-
-def inference(img,model):
+def inference(img,model_name):
   is_gpu = False 
     
-  device = CUDA(0) if is_gpu else "cpu"
-  clip_model, preprocess = clip.load("ViT-B/32", device=device, jit=False)
-  tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-    
-    
   prefix_length = 10
     
   model = ClipCaptionModel(prefix_length)
     
-  if model == "COCO":
+  if model_name == "COCO":
       model_path = 'coco_weights.pt'
   else:
       model_path = 'conceptual_weights.pt'
   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)
 
   use_beam_search = False