Update app.py

app.py

@@ -3,9 +3,7 @@ import torch
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from huggingface_hub import login
-import re
 import os
-from PIL import Image
 
 # Load Hugging Face token
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -29,6 +27,28 @@ MODELS = {
 USER_PFP = "user.png"  # Hugging Face user avatar
 AI_PFP = "ai_pfp.png"  # Replace with the path to your AI's image or a URL
 
+# Custom CSS for blue sliders and button
+st.markdown(
+    """
+    <style>
+    /* Blue slider */
+    .stSlider > div > div > div > div {
+        background-color: #1f78b4 !important;
+    }
+    /* Blue button */
+    .stButton > button {
+        background-color: #1f78b4 !important;
+        color: white !important;
+        border: none !important;
+    }
+    .stButton > button:hover {
+        background-color: #16609a !important;
+    }
+    </style>
+    """,
+    unsafe_allow_html=True,
+)
+
 class AtlasInferenceApp:
     def __init__(self):
         if "current_model" not in st.session_state:
@@ -81,7 +101,7 @@ class AtlasInferenceApp:
                 "config": {
                     "name": f"{MODELS[model_key]['name']} {model_size}",
                     "path": model_path,
-                    "system_prompt": os.getenv(MODELS[model_key]["system_prompt_env"]),  # Load system prompt from env
+                    "system_prompt": os.getenv(MODELS[model_key]["system_prompt_env"], "Default system prompt"),  # Load system prompt from env
                 }
             })
             return f"✅ {MODELS[model_key]['name']} {model_size} loaded successfully!"
@@ -89,7 +109,7 @@ class AtlasInferenceApp:
             return f"❌ Error: {str(e)}"
 
     def respond(self, message, max_tokens, temperature, top_p, top_k, image=None):
-        if not st.session_state.current_model["model"]:
+        if not st.session_state.current_model["model"] or not st.session_state.current_model["tokenizer"]:
             return "⚠️ Please select and load a model first"
 
         try: