Update app.py

app.py

@@ -181,14 +181,14 @@ def stage_5_package_and_upload(model_id: str, optimized_model_path: str, pipelin
 
 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:
-        yield {log_output: "Please enter a Model ID.", final_output: gr.Label(value="Idle", label="Status")}
+        yield {log_output: "Please enter a Model ID.", final_output: "Idle"}
         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: gr.Label(value={"label": f"RUNNING ({pipeline_type})"}, show_label=True),
+        final_output: f"RUNNING ({pipeline_type})",
         log_output: initial_log
     }
 
@@ -196,16 +196,19 @@ def run_amop_pipeline(model_id: str, pipeline_type: str, do_prune: bool, prune_p
     temp_model_dir = None
     try:
         repo_name_suffix = f"-amop-cpu-{pipeline_type.lower()}"
-        repo_id_for_link = f"{api.whoami()['name']}/{model_id.split('/')[-1]}{repo_name_suffix}"
+        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 pipeline_type == "ONNX":
             full_log += "Loading base model for pruning...\n"
-            yield {final_output: gr.Label(value="Loading model (1/5)"), log_output: full_log}
+            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: gr.Label(value="Pruning model (2/5)"), log_output: full_log}
+            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
@@ -217,7 +220,7 @@ def run_amop_pipeline(model_id: str, pipeline_type: str, do_prune: bool, prune_p
             tokenizer.save_pretrained(temp_model_dir)
             full_log += f"Saved intermediate model to temporary directory: {temp_model_dir}\n"
             
-            yield {final_output: gr.Label(value="Converting to ONNX (3/5)"), log_output: full_log}
+            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
@@ -225,7 +228,7 @@ def run_amop_pipeline(model_id: str, pipeline_type: str, do_prune: bool, prune_p
 
         elif pipeline_type == "GGUF":
             full_log += "[STAGE 1 & 2] Loading and Pruning are skipped for GGUF pipeline.\n"
-            yield {final_output: gr.Label(value="Converting to GGUF (3/5)"), log_output: full_log}
+            yield {final_output: "Converting to GGUF (3/5)", log_output: full_log}
             optimized_path, log = stage_3_4_gguf_quantize(model_id, gguf_quant_type)
             full_log += log
             options = {'pipeline_type': 'GGUF', 'quant_type': gguf_quant_type}
@@ -233,12 +236,12 @@ def run_amop_pipeline(model_id: str, pipeline_type: str, do_prune: bool, prune_p
         else:
             raise ValueError("Invalid pipeline type selected.")
 
-        yield {final_output: gr.Label(value="Packaging & Uploading (4/5)"), log_output: full_log}
+        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.Label(value="SUCCESS", label="Status"),
+            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"),
@@ -249,7 +252,7 @@ def run_amop_pipeline(model_id: str, pipeline_type: str, do_prune: bool, prune_p
         logging.error(f"AMOP Pipeline failed. Error: {e}", exc_info=True)
         full_log += f"\n[ERROR] Pipeline failed: {e}"
         yield {
-            final_output: gr.Label(value="ERROR", label="Status"),
+            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"),
@@ -273,7 +276,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             gr.Markdown("### 1. Select a Model")
             model_id_input = gr.Textbox(
                 label="Hugging Face Model ID", 
-                placeholder="e.g., gpt2, meta-llama/Llama-2-7b-chat-hf",
+                placeholder="e.g., gpt2, google/gemma-2b",
             )
             analyze_button = gr.Button("🔍 Analyze Model", variant="secondary")