Update pages/ImageToText.py

pages/ImageToText.py

@@ -11,7 +11,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 vitgpt_model.to(device)
 
-def generate_caption(processor, model, image, tokenizer=None, num_seq):
+def generate_caption(processor, model, image, num_seq, tokenizer=None):
     
     inputs = processor(images=image, return_tensors="pt").to(device)
     generated_ids = model.generate(pixel_values=inputs.pixel_values, 
@@ -30,7 +30,7 @@ def generate_caption(processor, model, image, tokenizer=None, num_seq):
     return generated_caption
 
 def generate_captions(image, num_seq):
-    caption_vitgpt = generate_caption(vitgpt_processor, vitgpt_model, image, vitgpt_tokenizer, num_seq)
+    caption_vitgpt = generate_caption(vitgpt_processor, vitgpt_model, image, num_seq, vitgpt_tokenizer)
     return caption_vitgpt
 
 st.title('Generate text to your image')
@@ -44,6 +44,6 @@ if uploaded_file is not None:
             image = Image.open(uploaded_file)
             st.image(image)
         with col2:
-            generated_caption = generate_caption(vitgpt_processor, vitgpt_model, image, vitgpt_tokenizer)
+            generated_caption = generate_caption(vitgpt_processor, vitgpt_model, image, num_seq, vitgpt_tokenizer)
             for i in generated_caption:
                 st.write(i)