try init model

app.py

@@ -44,7 +44,7 @@ st.sidebar.write('\n')
 
 with st.spinner('Generating image caption ...'):
 
-    caption, tokens, token_ids = predict(image)
+    caption, tokens, token_ids = predict_dummy(image)
 
     st.success(f'caption: {caption}')
     st.success(f'tokens: {tokens}')

model.py

@@ -25,9 +25,46 @@ shutil.copyfile(filepath, os.path.join(model_dir, 'flax_model.msgpack'))
 
 flax_vit_gpt2_lm = FlaxViTGPT2LMForConditionalGeneration.from_pretrained(model_dir)
 
+vit_model_name = 'google/vit-base-patch16-224-in21k'
+feature_extractor = ViTFeatureExtractor.from_pretrained(vit_model_name)
+
+gpt2_model_name = 'asi/gpt-fr-cased-small'
+tokenizer = GPT2Tokenizer.from_pretrained(gpt2_model_name)
+
+max_length = 32
+num_beams = 8
+gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
+
+
+@jax.jit
+def predict_fn(pixel_values):
+
+    return flax_vit_gpt2_lm.generate(pixel_values, **gen_kwargs)
+
 def predict(image):
-    return 'dummy caption!', ['dummy', 'caption', '!'], [1, 2, 3]
 
+    # batch dim is added automatically
+    encoder_inputs = feature_extractor(images=image, return_tensors="jax")
+    pixel_values = encoder_inputs.pixel_values
 
+    # generation
+    generation = predict_fn(pixel_values)
 
+    token_ids = np.array(generation.sequences)[0]
+    caption = tokenizer.decode(token_ids)
+
+    return caption, token_ids
+
+def init():
+
+    image_path = 'samples/val_000000039769.jpg'
+    image = Image.open(image_path)
+
+    caption, token_ids = predict(image)
+    image.close()
+
+def predict_dummy(image):
+    
+    return 'dummy caption!', ['dummy', 'caption', '!'], [1, 2, 3]
 
+init()