Update app.py

app.py

@@ -1,53 +1,31 @@
 import streamlit as st
-import torch
-from PIL import Image
-from fourm.models.fm import FM
-from torchvision import transforms
+import os
 
-# Set page to wide mode
-st.set_page_config(layout="wide")
-st.title("4M-7_B Multimodal Demo (CPU Mode)")
+# 1. Page Config
+st.set_page_config(page_title="4M CPU Boot", layout="wide")
 
-# 1. Load Model (Optimized for CPU)
-@st.cache_resource
-def load_4m_model():
-    # Load the 0.4B parameter model directly to CPU
-    model = FM.from_pretrained("EPFL-VILAB/4M-7_B_CC12M")
-    model.to("cpu")
-    model.eval()
-    return model
+st.title("4M-7_B Multimodal (CPU Optimization)")
 
-with st.spinner("Loading model to CPU (this takes about 1-2 mins)..."):
-    model = load_4m_model()
+# 2. Delayed Imports (Helps avoid crash on start)
+try:
+    import torch
+    from PIL import Image
+    from fourm.models.fm import FM
+    from torchvision import transforms
+    st.success("All libraries loaded successfully!")
+except ImportError as e:
+    st.error(f"Missing library: {e}")
+    st.stop()
 
-# 2. UI for Image Upload
-uploaded_file = st.file_uploader("Upload an image to analyze", type=["jpg", "png", "jpeg"])
+# 3. Load Model with local_files_only=False first time
+@st.cache_resource
+def load_model():
+    # This downloads the model weights (approx 1.6GB)
+    return FM.from_pretrained("EPFL-VILAB/4M-7_B_CC12M").to("cpu").eval()
 
-if uploaded_file:
-    img = Image.open(uploaded_file).convert('RGB')
-    
-    col1, col2 = st.columns(2)
-    with col1:
-        st.image(img, caption="Original Image")
-    
-    # Preprocessing
-    preprocess = transforms.Compose([
-        transforms.Resize(224),
-        transforms.CenterCrop(224),
-        transforms.ToTensor(),
-    ])
-    input_tensor = preprocess(img).unsqueeze(0)
+if st.button("Initialize Model"):
+    with st.spinner("Downloading/Loading weights to RAM..."):
+        model = load_model()
+        st.write("Model is ready for inference!")
 
-    # 3. Inference
-    with st.spinner("Processing on CPU..."):
-        with torch.no_grad():
-            # Generate depth/normals (target_modalities can be 'normal' or 'depth')
-            preds = model.generate({'rgb': input_tensor}, target_modalities=['normal'])
-            
-            # Post-process output
-            output_tensor = preds['normal'][0].clamp(0, 1)
-            output_img = transforms.ToPILImage()(output_tensor)
-    
-    with col2:
-        st.image(output_img, caption="Predicted Surface Normals")
-        st.success("Done!")
+# Add your inference code (the upload and predict part) below here