Fix repo_id

ui.py

@@ -1,5 +1,6 @@
 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
 
@@ -12,22 +13,49 @@ class UIController:
     """
 
     @staticmethod
-    def init_session(profile: Optional[gr.OAuthProfile] = None) -> Tuple[Any, ...]:
+    def init_session(profile: Optional[gr.OAuthProfile] = None, token: Optional[gr.OAuthToken] = 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, push_update, zip_update = UIController.update_hub_interactive(new_engine, username)
+        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']
+        )
 
         return (
             new_engine,
             new_engine.get_tools_json(),
             new_engine.config.MODEL_NAME,
             f"Ready. (Session {new_engine.session_id})",
-            repo_update, 
-            push_update,
-            zip_update, 
+            repo_update,        # Model Name Input
+            final_org_update,   # Owner Dropdown
+            push_update,        # Push Button
+            zip_update,         # Zip Button
             username
         )
 
@@ -76,23 +104,29 @@ class UIController:
         return gr.update(value=None, visible=False)
 
     @staticmethod
-    def upload_model(engine: FunctionGemmaEngine, repo_name: str, oauth_token: Optional[gr.OAuthToken]) -> str:
+    def upload_model(engine: FunctionGemmaEngine, owner: str, model_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 repo_name:
-            return "❌ Error: Please enter a repository name."
+        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()}"
         
         return engine.upload_model_to_hub(
-            repo_name=repo_name, 
+            repo_name=full_repo_id, 
             oauth_token=oauth_token.token, 
         )
 
     @staticmethod
-    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}`**"
+    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}`**"
         
     @staticmethod
     def update_hub_interactive(engine: Optional[FunctionGemmaEngine], username: Optional[str] = None):
@@ -100,9 +134,10 @@ class UIController:
         has_model_tuned = engine is not None and getattr(engine, 'has_model_tuned', False)
         
         return (
-            gr.update(interactive=is_logged_in), 
-            gr.update(interactive=is_logged_in and has_model_tuned),
-            gr.update(interactive=has_model_tuned)
+            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
         )
 
 # --- View / Layout Layer ---
@@ -190,17 +225,23 @@ 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.")
-                repo_name_input = gr.Textbox(
-                    label="Target Repository Name", value="functiongemma-270m-it-tuning-lab", placeholder="e.g., functiongemma-270m-it-tuned", interactive=False
-                )
+                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
+                    )
+                    
                 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": [repo_name_input, push_to_hub_btn, repo_id_preview, upload_status]
+        "hub_controls": [org_dropdown, model_name_input, push_to_hub_btn, repo_id_preview, upload_status]
     }
 
 # --- Main Build Function ---
@@ -218,19 +259,22 @@ 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] 
         
-        repo_input = export_ui["hub_controls"][0]
-        push_btn = export_ui["hub_controls"][1]
+        # Hub Controls
+        org_dropdown = export_ui["hub_controls"][0]
+        model_name_input = export_ui["hub_controls"][1]
+        push_btn = export_ui["hub_controls"][2]
         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), gr.update(interactive=False)]
+                   [gr.update(interactive=False) for _ in hub_inputs]
         
         def unlock_ui():
             """Unlocks general action buttons only. Zip/Push are handled by update_hub_interactive"""
@@ -239,23 +283,25 @@ def build_interface() -> gr.Blocks:
         # --- Event Wiring ---
 
         # 1. Initialization
-        demo.load(lock_ui, outputs=action_buttons + [push_btn, zip_btn]).then(
+        # Note: init_session now returns updates for both model_name_input and org_dropdown
+        demo.load(lock_ui, outputs=action_buttons + hub_inputs).then(
             fn=UIController.init_session,
-            inputs=None,
+            inputs=None, # Gradio automatically injects profile and token
             outputs=[
                 engine_state, 
                 data_ui["tools_editor"], 
                 train_ui["model_input"], 
                 train_ui["outputs"][0], # log output
-                repo_input, 
+                model_name_input,       # Text update
+                org_dropdown,           # Dropdown update (choices+val)
                 push_btn, 
-                zip_btn,                # Update Zip state based on initial engine state
+                zip_btn,
                 username_state
             ]
         ).then(
             fn=UIController.update_repo_preview,
-            inputs=[username_state, repo_input],
-            outputs=[export_ui["hub_controls"][2]] 
+            inputs=[org_dropdown, model_name_input],
+            outputs=[export_ui["hub_controls"][3]] 
         ).then(unlock_ui, outputs=action_buttons)
 
         # 2. Data Tab
@@ -279,10 +325,10 @@ def build_interface() -> gr.Blocks:
                 gr.update(visible=False), 
                 gr.update(interactive=False), # Lock Reload
                 gr.update(interactive=False), # Lock Eval
-                gr.update(interactive=False), # Lock Zip 
+                [gr.update(interactive=False) for _ in hub_inputs], # Lock Hub
                 gr.update(visible=True)       # Show Stop
             ),
-            outputs=[run_btn, reload_btn, eval_btn, zip_btn, stop_btn]
+            outputs=[run_btn, reload_btn, eval_btn, *hub_inputs, stop_btn]
         )
         train_run_event = train_run_event.then(
             fn=UIController.run_training,
@@ -300,7 +346,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=[repo_input, push_btn, zip_btn] 
+            outputs=hub_inputs
         )
 
         # 3b. Evaluation
@@ -308,7 +354,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 (optional, or lock)
+                gr.update(visible=False),     # Hide self
                 gr.update(visible=True)       # Show Stop
             ),
             outputs=[run_btn, reload_btn, eval_btn, stop_btn]
@@ -335,41 +381,47 @@ def build_interface() -> gr.Blocks:
             queue=False
         )
 
-        reload_btn.click(lock_ui, outputs=action_buttons + [push_btn, zip_btn]).then(
+        reload_btn.click(lock_ui, outputs=action_buttons + hub_inputs).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=[repo_input, push_btn, zip_btn]
+            outputs=hub_inputs
         )
 
         # 4. Export Tab
-        zip_btn.click(lock_ui, outputs=action_buttons + [push_btn, zip_btn]).then(
+        zip_btn.click(lock_ui, outputs=action_buttons + hub_inputs).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=[repo_input, push_btn, zip_btn]
+            outputs=hub_inputs
         )
 
-        repo_input.change(
+        # 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(
             fn=UIController.update_repo_preview,
-            inputs=[username_state, repo_input],
-            outputs=[export_ui["hub_controls"][2]]
+            inputs=[org_dropdown, model_name_input],
+            outputs=[export_ui["hub_controls"][3]]
         )
 
-        push_btn.click(lock_ui, outputs=action_buttons + [push_btn, zip_btn]).then(
+        push_btn.click(lock_ui, outputs=action_buttons + hub_inputs).then(
             fn=UIController.upload_model,
-            inputs=[engine_state, repo_input],
-            outputs=[export_ui["hub_controls"][3]]
+            inputs=[engine_state, org_dropdown, model_name_input], # oauth_token injected
+            outputs=[export_ui["hub_controls"][4]]
         ).then(unlock_ui, outputs=action_buttons).then(
             fn=UIController.update_hub_interactive,
             inputs=[engine_state, username_state],
-            outputs=[repo_input, push_btn, zip_btn]
+            outputs=hub_inputs
         )
 
     return demo