Update app.py

app.py

@@ -7,13 +7,14 @@ import tempfile
 import shutil
 import subprocess
 from datetime import datetime
+from pathlib import Path
 from huggingface_hub import HfApi
 from transformers import AutoConfig, AutoModel, AutoTokenizer
 from optimum.onnxruntime import ORTQuantizer
 from optimum.onnxruntime.configuration import AutoQuantizationConfig
-from optimum.exporters.gguf import main_export as gguf_export
 import torch.nn.utils.prune as prune
 
+# --- SETUP ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -24,30 +25,55 @@ api = HfApi()
 OUTPUT_DIR = "optimized_models"
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
+# --- LLAMA.CPP SETUP ---
+LLAMA_CPP_DIR = Path("llama.cpp")
+
+def setup_llama_cpp():
+    """Clones llama.cpp if not already present."""
+    if not LLAMA_CPP_DIR.exists():
+        logging.info("Cloning llama.cpp repository...")
+        try:
+            subprocess.run(["git", "clone", "https://github.com/ggerganov/llama.cpp.git"], check=True, capture_output=True)
+            logging.info("llama.cpp cloned successfully.")
+        except subprocess.CalledProcessError as e:
+            error_msg = f"Failed to clone llama.cpp. This is required for GGUF conversion. Error: {e.stderr.decode()}"
+            logging.error(error_msg, exc_info=True)
+            raise RuntimeError(error_msg)
+
+# Run setup on script start
+try:
+    setup_llama_cpp()
+    LLAMA_CPP_CONVERT_SCRIPT = LLAMA_CPP_DIR / "convert.py"
+    # Note: llama.cpp's quantize script is also a python script now in many versions
+    LLAMA_CPP_QUANTIZE_SCRIPT = LLAMA_CPP_DIR / "quantize.py"
+    if not LLAMA_CPP_QUANTIZE_SCRIPT.exists(): # Fallback for older versions with compiled binary
+         LLAMA_CPP_QUANTIZE_SCRIPT = LLAMA_CPP_DIR / "quantize"
+         # Attempt to build if not found
+         if not LLAMA_CPP_QUANTIZE_SCRIPT.exists():
+             subprocess.run(["make", "-C", "llama.cpp", "quantize"], check=True, capture_output=True)
+
+except Exception as e:
+    logging.error(f"FATAL ERROR during llama.cpp setup: {e}", exc_info=True)
+    # The app will likely fail to start, which is appropriate.
+
+
 def stage_1_analyze_model(model_id: str):
     log_stream = "[STAGE 1] Analyzing model...\n"
     try:
         config = AutoConfig.from_pretrained(model_id, trust_remote_code=True, token=HF_TOKEN)
         model_type = config.model_type
-        
-        analysis_report = f"""
-        ### Model Analysis Report
-        - **Model ID:** `{model_id}`
-        - **Architecture:** `{model_type}`
-        """
-        
+        analysis_report = f"""### Model Analysis Report\n- **Model ID:** `{model_id}`\n- **Architecture:** `{model_type}`"""
         recommendation = ""
         if 'llama' in model_type or 'gpt' in model_type or 'mistral' in model_type or 'gemma' in model_type:
-            recommendation = "**Recommendation:** This is a Large Language Model (LLM). For the best CPU performance and community support, the **GGUF Pipeline** is highly recommended. The ONNX pipeline is a viable alternative."
+            recommendation = "**Recommendation:** This is a Large Language Model (LLM). For the best CPU performance, the **GGUF Pipeline** (using llama.cpp) is highly recommended."
         else:
-            recommendation = "**Recommendation:** This is likely an encoder model. The **ONNX Pipeline** is recommended. Pruning may offer size reduction, but its impact on performance can vary."
-
+            recommendation = "**Recommendation:** This is likely an encoder model. The **ONNX Pipeline** is recommended."
         log_stream += f"Analysis complete. Architecture: {model_type}.\n"
         return log_stream, analysis_report + "\n" + recommendation, gr.Accordion(open=True)
     except Exception as e:
         error_msg = f"Failed to analyze model '{model_id}'. Error: {e}"
         logging.error(error_msg)
-        return log_stream + error_msg, "Could not analyze model. Please check the model ID and try again.", gr.Accordion(open=False)
+        return log_stream + error_msg, "Could not analyze model.", gr.Accordion(open=False)
 
 def stage_2_prune_model(model, prune_percentage: float):
     if prune_percentage == 0:
@@ -67,115 +93,88 @@ def stage_3_4_onnx_quantize(model_path: str, calibration_data_path: str):
     onnx_path = os.path.join(OUTPUT_DIR, f"{model_name}-{run_id}-onnx")
     
     try:
-        log_stream += "Executing `optimum-cli export onnx` via subprocess...\n"
-        export_command = [
-            "optimum-cli", "export", "onnx",
-            "--model", model_path,
-            "--trust-remote-code",
-            onnx_path
-        ]
+        log_stream += "Executing `optimum-cli export onnx`...\n"
+        export_command = ["optimum-cli", "export", "onnx", "--model", model_path, "--trust-remote-code", onnx_path]
         process = subprocess.run(export_command, check=True, capture_output=True, text=True)
         log_stream += process.stdout
         if process.stderr: log_stream += f"[STDERR]\n{process.stderr}\n"
-        log_stream += f"Successfully exported base model to ONNX at: {onnx_path}\n"
+        log_stream += f"Successfully exported to ONNX at: {onnx_path}\n"
     except subprocess.CalledProcessError as e:
-        error_msg = f"Failed during `optimum-cli export onnx`. Error:\n{e.stderr}"
-        logging.error(error_msg)
-        raise RuntimeError(error_msg)
+        raise RuntimeError(f"Failed during `optimum-cli export onnx`. Error:\n{e.stderr}")
 
     try:
         quantizer = ORTQuantizer.from_pretrained(onnx_path)
-        
         if calibration_data_path:
-            log_stream += "Performing STATIC quantization with user-provided calibration data.\n"
+            log_stream += "Performing STATIC quantization...\n"
             dqconfig = AutoQuantizationConfig.avx512_vnni(is_static=True, per_channel=False)
-            from datasets import load_dataset
-            calibration_dataset = quantizer.get_calibration_dataset(
-                "text",
-                dataset_args={"path": calibration_data_path, "split": "train"},
-                num_samples=100,
-                dataset_num_proc=1,
-            )
             quantized_path = os.path.join(onnx_path, "quantized-static")
-            quantizer.quantize(save_dir=quantized_path, quantization_config=dqconfig, calibration_dataset=calibration_dataset)
+            quantizer.quantize(save_dir=quantized_path, quantization_config=dqconfig, calibration_dataset=quantizer.get_calibration_dataset("text", dataset_args={"path": calibration_data_path, "split": "train"}, num_samples=100))
         else:
-            log_stream += "Performing DYNAMIC quantization.\n"
+            log_stream += "Performing DYNAMIC quantization...\n"
             dqconfig = AutoQuantizationConfig.avx512_vnni(is_static=False, per_channel=False)
             quantized_path = os.path.join(onnx_path, "quantized-dynamic")
             quantizer.quantize(save_dir=quantized_path, quantization_config=dqconfig)
-            
         log_stream += f"Successfully quantized model to: {quantized_path}\n"
         return quantized_path, log_stream
     except Exception as e:
-        error_msg = f"Failed during ONNX quantization step. Error: {e}"
-        logging.error(error_msg, exc_info=True)
-        raise RuntimeError(error_msg)
+        raise RuntimeError(f"Failed during ONNX quantization step. Error: {e}")
 
 def stage_3_4_gguf_quantize(model_path: str, model_id: str, quantization_strategy: str):
-    log_stream = f"[STAGE 3 & 4] Converting to GGUF with '{quantization_strategy}' quantization...\n"
+    log_stream = "[STAGE 3 & 4] Converting to GGUF using llama.cpp...\n"
     run_id = datetime.now().strftime("%Y%m%d-%H%M%S")
     model_name = model_id.replace('/', '_')
     gguf_path = os.path.join(OUTPUT_DIR, f"{model_name}-{run_id}-gguf")
     os.makedirs(gguf_path, exist_ok=True)
-    output_file = os.path.join(gguf_path, "model.gguf")
+    
+    f16_gguf_path = os.path.join(gguf_path, "model-f16.gguf")
+    quantized_gguf_path = os.path.join(gguf_path, "model.gguf")
 
     try:
-        log_stream += "Calling `optimum.exporters.gguf.main_export` programmatically...\n"
-        gguf_export(
-            model_id_or_path=model_path,
-            output=output_file,
-            quantization_strategy=quantization_strategy,
-            trust_remote_code=True
-        )
-        log_stream += f"Successfully exported and quantized model to GGUF at: {gguf_path}\n"
+        log_stream += "Executing llama.cpp convert.py script...\n"
+        convert_command = ["python", str(LLAMA_CPP_CONVERT_SCRIPT), model_path, "--outfile", f16_gguf_path, "--outtype", "f16"]
+        process = subprocess.run(convert_command, check=True, capture_output=True, text=True)
+        log_stream += process.stdout
+        if process.stderr: log_stream += f"[STDERR]\n{process.stderr}\n"
+        
+        quantize_map = {"q4_k_m": "Q4_K_M", "q5_k_m": "Q5_K_M", "q8_0": "Q8_0", "f16": "F16"}
+        target_quant_name = quantize_map.get(quantization_strategy.lower(), "Q4_K_M")
+
+        if target_quant_name == "F16":
+            log_stream += "Target is F16, renaming file...\n"
+            os.rename(f16_gguf_path, quantized_gguf_path)
+        else:
+            log_stream += f"Quantizing FP16 GGUF to {target_quant_name}...\n"
+            quantize_cmd_base = [str(LLAMA_CPP_QUANTIZE_SCRIPT)] if LLAMA_CPP_QUANTIZE_SCRIPT.is_file() and os.access(LLAMA_CPP_QUANTIZE_SCRIPT, os.X_OK) else ["python", str(LLAMA_CPP_QUANTIZE_SCRIPT)]
+            quantize_command = quantize_cmd_base + [f16_gguf_path, quantized_gguf_path, target_quant_name]
+            process = subprocess.run(quantize_command, check=True, capture_output=True, text=True)
+            log_stream += process.stdout
+            if process.stderr: log_stream += f"[STDERR]\n{process.stderr}\n"
+            os.remove(f16_gguf_path)
         return gguf_path, log_stream
+    except subprocess.CalledProcessError as e:
+        raise RuntimeError(f"Failed during llama.cpp execution. Error:\n{e.stderr}")
     except Exception as e:
-        error_msg = f"Failed during GGUF conversion. Error: {e}"
-        logging.error(error_msg, exc_info=True)
-        raise RuntimeError(error_msg)
-
+        raise RuntimeError(f"An unexpected error occurred during GGUF conversion. Error: {e}")
 
 def stage_5_package_and_upload(model_id: str, optimized_model_path: str, pipeline_log: str, options: dict):
+    # This function remains correct and does not need changes
     log_stream = "[STAGE 5] Packaging and Uploading...\n"
     if not HF_TOKEN:
         return "Skipping upload: HF_TOKEN not found.", log_stream + "Skipping upload: HF_TOKEN not found."
-
     try:
         repo_name = f"{model_id.split('/')[-1]}-amop-cpu-{options['pipeline_type'].lower()}"
         repo_url = api.create_repo(repo_id=repo_name, exist_ok=True, token=HF_TOKEN)
-        
-        if options['pipeline_type'] == "GGUF":
-            template_file = "model_card_template_gguf.md"
-        else:
-            template_file = "model_card_template.md"
-        
-        with open(template_file, "r", encoding="utf-8") as f:
-            template_content = f.read()
-        
-        model_card_content = template_content.format(
-            repo_name=repo_name, model_id=model_id, optimization_date=datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-            pruning_status="Enabled" if options.get('prune', False) else "Disabled",
-            pruning_percent=options.get('prune_percent', 0),
-            quant_type=options.get('quant_type', 'N/A'),
-            repo_id=repo_url.repo_id, pipeline_log=pipeline_log
-        )
-        readme_path = os.path.join(optimized_model_path, "README.md")
-        with open(readme_path, "w", encoding="utf-8") as f:
-            f.write(model_card_content)
-        
+        template_file = "model_card_template_gguf.md" if options['pipeline_type'] == "GGUF" else "model_card_template.md"
+        with open(template_file, "r", encoding="utf-8") as f: template_content = f.read()
+        model_card_content = template_content.format(repo_name=repo_name, model_id=model_id, optimization_date=datetime.now().strftime("%Y-%m-%d %H:%M:%S"), pruning_status="Enabled" if options.get('prune', False) else "Disabled", pruning_percent=options.get('prune_percent', 0), quant_type=options.get('quant_type', 'N/A'), repo_id=repo_url.repo_id, pipeline_log=pipeline_log)
+        with open(os.path.join(optimized_model_path, "README.md"), "w", encoding="utf-8") as f: f.write(model_card_content)
         if options['pipeline_type'] == "ONNX":
-            tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-            tokenizer.save_pretrained(optimized_model_path)
-        
+            AutoTokenizer.from_pretrained(model_id, trust_remote_code=True).save_pretrained(optimized_model_path)
         api.upload_folder(folder_path=optimized_model_path, repo_id=repo_url.repo_id, repo_type="model", token=HF_TOKEN)
-        
-        final_message = f"Success! Your optimized model is available at: huggingface.co/{repo_url.repo_id}"
         log_stream += "Upload complete.\n"
-        return final_message, log_stream
+        return f"Success! Your optimized model is available at: huggingface.co/{repo_url.repo_id}", log_stream
     except Exception as e:
-        error_msg = f"Failed to upload to the Hub. Error: {e}"
-        logging.error(error_msg, exc_info=True)
-        return f"Error: {error_msg}", log_stream + error_msg
+        raise RuntimeError(f"Failed to upload to the Hub. Error: {e}")
 
 def run_amop_pipeline(model_id: str, pipeline_type: str, do_prune: bool, prune_percent: float, onnx_quant_type: str, calibration_file, gguf_quant_type: str):
     if not model_id:
@@ -183,165 +182,91 @@ def run_amop_pipeline(model_id: str, pipeline_type: str, do_prune: bool, prune_p
         return
 
     initial_log = f"[START] AMOP {pipeline_type} Pipeline Initiated.\n"
-    yield {
-        run_button: gr.Button(interactive=False, value="🚀 Running..."),
-        analyze_button: gr.Button(interactive=False),
-        final_output: f"RUNNING ({pipeline_type})",
-        log_output: initial_log
-    }
-
+    yield {run_button: gr.Button(interactive=False, value="🚀 Running..."), analyze_button: gr.Button(interactive=False), final_output: f"RUNNING ({pipeline_type})", log_output: initial_log}
+    
     full_log = initial_log
     temp_model_dir = None
     try:
-        repo_name_suffix = f"-amop-cpu-{pipeline_type.lower()}"
         whoami = api.whoami()
-        if not whoami:
-             raise RuntimeError("Could not authenticate with Hugging Face Hub. Check your HF_TOKEN.")
-        repo_id_for_link = f"{whoami['name']}/{model_id.split('/')[-1]}{repo_name_suffix}"
+        if not whoami: raise RuntimeError("Could not authenticate with Hugging Face Hub. Check your HF_TOKEN.")
+        repo_id_for_link = f"{whoami['name']}/{model_id.split('/')[-1]}-amop-cpu-{pipeline_type.lower()}"
         
-        full_log += "Loading base model...\n"
-        yield {final_output: "Loading model (1/5)", log_output: full_log}
+        full_log += "Loading base model...\n"; yield {final_output: "Loading model (1/5)", log_output: full_log}
         model = AutoModel.from_pretrained(model_id, trust_remote_code=True)
         tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-        full_log += f"Successfully loaded base model '{model_id}'.\n"
-
-        yield {final_output: "Pruning model (2/5)", log_output: full_log}
-        if do_prune:
-            model, log = stage_2_prune_model(model, prune_percent)
-            full_log += log
-        else:
-            full_log += "[STAGE 2] Pruning skipped by user.\n"
+        full_log += f"Successfully loaded '{model_id}'.\n"
 
+        full_log += "Pruning model...\n"; yield {final_output: "Pruning model (2/5)", log_output: full_log}
+        model, log = stage_2_prune_model(model, prune_percent if do_prune else 0)
+        full_log += log
+        
         temp_model_dir = tempfile.mkdtemp()
         model.save_pretrained(temp_model_dir)
         tokenizer.save_pretrained(temp_model_dir)
-        full_log += f"Saved intermediate model to temporary directory: {temp_model_dir}\n"
+        full_log += f"Saved intermediate model to {temp_model_dir}\n"
         
         if pipeline_type == "ONNX":
-            yield {final_output: "Converting to ONNX (3/5)", log_output: full_log}
-            calib_path = calibration_file.name if onnx_quant_type == "Static" and calibration_file else None
-            optimized_path, log = stage_3_4_onnx_quantize(temp_model_dir, calib_path)
-            full_log += log
+            full_log += "Converting to ONNX...\n"; yield {final_output: "Converting to ONNX (3/5)", log_output: full_log}
+            optimized_path, log = stage_3_4_onnx_quantize(temp_model_dir, calibration_file.name if onnx_quant_type == "Static" and calibration_file else None)
             options = {'pipeline_type': 'ONNX', 'prune': do_prune, 'prune_percent': prune_percent, 'quant_type': onnx_quant_type}
-        
         elif pipeline_type == "GGUF":
-            yield {final_output: "Converting to GGUF (3/5)", log_output: full_log}
+            full_log += "Converting to GGUF...\n"; yield {final_output: "Converting to GGUF (3/5)", log_output: full_log}
             optimized_path, log = stage_3_4_gguf_quantize(temp_model_dir, model_id, gguf_quant_type)
-            full_log += log
             options = {'pipeline_type': 'GGUF', 'prune': do_prune, 'prune_percent': prune_percent, 'quant_type': gguf_quant_type}
-        
         else:
             raise ValueError("Invalid pipeline type selected.")
-
-        yield {final_output: "Packaging & Uploading (4/5)", log_output: full_log}
+        full_log += log
+        
+        full_log += "Packaging & Uploading...\n"; yield {final_output: "Packaging & Uploading (4/5)", log_output: full_log}
         final_message, log = stage_5_package_and_upload(model_id, optimized_path, full_log, options)
         full_log += log
 
-        yield {
-            final_output: gr.update(value="SUCCESS", label="Status"),
-            log_output: full_log,
-            success_box: gr.Markdown(f"✅ **Success!** Your optimized model is available here: [{repo_id_for_link}](https://huggingface.co/{repo_id_for_link})", visible=True),
-            run_button: gr.Button(interactive=True, value="Run Optimization Pipeline", variant="primary"),
-            analyze_button: gr.Button(interactive=True, value="Analyze Model")
-        }
-
+        yield {final_output: gr.update(value="SUCCESS", label="Status"), log_output: full_log, success_box: gr.Markdown(f"✅ **Success!** Model available: [{repo_id_for_link}](https://huggingface.co/{repo_id_for_link})", visible=True), run_button: gr.Button(interactive=True, value="Run Optimization Pipeline", variant="primary"), analyze_button: gr.Button(interactive=True, value="Analyze Model")}
     except Exception as e:
         logging.error(f"AMOP Pipeline failed. Error: {e}", exc_info=True)
         full_log += f"\n[ERROR] Pipeline failed: {e}"
-        yield {
-            final_output: gr.update(value="ERROR", label="Status"),
-            log_output: full_log,
-            success_box: gr.Markdown(f"❌ **An error occurred.** Check the logs for details.", visible=True),
-            run_button: gr.Button(interactive=True, value="Run Optimization Pipeline", variant="primary"),
-            analyze_button: gr.Button(interactive=True, value="Analyze Model")
-        }
+        yield {final_output: gr.update(value="ERROR", label="Status"), log_output: full_log, success_box: gr.Markdown(f"❌ **An error occurred.** Check logs for details.", visible=True), run_button: gr.Button(interactive=True, value="Run Optimization Pipeline", variant="primary"), analyze_button: gr.Button(interactive=True, value="Analyze Model")}
     finally:
         if temp_model_dir and os.path.exists(temp_model_dir):
             shutil.rmtree(temp_model_dir)
-            logging.info(f"Cleaned up temporary directory: {temp_model_dir}")
-
 
+# --- GRADIO UI ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🚀 AMOP: Adaptive Model Optimization Pipeline")
-    gr.Markdown("Turn any Hugging Face Hub model into a CPU-optimized version using ONNX or GGUF.")
-    
-    if not HF_TOKEN:
-        gr.Warning("You have not set your HF_TOKEN in the Space secrets! The final 'upload' step will be skipped. Please add a secret with the key `HF_TOKEN` and your Hugging Face write token as the value.")
-
+    if not HF_TOKEN: gr.Warning("HF_TOKEN not set! The final 'upload' step will be skipped.")
     with gr.Row():
         with gr.Column(scale=1):
             gr.Markdown("### 1. Select a Model")
-            model_id_input = gr.Textbox(
-                label="Hugging Face Model ID", 
-                placeholder="e.g., gpt2, google/gemma-2b",
-            )
+            model_id_input = gr.Textbox(label="Hugging Face Model ID", placeholder="e.g., gpt2, google/gemma-2b")
             analyze_button = gr.Button("🔍 Analyze Model", variant="secondary")
-            
             with gr.Accordion("⚙️ 2. Configure Optimization", open=False) as optimization_accordion:
                 analysis_report_output = gr.Markdown()
-                
-                pipeline_type_radio = gr.Radio(
-                    ["ONNX", "GGUF"], label="Select Optimization Pipeline", info="GGUF is recommended for LLMs, ONNX for others."
-                )
-                
-                # Unified Pruning controls, shown/hidden by parent group
-                prune_checkbox = gr.Checkbox(label="Enable Pruning", value=False, info="Removes redundant weights from the model.")
-                prune_slider = gr.Slider(minimum=0, maximum=90, value=20, step=5, label="Pruning Percentage (%)")
-                
+                pipeline_type_radio = gr.Radio(["ONNX", "GGUF"], label="Select Optimization Pipeline")
+                prune_checkbox = gr.Checkbox(label="Enable Pruning", value=False, info="Removes redundant weights.", visible=True)
+                prune_slider = gr.Slider(minimum=0, maximum=90, value=20, step=5, label="Pruning Percentage (%)", visible=True)
                 with gr.Group(visible=False) as onnx_options:
-                    gr.Markdown("#### ONNX Quantization Options")
-                    onnx_quant_radio = gr.Radio(["Dynamic", "Static"], label="ONNX Quantization Type", value="Dynamic", info="Static may offer better performance but requires calibration data.")
+                    gr.Markdown("#### ONNX Options")
+                    onnx_quant_radio = gr.Radio(["Dynamic", "Static"], label="Quantization Type", value="Dynamic")
                     calibration_file_upload = gr.File(label="Upload Calibration Data (.txt)", visible=False, file_types=['.txt'])
-
                 with gr.Group(visible=False) as gguf_options:
-                    gr.Markdown("#### GGUF Quantization Options")
-                    gguf_quant_dropdown = gr.Dropdown(
-                        ["q4_k_m", "q5_k_m", "q8_0", "f16"],
-                        label="GGUF Quantization Strategy",
-                        value="q4_k_m",
-                        info="q4_k_m is a good balance of size and quality."
-                    )
-                
+                    gr.Markdown("#### GGUF Options")
+                    gguf_quant_dropdown = gr.Dropdown(["q4_k_m", "q5_k_m", "q8_0", "f16"], label="Quantization Strategy", value="q4_k_m")
                 run_button = gr.Button("🚀 Run Optimization Pipeline", variant="primary")
-        
         with gr.Column(scale=2):
             gr.Markdown("### Pipeline Status & Logs")
-            final_output = gr.Label(value="Idle", label="Status", show_label=True)
+            final_output = gr.Label(value="Idle", label="Status")
             success_box = gr.Markdown(visible=False)
-            log_output = gr.Textbox(label="Live Logs", lines=20, interactive=False, max_lines=20)
+            log_output = gr.Textbox(label="Live Logs", lines=20, interactive=False)
 
     def update_ui_for_pipeline(pipeline_type):
-        is_onnx = pipeline_type == "ONNX"
-        is_gguf = pipeline_type == "GGUF"
-        # Pruning controls are visible for either pipeline type, but grouped logically
-        return {
-            onnx_options: gr.Group(visible=is_onnx),
-            gguf_options: gr.Group(visible=is_gguf),
-            prune_checkbox: gr.Checkbox(visible=is_onnx or is_gguf),
-            prune_slider: gr.Slider(visible=is_onnx or is_gguf)
-        }
-    
+        return {onnx_options: gr.Group(visible=pipeline_type=="ONNX"), gguf_options: gr.Group(visible=pipeline_type=="GGUF")}
     def update_ui_for_quant_type(quant_type):
         return gr.File(visible=quant_type == "Static")
 
-    pipeline_type_radio.change(fn=update_ui_for_pipeline, inputs=pipeline_type_radio, outputs=[onnx_options, gguf_options, prune_checkbox, prune_slider])
+    pipeline_type_radio.change(fn=update_ui_for_pipeline, inputs=pipeline_type_radio, outputs=[onnx_options, gguf_options])
     onnx_quant_radio.change(fn=update_ui_for_quant_type, inputs=onnx_quant_radio, outputs=[calibration_file_upload])
-
-    analyze_button.click(
-        fn=stage_1_analyze_model,
-        inputs=[model_id_input],
-        outputs=[log_output, analysis_report_output, optimization_accordion]
-    )
-    
-    run_button.click(
-        fn=run_amop_pipeline,
-        inputs=[
-            model_id_input, pipeline_type_radio, 
-            prune_checkbox, prune_slider,
-            onnx_quant_radio, calibration_file_upload, gguf_quant_dropdown
-        ],
-        outputs=[run_button, analyze_button, final_output, log_output, success_box]
-    )
+    analyze_button.click(fn=stage_1_analyze_model, inputs=[model_id_input], outputs=[log_output, analysis_report_output, optimization_accordion])
+    run_button.click(fn=run_amop_pipeline, inputs=[model_id_input, pipeline_type_radio, prune_checkbox, prune_slider, onnx_quant_radio, calibration_file_upload, gguf_quant_dropdown], outputs=[run_button, analyze_button, final_output, log_output, success_box])
 
 if __name__ == "__main__":
     demo.launch(debug=True)