fault-free model loading

app.py

@@ -63,22 +63,9 @@ if "model_name" not in st.session_state:
 if "llm" not in st.session_state:
     st.session_state.llm = None
 
-# Make sure models dir exists
+# Ensure model directory exists
 os.makedirs("models", exist_ok=True)
 
-# If the selected model file does not exist or is invalid, clean up and re-download
-def validate_or_download_model():
-    if not os.path.exists(model_path):
-        cleanup_old_models()
-        download_model()
-        return
-    try:
-        _ = Llama(model_path=model_path, n_ctx=16, n_threads=1)  # dummy check
-    except Exception as e:
-        st.warning(f"Model file was invalid or corrupt: {e}\nRedownloading...")
-        cleanup_old_models()
-        download_model()
-
 def cleanup_old_models():
     for f in os.listdir("models"):
         if f.endswith(".gguf") and f != selected_model["filename"]:
@@ -96,28 +83,40 @@ def download_model():
             local_dir_use_symlinks=False,
         )
 
-validate_or_download_model()
+def try_load_model(path):
+    try:
+        return Llama(model_path=path, n_ctx=1024, n_threads=2, n_threads_batch=2, n_batch=4, n_gpu_layers=0, use_mlock=False, use_mmap=True, verbose=False)
+    except Exception as e:
+        return str(e)
+
+def validate_or_download_model():
+    if not os.path.exists(model_path):
+        cleanup_old_models()
+        download_model()
+
+    # First load attempt
+    result = try_load_model(model_path)
+    if isinstance(result, str):
+        st.warning(f"Initial load failed: {result}\nAttempting re-download...")
+        try:
+            os.remove(model_path)
+        except:
+            pass
+        cleanup_old_models()
+        download_model()
+        result = try_load_model(model_path)
+        if isinstance(result, str):
+            st.error(f"Model still failed after re-download: {result}")
+            st.stop()
+        return result
+    return result
 
 # Load model if changed
 if st.session_state.model_name != selected_model_name:
     if st.session_state.llm is not None:
         del st.session_state.llm
         gc.collect()
-    try:
-        st.session_state.llm = Llama(
-            model_path=model_path,
-            n_ctx=1024,
-            n_threads=2,
-            n_threads_batch=2,
-            n_batch=4,
-            n_gpu_layers=0,
-            use_mlock=False,
-            use_mmap=True,
-            verbose=False,
-        )
-    except Exception as e:
-        st.error(f"Failed to load model: {e}")
-        st.stop()
+    st.session_state.llm = validate_or_download_model()
     st.session_state.model_name = selected_model_name
 
 llm = st.session_state.llm