Update app.py

app.py

@@ -11,17 +11,19 @@ model, _, transform = open_clip.create_model_and_transforms(
     pretrained="laion2b_s13b_b90k"
 )
 model.to(device)
+model.eval()
+traced_model = torch.jit.trace(model, torch.zeros((1, 3, 64, 64)).to(device))
 
 def output_generate(image):
     im = transform(image).unsqueeze(0).to(device)
     with torch.no_grad(), torch.cuda.amp.autocast():
-        generated = model.generate(im, seq_len=20)
+        generated = traced_model.generate(im, seq_len=20)
     return open_clip.decode(generated[0].detach()).split("<end_of_text>")[0].replace("<start_of_text>", "")
 
 def inference_caption(image):
     im = transform(image).unsqueeze(0).to(device)
     with torch.no_grad(), torch.cuda.amp.autocast():
-        generated = model.generate(
+        generated = traced_model.generate(
             im, 
             generation_type="beam_search",
             top_p=1.0, 
@@ -31,6 +33,8 @@ def inference_caption(image):
         )
     return open_clip.decode(generated[0].detach()).split("<end_of_text>")[0].replace("<start_of_text>", "")
 
+transform = open_clip.get_transforms("coca_ViT-B-32", image_size=128)
+
 image_input = gr.inputs.Image(type="pil")
 caption_output = gr.outputs.Textbox(label="Caption Output")
 caption_interface = gr.Interface(fn=inference_caption, inputs=image_input, outputs=caption_output, capture_session=True, title="CoCa: Contrastive Captioners", description="An open source implementation of CoCa: Contrastive Captioners are Image-Text Foundation Models.", examples=[path.as_posix() for path in sorted(pathlib.Path("images").glob("*.jpg"))], allow_flagging=False)