Update app.py

app.py

@@ -40,23 +40,22 @@ class Attention(nn.Module):
 # -----------------------
 # VQA Model
 # -----------------------
-# (Copy nguyên văn từ code gốc bạn cung cấp)
 class VQAModel(nn.Module):
-    # !! QUAN TRỌNG: Đảm bảo các giá trị mặc định này (hoặc giá trị bạn truyền vào khi load)
-    #    KHỚP VỚI CÁCH BẠN KHỞI TẠO MODEL KHI LƯU FILE .pth !!
     def __init__(self, vocab_size, embedding_dim=256, lstm_units=256, cnn_output_dim=512, attention_dim=256, max_seq_len=30):
         super(VQAModel, self).__init__()
         self.vocab_size = vocab_size
-        self.max_seq_len = max_seq_len # Lưu lại để dùng trong predict
+        self.max_seq_len = max_seq_len
 
-        # CNN Encoder (giống hệt lúc train)
-        self.cnn_net = nn.Sequential( # Đổi tên từ self.cnn thành self.cnn_net để tránh trùng tên biến local trong forward
+        # --- CNN Encoder: ĐỔI TÊN TRỞ LẠI THÀNH self.cnn ---
+        self.cnn = nn.Sequential( # Đổi tên lại thành self.cnn
             nn.Conv2d(3, 32, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
             nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
             nn.Conv2d(64, 128, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
             nn.Conv2d(128, cnn_output_dim, kernel_size=3, padding=1), nn.ReLU(),
             nn.AdaptiveAvgPool2d((1, 1))
         )
+        # ------------------------------------------------
+
         self.embedding = nn.Embedding(vocab_size, embedding_dim)
         self.caption_lstm = nn.LSTM(embedding_dim, lstm_units, batch_first=True)
         self.question_lstm = nn.LSTM(embedding_dim, lstm_units, batch_first=True)
@@ -64,41 +63,31 @@ class VQAModel(nn.Module):
         self.decoder_input_proj = nn.Linear(embedding_dim + 3 * lstm_units, lstm_units)
         self.decoder_lstm = nn.LSTM(lstm_units, lstm_units, batch_first=True)
         self.fc_out = nn.Linear(lstm_units, vocab_size)
-        self.dropout = nn.Dropout(0.5) # Tự động tắt khi model.eval()
+        self.dropout = nn.Dropout(0.5)
 
-    # Hàm forward không dùng trực tiếp khi predict, nhưng cần tồn tại
+    # Hàm forward không bị ảnh hưởng vì không gọi trực tiếp
     def forward(self, image, caption, question, answer_input):
-         # --- CNN Encoder ---
-        # Sử dụng self.cnn_net để gọi Sequential
-        cnn_features = self.cnn_net(image)  # (batch, cnn_output_dim, 1, 1)
-        cnn_features = cnn_features.view(cnn_features.size(0), -1)  # (batch, cnn_output_dim)
-
-        # --- Text Encoders ---
-        cap_embed = self.embedding(caption)  # (batch, cap_seq_len, embedding_dim)
-        cap_output, _ = self.caption_lstm(cap_embed)  # (batch, cap_seq_len, lstm_units)
-
-        q_embed = self.embedding(question)  # (batch, q_seq_len, embedding_dim)
-        q_output, _ = self.question_lstm(q_embed)  # (batch, q_seq_len, lstm_units)
-
-        # --- Attention ---
-        # Chắc chắn self.attention được gọi đúng
-        cap_attended = self.attention(cnn_features.unsqueeze(1), cap_output)  # (batch, lstm_units)
-        q_attended = self.attention(cnn_features.unsqueeze(1), q_output)      # (batch, lstm_units)
-
-        q_last = q_output[:, -1, :]  # (batch, lstm_units)
-
-        # Context vector: (batch, 3*lstm_units)
+         # Logic forward có thể vẫn dùng tên biến local cnn_features
+         # nhưng self.cnn để gọi mạng Sequential thì đã khớp tên
+        cnn_features = self.cnn(image) # Gọi self.cnn mới đúng tên
+        # ... (phần còn lại của forward giữ nguyên) ...
+        cnn_features = cnn_features.view(cnn_features.size(0), -1)
+        cap_embed = self.embedding(caption)
+        cap_output, _ = self.caption_lstm(cap_embed)
+        q_embed = self.embedding(question)
+        q_output, _ = self.question_lstm(q_embed)
+        cap_attended = self.attention(cnn_features.unsqueeze(1), cap_output)
+        q_attended = self.attention(cnn_features.unsqueeze(1), q_output)
+        q_last = q_output[:, -1, :]
         context = torch.cat([cap_attended, q_attended, q_last], dim=-1)
-
-        # --- Decoder với Teacher Forcing ---
-        answer_embed = self.embedding(answer_input)  # (batch, ans_seq_len, embedding_dim)
-        context_repeated = context.unsqueeze(1).repeat(1, answer_input.size(1), 1)  # (batch, ans_seq_len, 3*lstm_units)
-        decoder_in = torch.cat([answer_embed, context_repeated], dim=-1)  # (batch, ans_seq_len, embedding_dim + 3*lstm_units)
-        decoder_in = self.decoder_input_proj(decoder_in)  # (batch, ans_seq_len, lstm_units)
-
-        decoder_output, _ = self.decoder_lstm(decoder_in)  # (batch, ans_seq_len, lstm_units)
-        output = self.fc_out(self.dropout(decoder_output))  # (batch, ans_seq_len, vocab_size)
+        answer_embed = self.embedding(answer_input)
+        context_repeated = context.unsqueeze(1).repeat(1, answer_input.size(1), 1)
+        decoder_in = torch.cat([answer_embed, context_repeated], dim=-1)
+        decoder_in = self.decoder_input_proj(decoder_in)
+        decoder_output, _ = self.decoder_lstm(decoder_in)
+        output = self.fc_out(self.dropout(decoder_output))
         return output
+
 # ----------------------------------------------------------------------------
 
 # ============================================================================
@@ -218,8 +207,9 @@ def predict_vqa(image, question_str):
         with torch.no_grad(): # Tắt gradient calculation
             print("Encoding image...")
             # Sử dụng self.cnn_net thay vì self.cnn
-            cnn_features = model.cnn_net(image_tensor)
-            cnn_features = cnn_features.view(cnn_features.size(0), -1) # (1, cnn_output_dim)
+
+            cnn_features = model.cnn(image_tensor)
+            cnn_features = cnn_features.view(cnn_features.size(0), -1)
             print(f"CNN features shape: {cnn_features.shape}")
 
             print("Encoding question...")