Update src/streamlit_app.py

src/streamlit_app.py

@@ -4,7 +4,7 @@ import os
 from PIL import Image
 import torchvision.transforms as transforms
 from transformers import (
-    ConvNextForImageClassification,
+    ViTForImageClassification,
     AutoTokenizer,
     AutoModelForCausalLM
 )
@@ -24,10 +24,10 @@ cache_dir = "/tmp/cache"
 os.makedirs(cache_dir, exist_ok=True)
 os.environ["HUGGINGFACE_HUB_CACHE"] = cache_dir
 
-# Transform for ConvNeXt
+# Transform for ViT
 manual_transform = transforms.Compose([
-    transforms.Resize(224),
-    transforms.CenterCrop(196),
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.485, 0.456, 0.406],
                          std=[0.229, 0.224, 0.225]),
@@ -37,7 +37,7 @@ manual_transform = transforms.Compose([
 # Sidebar info
 st.sidebar.header("Models Used")
 st.sidebar.markdown("""
-- 🖼️ **Image Classifier**: `shingguy1/fine_tuned_convnext`  
+- 🖼️ **Image Classifier**: `shingguy1/fine_tuned_model`  
 - 💬 **Text Generator**: `TinyLlama/TinyLlama-1.1B-Chat-v1.0`
 """)
 
@@ -46,9 +46,9 @@ st.sidebar.markdown("""
 def load_models():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    # ConvNeXt for classification
-    model_convnext = ConvNextForImageClassification.from_pretrained(
-        "shingguy1/fine_tuned_convnext",
+    # ViT for classification
+    model_vit = ViTForImageClassification.from_pretrained(
+        "shingguy1/fine_tuned_model",
         cache_dir=cache_dir,
         token=hf_token
     ).to(device)
@@ -65,9 +65,9 @@ def load_models():
         device_map="auto"
     )
 
-    return model_convnext, tokenizer, model_llm, device
+    return model_vit, tokenizer, model_llm, device
 
-model_convnext, tokenizer, model_llm, device = load_models()
+model_vit, tokenizer, model_llm, device = load_models()
 
 # Image uploader
 uploaded_file = st.file_uploader("Upload a food image...", type=["jpg", "jpeg", "png"])
@@ -81,12 +81,12 @@ if uploaded_file is not None:
         # Predict food label
         input_tensor = manual_transform(image).unsqueeze(0).to(device)
         with torch.no_grad():
-            outputs = model_convnext(pixel_values=input_tensor)
+            outputs = model_vit(pixel_values=input_tensor)
         pred_idx = outputs.logits.argmax(-1).item()
-        pred_label = model_convnext.config.id2label[pred_idx]
+        pred_label = model_vit.config.id2label[pred_idx]
         st.success(f"🍴 Predicted Food: **{pred_label}**")
 
-        # Generate nutrition description with LLM
+        # Generate nutrition description with LLM (no echo)
         prompt = (
             f"Please provide a concise nutritional overview for a {pred_label}. "
             "Include typical serving size, approximate calories, macronutrient breakdown "
@@ -95,16 +95,19 @@ if uploaded_file is not None:
         st.subheader("🧾 Nutrition Information")
         st.write(f"🤖 Prompt to LLM:\n\n{prompt}")
 
-        input_ids = tokenizer(prompt, return_tensors="pt").to(model_llm.device)
-        with torch.no_grad():
-            output = model_llm.generate(
-                **input_ids,
-                max_new_tokens=300,
-                temperature=0.8,
-                top_p=0.9,
-                do_sample=True
-            )
-        caption = tokenizer.decode(output[0], skip_special_tokens=True).strip()
+        inputs = tokenizer(prompt, return_tensors="pt").to(model_llm.device)
+        input_len = inputs.input_ids.shape[1]
+        output_ids = model_llm.generate(
+            **inputs,
+            max_new_tokens=200,
+            temperature=0.8,
+            top_p=0.9,
+            do_sample=True,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.eos_token_id
+        )[0]
+        generated_ids = output_ids[input_len:]
+        caption = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
         st.info(caption)
 
     except Exception as e: