Spaces:
Running
Running
| import re | |
| from typing import Optional | |
| import gradio as gr | |
| from huggingface_hub import HfApi | |
| from .config import RUN_LOCALLY | |
| from .hf_utils import extract_quantization, get_gguf_files_from_repo, extract_username | |
| from .processing import process_split_request | |
| def load_gguf_files( | |
| source_repo_id: str, oauth_token: Optional[gr.OAuthToken] | |
| ) -> gr.Dropdown: | |
| """Load GGUF files from selected repository""" | |
| if not source_repo_id: | |
| return gr.Dropdown(choices=[], value=None, interactive=False) | |
| try: | |
| if RUN_LOCALLY: | |
| api = HfApi() | |
| else: | |
| token = oauth_token.token if oauth_token else None | |
| api = HfApi(token=token) | |
| gguf_files = get_gguf_files_from_repo(source_repo_id, api) | |
| if not gguf_files: | |
| return gr.Dropdown( | |
| choices=["No GGUF files found in this repository"], | |
| value=None, | |
| interactive=True, | |
| ) | |
| return gr.Dropdown( | |
| choices=gguf_files, | |
| value=gguf_files[0] if gguf_files else None, | |
| interactive=True, | |
| ) | |
| except Exception as e: | |
| return gr.Dropdown(choices=[f"Error: {str(e)}"], value=None, interactive=True) | |
| def create_ui() -> gr.Blocks: | |
| """Create and return the Gradio UI""" | |
| with gr.Blocks() as demo: | |
| gr.Markdown( | |
| value="# GGUF Splitter", | |
| elem_classes=["main-header"], | |
| ) | |
| if not RUN_LOCALLY: | |
| with gr.Group(): | |
| login_message = gr.Markdown( | |
| value="### Connect your Hugging Face account to split and upload GGUF models.\n\nDue to platform storage limitations, split files must be uploaded directly to your account after processing - they cannot be stored for later download.", | |
| elem_classes=["login-message"] | |
| ) | |
| login_btn = gr.LoginButton(size="lg") | |
| else: | |
| login_message = gr.Markdown(visible=False) | |
| login_btn = None | |
| gr.Markdown("### Step 1: Select Source Model") | |
| with gr.Group(): | |
| repo_selector = gr.Textbox( | |
| label="Model Repository", | |
| placeholder="Enter a Hugging Face model ID (e.g., Qwen/Qwen3-0.6B-GGUF)", | |
| info="Enter the full repository ID from Hugging Face", | |
| ) | |
| load_files_btn = gr.Button("Load GGUF Files", variant="secondary") | |
| gguf_selector = gr.Dropdown( | |
| label="GGUF File", | |
| choices=[], | |
| value=None, | |
| interactive=False, | |
| info="Select the GGUF file you want to split into shards", | |
| ) | |
| gr.Markdown("### Step 2: Configure Output") | |
| with gr.Group(): | |
| if not RUN_LOCALLY: | |
| repo_name = gr.Textbox( | |
| label="Target Repository Name", | |
| info="The sharded model will be uploaded to this repository", | |
| interactive=False, | |
| ) | |
| with gr.Row(): | |
| public_toggle = gr.Checkbox( | |
| label="Public Repository", | |
| value=True, | |
| info="Uncheck to create a private repository", | |
| interactive=False, | |
| scale=1, | |
| ) | |
| else: | |
| repo_name = gr.Textbox( | |
| visible=False, | |
| interactive=False, | |
| ) | |
| public_toggle = gr.Checkbox( | |
| visible=False, | |
| interactive=False, | |
| ) | |
| gr.Markdown("### Step 3: Split & " + ("Upload" if not RUN_LOCALLY else "Save")) | |
| process_btn = gr.Button( | |
| "Split and " + ("Upload GGUF" if not RUN_LOCALLY else "Save GGUF"), | |
| variant="primary", | |
| interactive=False, | |
| size="lg", | |
| elem_classes=["action-btn"], | |
| ) | |
| with gr.Accordion("Output Log", open=True): | |
| output_display = gr.Markdown( | |
| value="*Waiting for action...*", elem_classes=["status-box"] | |
| ) | |
| def update_components_on_login( | |
| profile: Optional[gr.OAuthProfile], oauth_token: Optional[gr.OAuthToken] | |
| ) -> dict: | |
| """Update component visibility and interactivity based on login state""" | |
| if profile and oauth_token: | |
| return { | |
| login_message: gr.Markdown(visible=False), | |
| repo_selector: gr.Textbox(interactive=True), | |
| load_files_btn: gr.Button(interactive=True), | |
| repo_name: gr.Textbox(interactive=True), | |
| gguf_selector: gr.Dropdown(interactive=True), | |
| public_toggle: gr.Checkbox(interactive=True), | |
| process_btn: gr.Button(interactive=True), | |
| } | |
| else: | |
| return { | |
| login_message: gr.Markdown(visible=True), | |
| repo_selector: gr.Textbox(interactive=False), | |
| load_files_btn: gr.Button(interactive=False), | |
| repo_name: gr.Textbox(interactive=False), | |
| gguf_selector: gr.Dropdown(interactive=False), | |
| public_toggle: gr.Checkbox(interactive=False), | |
| process_btn: gr.Button(interactive=False), | |
| } | |
| if login_btn is not None: | |
| demo.load( | |
| fn=update_components_on_login, | |
| inputs=[], | |
| outputs=[ | |
| login_message, | |
| repo_selector, | |
| load_files_btn, | |
| repo_name, | |
| gguf_selector, | |
| public_toggle, | |
| process_btn, | |
| ], | |
| ) | |
| else: | |
| demo.load( | |
| fn=lambda: { | |
| login_message: gr.Markdown(visible=False), | |
| repo_selector: gr.Textbox(interactive=True), | |
| load_files_btn: gr.Button(interactive=True), | |
| repo_name: gr.Textbox(interactive=True), | |
| gguf_selector: gr.Dropdown(interactive=True), | |
| public_toggle: gr.Checkbox(interactive=True), | |
| process_btn: gr.Button(interactive=True), | |
| }, | |
| outputs=[ | |
| login_message, | |
| repo_selector, | |
| load_files_btn, | |
| repo_name, | |
| gguf_selector, | |
| public_toggle, | |
| process_btn, | |
| ], | |
| ) | |
| def on_repo_selected(repo_id: str, oauth_token: Optional[gr.OAuthToken]): | |
| """Load GGUF files when a repository is selected""" | |
| return load_gguf_files(repo_id, oauth_token) | |
| def update_repo_name_on_selection( | |
| source_repo_id: str, | |
| gguf_filename: str, | |
| oauth_token: Optional[gr.OAuthToken], | |
| ): | |
| """Update repository name based on selected source repo and GGUF file""" | |
| if not source_repo_id or not gguf_filename: | |
| return gr.Textbox(value="") | |
| if not oauth_token: | |
| return gr.Textbox( | |
| value="", info="Sign in to auto-generate repository name" | |
| ) | |
| try: | |
| api = HfApi(token=oauth_token.token) | |
| user_info = api.whoami() | |
| username = extract_username(user_info) | |
| if not username: | |
| return gr.Textbox( | |
| value="", info="Unable to determine your Hugging Face username" | |
| ) | |
| model_name = source_repo_id.split("/")[-1] | |
| model_name = re.sub(r"-?GGUF$", "", model_name, flags=re.IGNORECASE) | |
| quantization = extract_quantization(gguf_filename) | |
| suffix = f"gguf-sharded-{quantization}-{model_name}" | |
| default_name = f"{username}/{suffix}" | |
| return gr.Textbox(value=default_name) | |
| except Exception: | |
| return gr.Textbox( | |
| value="", info="Unable to auto-generate repository name" | |
| ) | |
| load_files_btn.click( | |
| fn=on_repo_selected, inputs=[repo_selector], outputs=gguf_selector | |
| ) | |
| gguf_selector.change( | |
| fn=update_repo_name_on_selection, | |
| inputs=[repo_selector, gguf_selector], | |
| outputs=repo_name, | |
| ) | |
| process_btn.click( | |
| fn=process_split_request, | |
| inputs=[ | |
| repo_name, | |
| repo_selector, | |
| gguf_selector, | |
| public_toggle, | |
| output_display, | |
| ], | |
| outputs=output_display, | |
| ) | |
| return demo | |