Update src/streamlit_app.py

src/streamlit_app.py

@@ -5,7 +5,6 @@ from PIL import Image
 import torchvision.transforms as transforms
 from transformers import (
     ConvNextForImageClassification,
-    ConvNextImageProcessor,
     AutoTokenizer,
     AutoModelForCausalLM
 )
@@ -17,14 +16,18 @@ 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 for a compatible model
-image_processor = ConvNextImageProcessor.from_pretrained("facebook/convnext-base-224-22k-1k")
+# Transform for ConvNeXt
+manual_transform = transforms.Compose([
+    transforms.Resize(224),
+    transforms.CenterCrop(196),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    transforms.ConvertImageDtype(torch.float32)
+])
 
 # Sidebar info
 st.sidebar.header("Models Used")
@@ -37,36 +40,26 @@ st.sidebar.markdown("""
 @st.cache_resource
 def load_models():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    st.info(f"Using device: {device}")
 
-    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()
-
-    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()
+    # ConvNeXt for classification
+    model_convnext = ConvNextForImageClassification.from_pretrained(
+        "shingguy1/food-calorie-convnext",
+        cache_dir=cache_dir,
+        token=hf_token
+    ).to(device)
+
+    # 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.float32,
+        device_map="auto"
+    )
 
     return model_convnext, tokenizer, model_llm, device
 
-with st.spinner("Loading models..."):
-    model_convnext, tokenizer, model_llm, device = load_models()
+model_convnext, tokenizer, model_llm, device = load_models()
 
 # Upload image
 uploaded_file = st.file_uploader("Upload a food image...", type=["jpg", "jpeg", "png"])
@@ -77,32 +70,24 @@ if uploaded_file is not None:
         st.image(image, caption="Uploaded Image", use_column_width=True)
 
         # Predict with ConvNeXt
-        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}**")
+        input_tensor = manual_transform(image).unsqueeze(0).to(device)
+        with torch.no_grad():
+            outputs = model_convnext(pixel_values=input_tensor)
+        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"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'."
+        prompt = f"Give the calories, macros, and nutritional facts of a {pred_label}."
         st.subheader("🧾 Nutrition Information")
         st.write(f"🤖 Prompt: `{prompt}`")
 
-        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 if caption else "No nutritional information generated.")
+        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)
+
+        st.info(caption)
 
     except Exception as e:
         st.error(f"Something went wrong: {e}")