Update app.py

app.py

@@ -45,32 +45,34 @@ resblock.load_state_dict(torch.load('./model_chkpt/step2_resblock.pth',map_locat
 def model_generate_ans(img,val_q):
 
     max_generate_length = 30
-    
-    # image
-    image_processed  = processor(images=img, return_tensors="pt").to(device)
-    clip_val_outputs = clip_model(**image_processed).last_hidden_state[:,1:,:]
-    val_image_embeds = projection(clip_val_outputs)
-    val_image_embeds = resblock(val_image_embeds).to(torch.float16)
-
-    img_token_tensor = torch.tensor(IMAGE_TOKEN_ID).to(device)
-    img_token_embeds = merged_model.model.embed_tokens(img_token_tensor).unsqueeze(0).unsqueeze(0)
 
-    val_q_tokenised = tokenizer(val_q, return_tensors="pt", return_attention_mask=False)['input_ids'].squeeze(0)
-    val_q_embeds    = merged_model.model.embed_tokens(val_q_tokenised).unsqueeze(0)
+    with torch.no_grad():
     
-    val_combined_embeds = torch.cat([val_image_embeds, img_token_embeds, val_q_embeds], dim=1) # 4, 69, 2560
+        # image
+        image_processed  = processor(images=img, return_tensors="pt").to(device)
+        clip_val_outputs = clip_model(**image_processed).last_hidden_state[:,1:,:]
+        val_image_embeds = projection(clip_val_outputs)
+        val_image_embeds = resblock(val_image_embeds).to(torch.float16)
     
-    predicted_caption = torch.full((1,max_generate_length),50256)
-
-    for g in range(max_generate_length):
-        phi_output_logits = merged_model(inputs_embeds=val_combined_embeds)['logits'] # 4, 69, 51200
-        predicted_word_token_logits = phi_output_logits[:, -1, :].unsqueeze(1) # 4,1,51200
-        predicted_word_token = torch.argmax(predicted_word_token_logits, dim = -1) # 4,1
-        predicted_caption[:,g] = predicted_word_token.view(1,-1).to(device)
-        next_token_embeds = phi_model.model.embed_tokens(predicted_word_token) # 4,1,2560
-        val_combined_embeds   = torch.cat([val_combined_embeds, next_token_embeds], dim=1)
+        img_token_tensor = torch.tensor(IMAGE_TOKEN_ID).to(device)
+        img_token_embeds = merged_model.model.embed_tokens(img_token_tensor).unsqueeze(0).unsqueeze(0)
+    
+        val_q_tokenised = tokenizer(val_q, return_tensors="pt", return_attention_mask=False)['input_ids'].squeeze(0)
+        val_q_embeds    = merged_model.model.embed_tokens(val_q_tokenised).unsqueeze(0)
+        
+        val_combined_embeds = torch.cat([val_image_embeds, img_token_embeds, val_q_embeds], dim=1) # 4, 69, 2560
         
-    predicted_captions_decoded = tokenizer.batch_decode(predicted_caption,ignore_index = 50256)
+        predicted_caption = torch.full((1,max_generate_length),50256)
+    
+        for g in range(max_generate_length):
+            phi_output_logits = merged_model(inputs_embeds=val_combined_embeds)['logits'] # 4, 69, 51200
+            predicted_word_token_logits = phi_output_logits[:, -1, :].unsqueeze(1) # 4,1,51200
+            predicted_word_token = torch.argmax(predicted_word_token_logits, dim = -1) # 4,1
+            predicted_caption[:,g] = predicted_word_token.view(1,-1)
+            next_token_embeds = phi_model.model.embed_tokens(predicted_word_token) # 4,1,2560
+            val_combined_embeds   = torch.cat([val_combined_embeds, next_token_embeds], dim=1)
+            
+        predicted_captions_decoded = tokenizer.batch_decode(predicted_caption,ignore_index = 50256)
     
     return predicted_captions_decoded