Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,97 +1,18 @@
 import streamlit as st
-import torch
-import os
+import io
 from PIL import Image
-import torchvision.transforms as transforms
-from transformers import (
-    ConvNextForImageClassification,
-    AutoTokenizer,
-    AutoModelForCausalLM
-)
 
-# Set Streamlit UI
-st.set_page_config(page_title="🍽️ Food Nutrition Estimator", page_icon="🥗", layout="centered")
-st.title("🍽️ Food Nutrition Estimator")
-st.markdown("Upload a food image and get nutritional information generated by AI!")
+st.set_page_config(page_title="🧪 Upload Debug Test", layout="centered")
+st.title("🧪 Upload Debug Test")
 
-# Environment & cache setup
-hf_token = os.getenv("HF_TOKEN")
-cache_dir = "/tmp/cache"
-os.makedirs(cache_dir, exist_ok=True)
-os.environ["HUGGINGFACE_HUB_CACHE"] = cache_dir
-
-# 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")
-st.sidebar.markdown("""
-- 🖼️ **Image Classifier**: `shingguy1/food-calorie-convnext`  
-- 💬 **Text Generator**: `TinyLlama/TinyLlama-1.1B-Chat-v1.0`
-""")
-
-# Load models
-@st.cache_resource
-def load_models():
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    # 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
-
-model_convnext, tokenizer, model_llm, device = load_models()
-
-# Upload image
-uploaded_file = st.file_uploader("Upload a food image...", type=["jpg", "jpeg", "png"])
+uploaded_file = st.file_uploader("Upload an image...", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
+    st.write("Upload received!")
     try:
-        image = Image.open(uploaded_file).convert("RGB")
-        st.image(image, caption="Uploaded Image", use_column_width=True)
-
-        # Predict with ConvNeXt
-        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"Give the calories, macros, and nutritional facts of a {pred_label}."
-        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)
-
-        st.info(caption)
-
+        file_bytes = uploaded_file.read()
+        image = Image.open(io.BytesIO(file_bytes)).convert("RGB")
+        st.image(image, caption="Image Loaded", use_column_width=True)
+        st.success("✅ Image loaded successfully.")
     except Exception as e:
-        st.error(f"Something went wrong: {e}")
-
-# Footer
-st.markdown("---")
-st.markdown("Built with ❤️ using Streamlit and Hugging Face by **shingguy1**")
+        st.error(f"❌ Failed to open image: {e}")