Upload 2 files

app.py

@@ -52,7 +52,7 @@ with st.sidebar:
 
     if st.button("🚀 Initialize Engine"):
         manager = BitNetManager()
-        if manager.setup_engine():
+        if manager.setup_engine(model_id=model_id):
             st.session_state.engine_ready = True
             st.success("BitNet.cpp Ready!")
         else:

model_manager.py

@@ -57,13 +57,19 @@ class BitNetManager:
                     f.write(patched_setup)
                 st.success("Patch applied to setup_env.py (unquoted)!")
 
-    def setup_engine(self):
+    def setup_engine(self, model_id="1bitLLM/bitnet_b1_58-3B"):
         """Clone and compile utilizing official setup_env.py with log streaming."""
-        # Check if already compiled
+        model_name = model_id.split("/")[-1]
+        model_path = os.path.join(self.bitnet_dir, "models", model_name, "ggml-model-i2_s.gguf")
         binary = self.get_binary_path()
-        if binary and os.path.exists(binary):
-            st.success("BitNet engine is already compiled and ready!")
+        
+        # Check if already compiled AND model exists
+        if binary and os.path.exists(binary) and os.path.exists(model_path):
+            st.success(f"BitNet engine and model ({model_name}) are ready!")
             return True
+        
+        if binary and os.path.exists(binary):
+            st.info(f"Engine binary found, but model weights for {model_name} are missing. Starting setup...")
 
         if not os.path.exists(self.bitnet_dir):
             st.info("Cloning BitNet repository...")
@@ -129,19 +135,19 @@ class BitNetManager:
         return None
 
     def download_model(self, model_id="1bitLLM/bitnet_b1_58-3B", filename="ggml-model-i2_s.gguf"):
-        """Download model weights if not already handled by setup_env.py."""
-        # setup_env.py usually downloads to models/ (relative to bitnet_dir)
-        local_model_path = os.path.join(self.bitnet_dir, "models", model_id, filename)
+        """Locate the model weights. These are generated locally by setup_env.py."""
+        # setup_env.py downloads weights to models/<model_name>/
+        # e.g. models/bitnet_b1_58-3B/
+        model_name = model_id.split("/")[-1]
+        local_model_path = os.path.join(self.bitnet_dir, "models", model_name, filename)
+        
         if os.path.exists(local_model_path):
+            st.success(f"Found local model: {model_name}")
             return local_model_path
             
-        st.info(f"Downloading model {model_id} via Hub...")
-        try:
-            path = hf_hub_download(repo_id=model_id, filename=filename)
-            return path
-        except Exception as e:
-            st.error(f"Model download failed: {e}")
-            return None
+        st.error(f"Model file not found at {local_model_path}")
+        st.info("The GGUF model must be generated by the 'Initialize Engine' process. Please run it again to download and convert the weights.")
+        return None
 
     def run_inference(self, prompt, model_path):
         """Execute the bitnet binary with the provided prompt."""