Update app.py

app.py

@@ -5,13 +5,30 @@ from PIL import Image
 from torchvision import transforms
 import json
 import os
+from huggingface_hub import hf_hub_download
 
+# ---------------------
+# Config
+# ---------------------
 IMG_SIZE = 224
 SEQ_LEN = 32
 VOCAB_SIZE = 75460
-
+REPO_ID = "hackergeek/RADIOCAP13"             # your HF repo
+WEIGHTS_FILENAME = "pytorch_model.bin"
+VOCAB_FILENAME   = "vocab.json"
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# ---------------------
+# Download model files (if not present)
+# ---------------------
+# Download weights
+weights_path = hf_hub_download(repo_id=REPO_ID, filename=WEIGHTS_FILENAME)
+# Download vocab
+vocab_path   = hf_hub_download(repo_id=REPO_ID, filename=VOCAB_FILENAME)
+
+# ---------------------
+# Preprocessing & Tokenizer
+# ---------------------
 transform = transforms.Compose([
     transforms.Resize((IMG_SIZE, IMG_SIZE)),
     transforms.ToTensor(),
@@ -33,68 +50,79 @@ class SimpleTokenizer:
 
     @classmethod
     def load(cls, path):
-        with open(f"{path}/vocab.json", "r") as f:
+        with open(path, "r") as f:
             word2idx = json.load(f)
         return cls(word2idx)
 
+# ---------------------
+# Decoder
+# ---------------------
 class BiasDecoder(torch.nn.Module):
     def __init__(self, feature_dim=768, vocab_size=VOCAB_SIZE):
         super().__init__()
-        self.token_emb = torch.nn.Embedding(vocab_size, feature_dim)
-        self.pos_emb = torch.nn.Embedding(SEQ_LEN-1, feature_dim)
-        self.final_layer = torch.nn.Linear(feature_dim, vocab_size)
+        self.token_emb    = torch.nn.Embedding(vocab_size, feature_dim)
+        self.pos_emb      = torch.nn.Embedding(SEQ_LEN-1, feature_dim)
+        self.final_layer  = torch.nn.Linear(feature_dim, vocab_size)
 
     def forward(self, img_feat, target_seq):
-        x = self.token_emb(target_seq)
-        pos = torch.arange(x.size(1), device=x.device).clamp(max=self.pos_emb.num_embeddings - 1)
-        x = x + self.pos_emb(pos)
-        x = x + img_feat.unsqueeze(1)
+        x   = self.token_emb(target_seq)
+        pos = torch.arange(x.size(1), device=x.device).clamp(max=self.pos_emb.num_embeddings-1)
+        x   = x + self.pos_emb(pos)
+        x   = x + img_feat.unsqueeze(1)
         return self.final_layer(x)
 
-# Load ViT
+# ---------------------
+# Load models
+# ---------------------
 vit = ViTModel.from_pretrained("google/vit-base-patch16-224-in21k").to(device)
 vit.eval()
 
-# Load decoder weights from RADIOCAP13 folder
 decoder = BiasDecoder().to(device)
-decoder.load_state_dict(torch.load("main/pytorch_model.bin", map_location=device))
+decoder.load_state_dict(torch.load(weights_path, map_location=device))
 decoder.eval()
 
-# Load tokenizer from same folder
-tokenizer = SimpleTokenizer.load("RADIOCAP13")
-pad_idx = tokenizer.word2idx["<PAD>"]
+tokenizer = SimpleTokenizer.load(vocab_path)
+pad_idx   = tokenizer.word2idx["<PAD>"]
 
+# ---------------------
+# Caption generation
+# ---------------------
 @torch.no_grad()
-def generate_caption(img):
+def generate_caption(img, max_len=SEQ_LEN, beam_size=3):
     img_tensor = preprocess_image(img).unsqueeze(0).to(device)
-    img_feat = vit(pixel_values=img_tensor).pooler_output
+    img_feat   = vit(pixel_values=img_tensor).pooler_output
 
     beams = [([tokenizer.word2idx["<SOS>"]], 0.0)]
-    beam_size = 3
 
-    for _ in range(SEQ_LEN - 1):
+    for _ in range(max_len - 1):
         candidates = []
         for seq, score in beams:
-            inp = torch.tensor(seq + [pad_idx] * (SEQ_LEN - len(seq)), device=device).unsqueeze(0)
+            inp    = torch.tensor(seq + [pad_idx] * (SEQ_LEN - len(seq)), device=device).unsqueeze(0)
             logits = decoder(img_feat, inp)
-            probs = torch.nn.functional.log_softmax(logits[0, len(seq)-1], dim=-1)
+            probs  = torch.nn.functional.log_softmax(logits[0, len(seq)-1], dim=-1)
             top_p, top_i = torch.topk(probs, beam_size)
+
             for i in range(beam_size):
                 candidates.append((seq + [top_i[i].item()], score + top_p[i].item()))
+
         beams = sorted(candidates, key=lambda x: x[1], reverse=True)[:beam_size]
+
         if all(s[-1] == tokenizer.word2idx["<EOS>"] for s, _ in beams):
             break
 
     words = [tokenizer.idx2word.get(i, "<UNK>") for i in beams[0][0][1:] if i != pad_idx]
     return " ".join(words)
 
+# ---------------------
+# Gradio interface
+# ---------------------
 with gr.Blocks() as demo:
     gr.Markdown("# RADIOCAP13 — Image Captioning Demo")
     gr.Markdown(f"**Device:** {'GPU 🚀' if torch.cuda.is_available() else 'CPU 🐢'}")
 
     img_in = gr.Image(type="pil", label="Upload an Image")
-    out = gr.Textbox(label="Generated Caption")
-    btn = gr.Button("Generate Caption")
+    out    = gr.Textbox(label="Generated Caption")
+    btn    = gr.Button("Generate Caption")
     status = gr.Markdown("Ready.")
 
     def wrapped(img):