Update app.py

app.py

@@ -1,164 +1,207 @@
 import os
 import subprocess
 import streamlit as st
-from huggingface_hub import snapshot_download, login, HfApi
+from huggingface_hub import snapshot_download, HfApi
 
-if "quantized_model_path" not in st.session_state:
-    st.session_state.quantized_model_path = None
-if "upload_to_hf" not in st.session_state:
-    st.session_state.upload_to_hf = False
+# ============================================================
+# SESSION STATE
+# ============================================================
 
-def check_directory_path(directory_name: str) -> str:
-    if os.path.exists(directory_name):
-        path = os.path.abspath(directory_name)
-        return str(path)
+if "quantized_models" not in st.session_state:
+    st.session_state.quantized_models = []
 
+# ============================================================
+# CONFIG
+# ============================================================
 
+MODELS_LIST = ['rahul7star/Qwen3-4B-Thinking-2509-Genius-Coder-AI-Full']
 
-models_list=['rahul7star/Qwen3-4B-Thinking-2509-Genius-Coder-AI-Full']
-# Define quantization types
-QUANT_TYPES = ["Q2_K", "Q3_K_l", "Q3_K_M", "Q3_K_S", "Q4_0", "Q4_1", "Q4_K_M", "Q4_K_S", "Q5_0", "Q5_1", "Q5_K_M", "Q5_K_S", "Q6_K", "Q8_0", "BF16", "F16", "F32"]
+QUANT_TYPES = [
+    "Q2_K", "Q3_K_l", "Q3_K_M", "Q3_K_S",
+    "Q4_0", "Q4_1", "Q4_K_M", "Q4_K_S",
+    "Q5_0", "Q5_1", "Q5_K_M", "Q5_K_S",
+    "Q6_K", "Q8_0", "BF16", "F16", "F32"
+]
 
-model_dir_path = check_directory_path("/app/llama.cpp")
+LLAMA_CPP_PATH = "/app/llama.cpp"
+CONVERT_SCRIPT = f"{LLAMA_CPP_PATH}/convert_hf_to_gguf.py"
+QUANTIZE_BIN = f"{LLAMA_CPP_PATH}/build/bin/llama-quantize"
+
+# ============================================================
+# UTILS
+# ============================================================
+
+def check_dependencies():
+    if not os.path.exists(CONVERT_SCRIPT):
+        st.error("❌ convert_hf_to_gguf.py not found")
+        st.stop()
+    if not os.path.exists(QUANTIZE_BIN):
+        st.error("❌ llama-quantize binary not found")
+        st.stop()
 
 def download_model(hf_model_name, output_dir="/tmp/models"):
-    """
-    Downloads a Hugging Face model and saves it locally.
-    """
-    st.write(f"📥 Downloading `{hf_model_name}` from Hugging Face...")
-    os.makedirs(output_dir, exist_ok=True)
-    snapshot_download(repo_id=hf_model_name, local_dir=output_dir, local_dir_use_symlinks=False)
-    st.success("✅ Model downloaded successfully!")
-
-def convert_to_gguf(model_dir, output_file):
-    """
-    Converts a Hugging Face model to GGUF format.
-    """
-    st.write(f"🔄 Converting `{model_dir}` to GGUF format...")
-    os.makedirs(os.path.dirname(output_file), exist_ok=True)
+    st.write(f"📥 Downloading `{hf_model_name}` ...")
+    model_path = snapshot_download(
+        repo_id=hf_model_name,
+        local_dir=output_dir,
+        local_dir_use_symlinks=False
+    )
+    st.success("✅ Model downloaded")
+    return model_path
+
+def convert_to_gguf(model_path, output_file):
+    st.write("🔄 Converting to GGUF...")
     cmd = [
-        "python3", "/app/llama.cpp/convert_hf_to_gguf.py", model_dir,
-        "--outfile", output_file
+        "python3",
+        CONVERT_SCRIPT,
+        model_path,
+        "--outfile",
+        output_file
     ]
-    process = subprocess.run(cmd, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-    if process.returncode == 0:
-        st.success(f"✅ Conversion complete: `{output_file}`")
-    else:
-        st.error(f"❌ Conversion failed: {process.stderr}")
-
-def quantize_llama(model_path, quantized_output_path, quant_type):
-    """
-    Quantizes a GGUF model.
-    """
-    st.write(f"⚡ Quantizing `{model_path}` with `{quant_type}` precision...")
-    os.makedirs(os.path.dirname(quantized_output_path), exist_ok=True)
-    quantize_path = "/app/llama.cpp/build/bin/llama-quantize"
-    
+    result = subprocess.run(cmd, capture_output=True, text=True)
+    if result.returncode != 0:
+        st.error(result.stderr)
+        raise RuntimeError("Conversion failed")
+    st.success("✅ GGUF created")
+
+def quantize_model(gguf_file, quant_type):
+    output_file = gguf_file.replace(".gguf", f"-{quant_type}.gguf")
+
+    st.write(f"⚡ Quantizing → {quant_type}")
     cmd = [
-        "/app/llama.cpp/build/bin/llama-quantize", 
-        model_path, 
-        quantized_output_path,
+        QUANTIZE_BIN,
+        gguf_file,
+        output_file,
         quant_type
     ]
-    
-    process = subprocess.run(cmd, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-
-    if process.returncode == 0:
-        st.success(f"✅ Quantized model saved at `{quantized_output_path}`")
-    else:
-        st.error(f"❌ Quantization failed: {process.stderr}")
-
-def automate_llama_quantization(hf_model_name, quant_type):
-    """
-    Orchestrates the entire quantization process.
-    """
-    output_dir = "/tmp/models"
-    gguf_file = os.path.join(output_dir, f"{hf_model_name.replace('/', '_')}.gguf")
-    quantized_file = gguf_file.replace(".gguf", f"-{quant_type}.gguf")
-
-    progress_bar = st.progress(0)
-
-    # Step 1: Download
-    st.write("### Step 1: Downloading Model")
-    download_model(hf_model_name, output_dir)
-    progress_bar.progress(33)
-
-    # Step 2: Convert to GGUF
-    st.write("### Step 2: Converting Model to GGUF Format")
-    convert_to_gguf(output_dir, gguf_file)
-    progress_bar.progress(66)
-
-    # Step 3: Quantize Model
-    st.write("### Step 3: Quantizing Model")
-    quantize_llama(gguf_file, quantized_file, quant_type.lower())
-    progress_bar.progress(100)
-
-    st.success(f"🎉 All steps completed! Quantized model available at: `{quantized_file}`")
-    return quantized_file
-
-def upload_to_huggingface(file_path, repo_id, token):
-    """
-    Uploads a file to Hugging Face Hub.
-    """
-    try:
-        # Log in to Hugging Face
-        login(token=token)
-
-        # Initialize HfApi
-        api = HfApi()
-
-        # Create the repository if it doesn't exist
-        api.create_repo(repo_id, exist_ok=True, repo_type="model")
-
-        # Upload the file
-        api.upload_file(
-            path_or_fileobj=file_path,
-            path_in_repo=os.path.basename(file_path),
-            repo_id=repo_id,
-        )
-        st.success(f"✅ File uploaded to Hugging Face: {repo_id}")
-
-        # Reset session state and rerun
-        st.session_state.quantized_model_path = None
-        st.session_state.upload_to_hf = False
-        st.rerun()
-    except Exception as e:
-        st.error(f"❌ Failed to upload file: {e}")
-
-st.title("🦙 LLaMA Model Quantization (llama.cpp)")
-
-
-selected_model = st.selectbox("Select the Hugging Face Model", models_list, index=None)
-hf_model_name = selected_model if selected_model else st.text_input("Enter Hugging Face Model (If not there in the above list)")
-
-quant_type = st.selectbox("Select Quantization Type", QUANT_TYPES)
-start_button = st.button("🚀 Start Quantization")
-
-if start_button:
-    if hf_model_name and quant_type:
-        with st.spinner("Processing..."):
-            st.session_state.quantized_model_path = automate_llama_quantization(hf_model_name, quant_type)
-    else:
-        st.warning("Please select/enter the necessary fields.")
-
-if st.session_state.quantized_model_path:
-    with open(st.session_state.quantized_model_path, "rb") as f:
-        if st.download_button("⬇️ Download Quantized Model", f, file_name=os.path.basename(st.session_state.quantized_model_path)):
-            st.session_state.quantized_model_path = None
-            st.session_state.upload_to_hf = False
-            st.rerun()
-    
-    # Checkbox for upload section
-    st.session_state.upload_to_hf = st.checkbox("Upload to Hugging Face", value=st.session_state.upload_to_hf)
-    
-    if st.session_state.upload_to_hf:
-        st.write("### Upload to Hugging Face")
-        repo_id = st.text_input("Enter Hugging Face Repository ID (e.g., 'username/repo-name')")
-        hf_token = st.text_input("Enter Hugging Face Token", type="password")
-        
-        if st.button("📤 Upload to Hugging Face"):
-            if repo_id and hf_token:
-                with st.spinner("Uploading..."):
-                    upload_to_huggingface(st.session_state.quantized_model_path, repo_id, hf_token)
-            else:
-                st.warning("Please provide a valid repository ID and Hugging Face token.")
+
+    result = subprocess.run(cmd, capture_output=True, text=True)
+
+    if result.returncode != 0:
+        st.error(result.stderr)
+        return None
+
+    st.success(f"✅ {quant_type} done")
+    return output_file
+
+def upload_to_huggingface(file_path, repo_id):
+    hf_token = os.getenv("HF_TOKEN")
+
+    if not hf_token:
+        st.error("❌ HF_TOKEN not found in environment variables")
+        return
+
+    api = HfApi(token=hf_token)
+
+    api.create_repo(repo_id, exist_ok=True, repo_type="model")
+
+    api.upload_file(
+        path_or_fileobj=file_path,
+        path_in_repo=os.path.basename(file_path),
+        repo_id=repo_id,
+    )
+
+    st.success(f"🚀 Uploaded to https://huggingface.co/{repo_id}")
+
+# ============================================================
+# UI
+# ============================================================
+
+st.title("🦙 LLaMA.cpp Multi-Quantization Tool")
+
+check_dependencies()
+
+# Model selection
+selected_model = st.selectbox(
+    "Select Hugging Face Model",
+    MODELS_LIST,
+    index=None
+)
+
+hf_model_name = selected_model or st.text_input(
+    "Or Enter Custom HF Model ID"
+)
+
+# Multi-checkbox quant selection
+st.subheader("Select Quantization Types")
+
+selected_quants = []
+cols = st.columns(4)
+
+for i, quant in enumerate(QUANT_TYPES):
+    with cols[i % 4]:
+        if st.checkbox(quant):
+            selected_quants.append(quant)
+
+# Start button
+if st.button("🚀 Start Quantization"):
+
+    if not hf_model_name:
+        st.warning("Please enter a model name")
+        st.stop()
+
+    if not selected_quants:
+        st.warning("Select at least one quant type")
+        st.stop()
+
+    with st.spinner("Processing..."):
+        try:
+            base_dir = "/tmp/models"
+            os.makedirs(base_dir, exist_ok=True)
+
+            model_path = download_model(hf_model_name, base_dir)
+
+            gguf_file = os.path.join(
+                base_dir,
+                hf_model_name.replace("/", "_") + ".gguf"
+            )
+
+            convert_to_gguf(model_path, gguf_file)
+
+            st.session_state.quantized_models = []
+
+            for quant in selected_quants:
+                quant_file = quantize_model(gguf_file, quant)
+                if quant_file:
+                    st.session_state.quantized_models.append(quant_file)
+
+            st.success("🎉 All quantizations completed")
+
+        except Exception as e:
+            st.error(f"❌ Error: {str(e)}")
+
+# ============================================================
+# DOWNLOAD + UPLOAD SECTION
+# ============================================================
+
+if st.session_state.quantized_models:
+
+    st.subheader("📦 Generated Models")
+
+    for file_path in st.session_state.quantized_models:
+
+        with open(file_path, "rb") as f:
+            st.download_button(
+                label=f"⬇️ Download {os.path.basename(file_path)}",
+                data=f,
+                file_name=os.path.basename(file_path),
+                key=file_path
+            )
+
+    st.divider()
+
+    st.subheader("🚀 Upload to Hugging Face")
+
+    repo_id = st.text_input(
+        "Target Repository (e.g. username/model-quant)"
+    )
+
+    if st.button("📤 Upload All to HF"):
+        if not repo_id:
+            st.warning("Enter repository ID")
+        else:
+            with st.spinner("Uploading..."):
+                for file_path in st.session_state.quantized_models:
+                    upload_to_huggingface(file_path, repo_id)
+
+                st.success("✅ All files uploaded successfully")