Update app.py

app.py

@@ -24,10 +24,6 @@ def process(model,image, tokenizer, device):
     model.eval()
     with torch.no_grad():
         encoder_input = image.unsqueeze(0).to(device) # (b, seq_len)
-        # decoder_input = batch['decoder_input'].to(device) # (B, seq_len)
-        # encoder_mask = batch['encoder_mask'].to(device) # (B, 1, 1, seq_len)
-        # decoder_mask = batch['decoder_mask'].to(device) # (B, 1, seq_len, seq_len)
-
         model_out = greedy_decode(model, encoder_input, None, tokenizer, 196,device)
         model_text  = tokenizer.decode(model_out.detach().cpu().numpy())
         return model_text
@@ -111,13 +107,6 @@ def greedy_decode(model, source, source_mask, tokenizer_tgt, max_len, device):
         
         # Append next word
         decoder_input = torch.cat([decoder_input, next_word.unsqueeze(0)], dim=1)
-        # # get next token
-        # prob = model.project(out[:, -1])
-        # _, next_word = torch.max(prob, dim=1)
-        # # print(f'prob: {prob.shape}')
-        # decoder_input = torch.cat(
-        #     [decoder_input, torch.empty(1, 1).long().fill_(next_word.item()).to(device)], dim=1
-        # )
 
         if next_word.item() == eos_idx:
             break
@@ -127,7 +116,7 @@ def greedy_decode(model, source, source_mask, tokenizer_tgt, max_len, device):
 def image_base64(image):
     
 
-    # with open('C:/AI/projects/vision_model_pretrained/validation/content/memory_image_23330.jpg', 'rb') as image_file:
+ 
     base64_bytes = base64.b64encode(image_file.read())
        
 
@@ -135,27 +124,11 @@ def image_base64(image):
     return base64_string
 
 
-def start():
-    print('start')
-    accelerator = Accelerator()
-    device = accelerator.device
-
-    config = get_config()
-    tokenizer = get_or_build_tokenizer(config)
-    model = get_model(config, len(tokenizer))
-    model = accelerator.prepare(model)
-    accelerator.load_state('model.tensors')
-
-    image = image_base64()
-
-    
 
-    process(model, image, tokenizer, device)
 
-# start()
 
 def main():
-    st.title("Image Captioning with Transformer Models")
+    st.title("Image Captioning with Vision Transformer")
     image = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
     
     if image is not None:
@@ -173,8 +146,7 @@ def main():
             model = get_model(config, len(tokenizer))
             model = accelerator.prepare(model)
             accelerator.load_state('models/')
-            # model = get_model(config, len(tokenizer))
-            # model.to(device)
+         
             
            
             text_output = process(model, image, tokenizer, device)