Update src/streamlit_app.py

src/streamlit_app.py

@@ -17,12 +17,14 @@ st.markdown("Upload a food image and get nutritional information generated by AI
 
 # Environment & cache setup
 hf_token = os.getenv("HF_TOKEN")
+if not hf_token:
+    st.warning("HF_TOKEN not set. Please set the environment variable HF_TOKEN to access private models.")
 cache_dir = "/tmp/cache"
 os.makedirs(cache_dir, exist_ok=True)
 os.environ["HUGGINGFACE_HUB_CACHE"] = cache_dir
 
-# Use ConvNeXt's official image processor
-image_processor = ConvNextImageProcessor.from_pretrained("facebook/convnext-base-224-22k")
+# Use ConvNeXt's official image processor for a compatible model
+image_processor = ConvNextImageProcessor.from_pretrained("facebook/convnext-base-224-22k-1k")
 
 # Sidebar info
 st.sidebar.header("Models Used")
@@ -35,66 +37,75 @@ st.sidebar.markdown("""
 @st.cache_resource
 def load_models():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model_convnext = ConvNextForImageClassification.from_pretrained(
-        "shingguy1/food-calorie-convnext",
-        cache_dir=cache_dir,
-        token=hf_token
-    ).to(device)
-    tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", cache_dir=cache_dir)
-    model_llm = AutoModelForCausalLM.from_pretrained(
-        "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
-        cache_dir=cache_dir,
-        torch_dtype=torch.float32,
-        device_map="auto"
-    )
-    return model_convnext, tokenizer, model_llm, device
+    st.info(f"Using device: {device}")
 
-model_convnext, tokenizer, model_llm, device = load_models()
+    try:
+        # ConvNeXt for classification
+        model_convnext = ConvNextForImageClassification.from_pretrained(
+            "shingguy1/food-calorie-convnext",
+            cache_dir=cache_dir,
+            token=hf_token
+        ).to(device)
+    except Exception as e:
+        st.error(f"Failed to load ConvNeXt model: {e}")
+        st.stop()
 
-# Initialize session state for uploaded file
-if "uploaded_file" not in st.session_state:
-    st.session_state.uploaded_file = None
+    try:
+        # TinyLlama for nutritional facts
+        tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", cache_dir=cache_dir)
+        model_llm = AutoModelForCausalLM.from_pretrained(
+            "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+            cache_dir=cache_dir,
+            torch_dtype=torch.float16 if device.type == "cuda" else torch.float32,
+            device_map="auto"
+        )
+    except Exception as e:
+        st.error(f"Failed to load TinyLlama model: {e}")
+        st.stop()
 
-# Upload image
-uploaded_file = st.file_uploader("Upload a food image...", type=["jpg", "jpeg", "png"], key="file_uploader")
+    return model_convnext, tokenizer, model_llm, device
 
-# Update session state when a file is uploaded
-if uploaded_file is not None and st.session_state.uploaded_file != uploaded_file:
-    st.session_state.uploaded_file = uploaded_file
-    st.write(f"File uploaded: {uploaded_file.name}, Size: {uploaded_file.size} bytes")
-else:
-    st.warning("No file uploaded yet. Please upload a .jpg, .jpeg, or .png image.")
+with st.spinner("Loading models..."):
+    model_convnext, tokenizer, model_llm, device = load_models()
+
+# Upload image
+uploaded_file = st.file_uploader("Upload a food image...", type=["jpg", "jpeg", "png"])
 
-# Process uploaded file
-if st.session_state.uploaded_file is not None:
+if uploaded_file is not None:
     try:
-        image = Image.open(st.session_state.uploaded_file).convert("RGB")
+        image = Image.open(uploaded_file).convert("RGB")
         st.image(image, caption="Uploaded Image", use_column_width=True)
 
         # Predict with ConvNeXt
-        inputs = image_processor(image, return_tensors="pt").to(device)
-        with torch.no_grad():
-            outputs = model_convnext(**inputs)
-        pred_idx = outputs.logits.argmax(-1).item()
-        pred_label = model_convnext.config.id2label[pred_idx]
-        st.success(f"🍴 Predicted Food: **{pred_label}**")
+        with st.spinner("Classifying food..."):
+            inputs = image_processor(image, return_tensors="pt").to(device)
+            with torch.no_grad():
+                outputs = model_convnext(**inputs)
+            pred_idx = outputs.logits.argmax(-1).item()
+            pred_label = model_convnext.config.id2label[pred_idx]
+            st.success(f"🍴 Predicted Food: **{pred_label}**")
 
         # Generate nutrition caption using TinyLlama
-        prompt = f"Give the calories, macros, and nutritional facts of a {pred_label}."
+        prompt = f"Provide the calories, protein, fat, and carbs for a typical serving of {pred_label}. Format the response as: 'Calories: X kcal, Protein: Y g, Fat: Z g, Carbs: W g'."
         st.subheader("🧾 Nutrition Information")
         st.write(f"🤖 Prompt: `{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=100)
-        caption = tokenizer.decode(output[0], skip_special_tokens=True)
+        with st.spinner("Generating nutritional facts..."):
+            input_ids = tokenizer(prompt, return_tensors="pt").to(device)
+            with torch.no_grad():
+                output = model_llm.generate(
+                    **input_ids,
+                    max_new_tokens=150,
+                    temperature=0.7,
+                    top_p=0.9
+                )
+            caption = tokenizer.decode(output[0], skip_special_tokens=True)
+            caption = caption.replace(prompt, "").strip()  # Remove prompt if echoed
 
-        st.info(caption)
+        st.info(caption if caption else "No nutritional information generated.")
 
     except Exception as e:
         st.error(f"Something went wrong: {e}")
-else:
-    st.info("Please upload an image to get started.")
 
 # Footer
 st.markdown("---")