Update app.py

app.py

@@ -7,7 +7,7 @@ import json
 from torch import nn
 from huggingface_hub import hf_hub_download
 
-# Định nghĩa mô hình (giống lúc train)
+# Định nghĩa mô hình
 class VQAModel(nn.Module):
     def __init__(self, num_answers):
         super(VQAModel, self).__init__()
@@ -29,48 +29,59 @@ class VQAModel(nn.Module):
 repo_id = "duyan2803/vqa-model-vit-bert"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load config
-config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
-with open(config_path, "r") as f:
-    config = json.load(f)
-num_answers = config["num_answers"]
+try:
+    # Load config
+    config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
+    with open(config_path, "r") as f:
+        config = json.load(f)
+    num_answers = config["num_answers"]
 
-# Load weights
-weights_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin")
-model = VQAModel(num_answers=num_answers)
-model.load_state_dict(torch.load(weights_path, map_location=device, weights_only=True))
-model.to(device)
-model.eval()
+    # Load weights
+    weights_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin")
+    model = VQAModel(num_answers=num_answers)
+    state_dict = torch.load(weights_path, map_location=device, weights_only=True)
+    model.load_state_dict(state_dict)
+    model.to(device)
+    model.eval()
+    print("Đã load mô hình thành công!")
 
-# Load tokenizer
-tokenizer = BertTokenizer.from_pretrained(repo_id)
+    # Load tokenizer
+    tokenizer = BertTokenizer.from_pretrained(repo_id)
+    print("Đã load tokenizer thành công!")
 
-# Load answer list
-answer_list_path = hf_hub_download(repo_id=repo_id, filename="answer_list.json")
-with open(answer_list_path, "r") as f:
-    answer_list = json.load(f)
+    # Load answer list
+    answer_list_path = hf_hub_download(repo_id=repo_id, filename="answer_list.json")
+    with open(answer_list_path, "r") as f:
+        answer_list = json.load(f)
+    print("Đã load answer list thành công!")
+except Exception as e:
+    print(f"Lỗi khi load mô hình hoặc file: {str(e)}")
+    raise e
 
 # Hàm dự đoán
 def predict(image, question):
-    # Xử lý ảnh
-    transform = transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-    ])
-    image_tensor = transform(image).unsqueeze(0).to(device)
+    try:
+        # Xử lý ảnh
+        transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ])
+        image_tensor = transform(image).unsqueeze(0).to(device)
 
-    # Xử lý câu hỏi
-    tokenized = tokenizer(question, padding='max_length', truncation=True, max_length=32, return_tensors='pt')
-    input_ids = tokenized['input_ids'].to(device)
-    attention_mask = tokenized['attention_mask'].to(device)
+        # Xử lý câu hỏi
+        tokenized = tokenizer(question, padding='max_length', truncation=True, max_length=32, return_tensors='pt')
+        input_ids = tokenized['input_ids'].to(device)
+        attention_mask = tokenized['attention_mask'].to(device)
 
-    # Dự đoán
-    with torch.no_grad():
-        output = model(image_tensor, input_ids, attention_mask)
-        pred_idx = output.argmax(dim=1).item()
-    
-    return answer_list[pred_idx]
+        # Dự đoán
+        with torch.no_grad():
+            output = model(image_tensor, input_ids, attention_mask)
+            pred_idx = output.argmax(dim=1).item()
+        
+        return answer_list[pred_idx]
+    except Exception as e:
+        return f"Lỗi khi dự đoán: {str(e)}"
 
 # Giao diện Gradio
 interface = gr.Interface(

requirements.txt

@@ -1,6 +1,6 @@
-torch
-transformers
-torchvision
+torch==2.0.1
+transformers>=4.32.0
+torchvision==0.15.2
 pillow
-gradio
-huggingface_hub
+gradio==4.0.2
+huggingface_hub>=0.29.0