Update app.py

app.py

@@ -14,7 +14,6 @@ def prune_to_single_layer(model_id: str):
     status_lines = []
     status_lines.append(f"Loading base model: {model_id}")
     
-    # Initialize as None to prevent UnboundLocalError in the finally block
     model = None
     tokenizer = None
         
@@ -65,8 +64,8 @@ def prune_to_single_layer(model_id: str):
                 outputs = model.generate(
                     **inputs.to(model.device),
                     max_new_tokens=40,
-                    do_sample=False, # Temperature removed to prevent conflict
-                    pad_token_id=tokenizer.eos_token_id # Prevents warnings/crashes on Llama
+                    do_sample=False, 
+                    pad_token_id=tokenizer.eos_token_id 
                 )
             text = tokenizer.decode(outputs[0], skip_special_tokens=True)
             status_lines.append("\nQuick generation test (should be at least semi-coherent):")
@@ -86,7 +85,6 @@ def prune_to_single_layer(model_id: str):
         return "\n".join(status_lines) + f"\n\n❌ Failed: {err_msg}"
         
     finally:
-        # Safely try to free memory even on failure
         if model is not None:
             del model
         if tokenizer is not None:
@@ -98,7 +96,8 @@ def prune_to_single_layer(model_id: str):
 # ────────────────────────────────────────────────
 CSS = """.gradio-container { max-width: 780px !important; }"""
 
-with gr.Blocks(title="Minimal Single-Layer Pruner", css=CSS, theme=gr.themes.Default()) as demo:
+# FIX 1: Removed css and theme from Blocks()
+with gr.Blocks(title="Minimal Single-Layer Pruner") as demo:
     gr.Markdown("""
     # Single-Layer Pruner (test version)
     Loads a small model → keeps **only the last layer** → shows result + quick generation test.
@@ -117,11 +116,11 @@ with gr.Blocks(title="Minimal Single-Layer Pruner", css=CSS, theme=gr.themes.Def
         value="Qwen/Qwen2.5-0.5B-Instruct"
     )
         
+    # FIX 2: Removed show_copy_button=True
     status = gr.Textbox(
         label="Pruning log",
         lines=18,
-        interactive=False,
-        show_copy_button=True
+        interactive=False
     )
         
     btn = gr.Button("Prune to 1 layer → Test", variant="primary")
@@ -140,4 +139,5 @@ with gr.Blocks(title="Minimal Single-Layer Pruner", css=CSS, theme=gr.themes.Def
     """)
 
 if __name__ == "__main__":
-    demo.launch()
+    # FIX 3: Moved css and theme into launch()
+    demo.launch(css=CSS, theme=gr.themes.Default())