revert

ui.py

@@ -1,6 +1,5 @@
 import gradio as gr
 from typing import Optional, Tuple, Generator, List, Any
-from huggingface_hub import HfApi
 from config import AppConfig
 from engine import FunctionGemmaEngine
 
@@ -13,49 +12,22 @@ class UIController:
     """
 
     @staticmethod
-    def init_session(profile: Optional[gr.OAuthProfile] = None, token: Optional[gr.OAuthToken] = None) -> Tuple[Any, ...]:
+    def init_session(profile: Optional[gr.OAuthProfile] = None) -> Tuple[Any, ...]:
         config = AppConfig()
         new_engine = FunctionGemmaEngine(config)
         username = profile.username if profile else None
         
-        # Fetch available namespaces (User + Orgs) if logged in
-        namespaces = []
-        default_owner = None
-        
-        if token:
-            try:
-                api = HfApi(token=token.token)
-                user_info = api.whoami()
-                # Add user's own namespace first
-                namespaces.append(user_info['name'])
-                # Add organizations
-                if 'orgs' in user_info:
-                    namespaces.extend([org['name'] for org in user_info['orgs']])
-                
-                default_owner = namespaces[0] if namespaces else None
-            except Exception as e:
-                print(f"Error fetching namespaces: {e}")
-                pass
-
         # Calculate initial interactivity state
-        repo_update, org_interactivity, push_update, zip_update = UIController.update_hub_interactive(new_engine, username)
-        
-        # Combine interactivity with the fetched choices for the Org Dropdown
-        final_org_update = gr.update(
-            choices=namespaces,
-            value=default_owner,
-            interactive=org_interactivity['interactive']
-        )
+        repo_update, push_update, zip_update = UIController.update_hub_interactive(new_engine, username)
 
         return (
             new_engine,
             new_engine.get_tools_json(),
             new_engine.config.MODEL_NAME,
             f"Ready. (Session {new_engine.session_id})",
-            repo_update,        # Model Name Input
-            final_org_update,   # Owner Dropdown
-            push_update,        # Push Button
-            zip_update,         # Zip Button
+            repo_update, 
+            push_update,
+            zip_update, 
             username
         )
 
@@ -104,29 +76,23 @@ class UIController:
         return gr.update(value=None, visible=False)
 
     @staticmethod
-    def upload_model(engine: FunctionGemmaEngine, owner: str, model_name: str, oauth_token: Optional[gr.OAuthToken]) -> str:
+    def upload_model(engine: FunctionGemmaEngine, repo_name: str, oauth_token: Optional[gr.OAuthToken]) -> str:
         if oauth_token is None:
             return "❌ Error: You must log in (top right) to upload models."
-        if not owner:
-            return "❌ Error: Please select an Owner (User or Organization)."
-        if not model_name:
-            return "❌ Error: Please enter a model name."
-        
-        # Construct the full repo_id (e.g. "google/functiongemma-tuned")
-        full_repo_id = f"{owner}/{model_name.strip()}"
+        if not repo_name:
+            return "❌ Error: Please enter a repository name."
         
         return engine.upload_model_to_hub(
-            repo_name=full_repo_id, 
+            repo_name=repo_name, 
             oauth_token=oauth_token.token, 
         )
 
     @staticmethod
-    def update_repo_preview(owner: str, model_name: str) -> str:
-        if not owner:
-            return "Target Repository: (Select Owner First)"
-        
-        clean_name = model_name.strip() if model_name else "..."
-        return f"Target Repository: **`{owner}/{clean_name}`**"
+    def update_repo_preview(username: Optional[str], repo_name: str) -> str:
+        if not username:
+            return "⚠️ Sign in to see the target repository path."
+        clean_repo = repo_name.strip() if repo_name else "..."
+        return f"Target Repository: **`{username}/{clean_repo}`**"
         
     @staticmethod
     def update_hub_interactive(engine: Optional[FunctionGemmaEngine], username: Optional[str] = None):
@@ -134,10 +100,9 @@ class UIController:
         has_model_tuned = engine is not None and getattr(engine, 'has_model_tuned', False)
         
         return (
-            gr.update(interactive=is_logged_in), # Model Name Input
-            gr.update(interactive=is_logged_in), # Owner Dropdown
-            gr.update(interactive=is_logged_in and has_model_tuned), # Push Button
-            gr.update(interactive=has_model_tuned) # Zip Button
+            gr.update(interactive=is_logged_in), 
+            gr.update(interactive=is_logged_in and has_model_tuned),
+            gr.update(interactive=has_model_tuned)
         )
 
 # --- View / Layout Layer ---
@@ -225,23 +190,17 @@ def _render_export_tab(engine_state, username_state):
             
             with gr.Column():
                 gr.Markdown("#### Option B: Save to Hugging Face Hub")
-                gr.Markdown("Publish your fine-tuned model to your personal Hugging Face account or an Organization.")
-                
-                with gr.Row():
-                    org_dropdown = gr.Dropdown(
-                        label="Owner", choices=[], interactive=False, scale=1
-                    )
-                    model_name_input = gr.Textbox(
-                        label="Model Name", value="functiongemma-270m-it-tuning-lab", placeholder="e.g., functiongemma-tuned", interactive=False, scale=2
-                    )
-                    
+                gr.Markdown("Publish your fine-tuned model to your personal Hugging Face account.")
+                repo_name_input = gr.Textbox(
+                    label="Target Repository Name", value="functiongemma-270m-it-tuning-lab", placeholder="e.g., functiongemma-270m-it-tuned", interactive=False
+                )
                 push_to_hub_btn = gr.Button("Save to Hugging Face Hub", variant="secondary", interactive=False)
                 repo_id_preview = gr.Markdown("Target Repository: (Waiting for input...)")
                 upload_status = gr.Markdown("")
 
     return {
         "zip_controls": [zip_btn, download_file],
-        "hub_controls": [org_dropdown, model_name_input, push_to_hub_btn, repo_id_preview, upload_status]
+        "hub_controls": [repo_name_input, push_to_hub_btn, repo_id_preview, upload_status]
     }
 
 # --- Main Build Function ---
@@ -259,22 +218,19 @@ def build_interface() -> gr.Blocks:
             export_ui = _render_export_tab(engine_state, username_state)
 
         # Helpers for UI State
+        # 'action_buttons' now ONLY contains buttons that should always be enabled after a process
+        # Zip and Push buttons are excluded here because their state depends on has_model_tuned
         run_btn, stop_btn, reload_btn, eval_btn = train_ui["buttons"]
         action_buttons = [reload_btn, run_btn, eval_btn] 
         
-        # Hub Controls
-        org_dropdown = export_ui["hub_controls"][0]
-        model_name_input = export_ui["hub_controls"][1]
-        push_btn = export_ui["hub_controls"][2]
+        repo_input = export_ui["hub_controls"][0]
+        push_btn = export_ui["hub_controls"][1]
         zip_btn = export_ui["zip_controls"][0]
 
-        # The list of Hub-related inputs that need updating
-        hub_inputs = [model_name_input, org_dropdown, push_btn, zip_btn]
-
         def lock_ui():
             """Locks all buttons (including Zip/Push) during processing"""
             return [gr.update(interactive=False) for _ in action_buttons] + \
-                   [gr.update(interactive=False) for _ in hub_inputs]
+                   [gr.update(interactive=False), gr.update(interactive=False)]
         
         def unlock_ui():
             """Unlocks general action buttons only. Zip/Push are handled by update_hub_interactive"""
@@ -283,24 +239,23 @@ def build_interface() -> gr.Blocks:
         # --- Event Wiring ---
 
         # 1. Initialization
-        demo.load(lock_ui, outputs=action_buttons + hub_inputs).then(
+        demo.load(lock_ui, outputs=action_buttons + [push_btn, zip_btn]).then(
             fn=UIController.init_session,
-            inputs=None, # Gradio automatically injects profile and token
+            inputs=None,
             outputs=[
                 engine_state, 
                 data_ui["tools_editor"], 
                 train_ui["model_input"], 
                 train_ui["outputs"][0], # log output
-                model_name_input,       # Text update
-                org_dropdown,           # Dropdown update (choices+val)
+                repo_input, 
                 push_btn, 
-                zip_btn,
+                zip_btn,                # Update Zip state based on initial engine state
                 username_state
             ]
         ).then(
             fn=UIController.update_repo_preview,
-            inputs=[org_dropdown, model_name_input],
-            outputs=[export_ui["hub_controls"][3]] 
+            inputs=[username_state, repo_input],
+            outputs=[export_ui["hub_controls"][2]] 
         ).then(unlock_ui, outputs=action_buttons)
 
         # 2. Data Tab
@@ -320,14 +275,14 @@ def build_interface() -> gr.Blocks:
         
         # 3a. Training
         train_run_event = run_btn.click(
-            fn=lambda: [
-                gr.update(visible=False),                     # run_btn
-                gr.update(interactive=False),                 # reload_btn
-                gr.update(interactive=False)                  # eval_btn
-            ] + [gr.update(interactive=False) for _ in hub_inputs] + [ # Unpack Hub Inputs list
-                gr.update(visible=True)                       # stop_btn
-            ],
-            outputs=[run_btn, reload_btn, eval_btn, *hub_inputs, stop_btn]
+            fn=lambda: (
+                gr.update(visible=False), 
+                gr.update(interactive=False), # Lock Reload
+                gr.update(interactive=False), # Lock Eval
+                gr.update(interactive=False), # Lock Zip 
+                gr.update(visible=True)       # Show Stop
+            ),
+            outputs=[run_btn, reload_btn, eval_btn, zip_btn, stop_btn]
         )
         train_run_event = train_run_event.then(
             fn=UIController.run_training,
@@ -345,7 +300,7 @@ def build_interface() -> gr.Blocks:
             # Final check determines if Zip/Push should unlock
             fn=UIController.update_hub_interactive,
             inputs=[engine_state, username_state],
-            outputs=hub_inputs
+            outputs=[repo_input, push_btn, zip_btn] 
         )
 
         # 3b. Evaluation
@@ -353,7 +308,7 @@ def build_interface() -> gr.Blocks:
             fn=lambda: (
                 gr.update(interactive=False), # Lock Run
                 gr.update(interactive=False), # Lock Reload
-                gr.update(visible=False),     # Hide self
+                gr.update(visible=False),     # Hide self (optional, or lock)
                 gr.update(visible=True)       # Show Stop
             ),
             outputs=[run_btn, reload_btn, eval_btn, stop_btn]
@@ -380,47 +335,41 @@ def build_interface() -> gr.Blocks:
             queue=False
         )
 
-        reload_btn.click(lock_ui, outputs=action_buttons + hub_inputs).then(
+        reload_btn.click(lock_ui, outputs=action_buttons + [push_btn, zip_btn]).then(
             fn=UIController.handle_reset,
             inputs=[engine_state, train_ui["model_input"]],
             outputs=[train_ui["outputs"][0]]
         ).then(unlock_ui, outputs=action_buttons).then(
             fn=UIController.update_hub_interactive,
             inputs=[engine_state, username_state],
-            outputs=hub_inputs
+            outputs=[repo_input, push_btn, zip_btn]
         )
 
         # 4. Export Tab
-        zip_btn.click(lock_ui, outputs=action_buttons + hub_inputs).then(
+        zip_btn.click(lock_ui, outputs=action_buttons + [push_btn, zip_btn]).then(
             fn=UIController.zip_model,
             inputs=[engine_state],
             outputs=[export_ui["zip_controls"][1]]
         ).then(unlock_ui, outputs=action_buttons).then(
             fn=UIController.update_hub_interactive,
             inputs=[engine_state, username_state],
-            outputs=hub_inputs
+            outputs=[repo_input, push_btn, zip_btn]
         )
 
-        # Update preview when dropdown or text changes
-        org_dropdown.change(
-            fn=UIController.update_repo_preview,
-            inputs=[org_dropdown, model_name_input],
-            outputs=[export_ui["hub_controls"][3]]
-        )
-        model_name_input.change(
+        repo_input.change(
             fn=UIController.update_repo_preview,
-            inputs=[org_dropdown, model_name_input],
-            outputs=[export_ui["hub_controls"][3]]
+            inputs=[username_state, repo_input],
+            outputs=[export_ui["hub_controls"][2]]
         )
 
-        push_btn.click(lock_ui, outputs=action_buttons + hub_inputs).then(
+        push_btn.click(lock_ui, outputs=action_buttons + [push_btn, zip_btn]).then(
             fn=UIController.upload_model,
-            inputs=[engine_state, org_dropdown, model_name_input], # oauth_token injected
-            outputs=[export_ui["hub_controls"][4]]
+            inputs=[engine_state, repo_input],
+            outputs=[export_ui["hub_controls"][3]]
         ).then(unlock_ui, outputs=action_buttons).then(
             fn=UIController.update_hub_interactive,
             inputs=[engine_state, username_state],
-            outputs=hub_inputs
+            outputs=[repo_input, push_btn, zip_btn]
         )
 
     return demo