mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-10-29 06:04:23 +08:00
259 lines
10 KiB
Python
259 lines
10 KiB
Python
"""
|
|
Download screens.
|
|
"""
|
|
|
|
from pathlib import Path
|
|
from typing import Optional, List, Dict, Any
|
|
|
|
from textual.app import ComposeResult
|
|
from textual.containers import Container, Horizontal, Vertical, ScrollableContainer, Center
|
|
from textual.widgets import Button, Checkbox, DataTable, Input, Label, RichLog, Select, Static
|
|
from textual.screen import ModalScreen
|
|
from textual.worker import Worker, WorkerState
|
|
|
|
|
|
class DownloadScreen(ModalScreen):
|
|
"""Modal screen for downloading models"""
|
|
|
|
CSS_PATH = str(Path(__file__).parent.parent / "styles" / "download.tcss")
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.repo_type = "huggingface"
|
|
self.verified = False
|
|
self.files = []
|
|
self.total_size = 0
|
|
|
|
def compose(self) -> ComposeResult:
|
|
from kt_kernel.cli.config.settings import get_settings
|
|
|
|
settings = get_settings()
|
|
model_paths = settings.get_model_paths()
|
|
default_dir = str(model_paths[0]) if model_paths else ""
|
|
|
|
with Container(id="download-dialog"):
|
|
yield Label("[bold cyan]Download Model[/bold cyan]")
|
|
yield Label("")
|
|
|
|
with ScrollableContainer(id="download-content"):
|
|
yield Static("", id="download-status")
|
|
|
|
with Vertical(id="download-params"):
|
|
# Repository type selection
|
|
yield Label("[bold]Repository Source:[/bold]")
|
|
with Horizontal():
|
|
yield Button("HuggingFace", id="btn-hf", variant="primary")
|
|
yield Button("ModelScope", id="btn-ms")
|
|
|
|
# Repository ID input
|
|
yield Label("[bold]Repository ID:[/bold]")
|
|
yield Input(placeholder="e.g., deepseek-ai/DeepSeek-V3", id="input-repo")
|
|
|
|
# File pattern input
|
|
yield Label("[bold]File Pattern:[/bold]")
|
|
yield Input(
|
|
value="*", placeholder="* for all files, or *.safetensors, *.gguf, etc.", id="input-pattern"
|
|
)
|
|
|
|
# Model name input (added)
|
|
yield Label("[bold]Model Name:[/bold]")
|
|
yield Input(placeholder="Auto-filled after verification", id="input-model-name")
|
|
|
|
# Save path input (added)
|
|
yield Label("[bold]Save Path:[/bold]")
|
|
yield Input(value=default_dir, placeholder="Path to save the model", id="input-save-path")
|
|
|
|
yield Label("")
|
|
with Horizontal(id="download-buttons"):
|
|
yield Button("Verify Repository", id="btn-verify", variant="primary")
|
|
yield Button("Download", id="btn-download", variant="success", disabled=True)
|
|
yield Button("Cancel", id="btn-cancel")
|
|
|
|
def on_button_pressed(self, event: Button.Pressed) -> None:
|
|
if event.button.id == "btn-hf":
|
|
self.repo_type = "huggingface"
|
|
event.button.variant = "primary"
|
|
self.query_one("#btn-ms", Button).variant = "default"
|
|
self.verified = False
|
|
self.query_one("#btn-download", Button).disabled = True
|
|
elif event.button.id == "btn-ms":
|
|
self.repo_type = "modelscope"
|
|
event.button.variant = "primary"
|
|
self.query_one("#btn-hf", Button).variant = "default"
|
|
self.verified = False
|
|
self.query_one("#btn-download", Button).disabled = True
|
|
elif event.button.id == "btn-verify":
|
|
self._verify_repository()
|
|
elif event.button.id == "btn-download":
|
|
self._start_download()
|
|
elif event.button.id == "btn-cancel":
|
|
self.dismiss(None)
|
|
|
|
def _verify_repository(self) -> None:
|
|
"""Verify repository and list files"""
|
|
repo_input = self.query_one("#input-repo", Input)
|
|
pattern_input = self.query_one("#input-pattern", Input)
|
|
status_widget = self.query_one("#download-status", Static)
|
|
|
|
repo_id = repo_input.value.strip()
|
|
pattern = pattern_input.value.strip() or "*"
|
|
|
|
# Validate repo_id
|
|
if not repo_id:
|
|
status_widget.update("[red]Please enter a repository ID[/red]")
|
|
return
|
|
|
|
if "/" not in repo_id:
|
|
status_widget.update("[red]Repository ID must be in format: namespace/model-name[/red]")
|
|
return
|
|
|
|
# Show verifying message
|
|
status_widget.update(f"[yellow]Verifying repository: {repo_id}...[/yellow]")
|
|
self.query_one("#btn-verify", Button).disabled = True
|
|
|
|
# Run verification in worker
|
|
from kt_kernel.cli.utils.download_helper import (
|
|
verify_repo_exists,
|
|
list_remote_files_hf,
|
|
list_remote_files_ms,
|
|
filter_files_by_pattern,
|
|
calculate_total_size,
|
|
)
|
|
from kt_kernel.cli.utils.model_verifier import check_huggingface_connectivity
|
|
|
|
try:
|
|
# Check connectivity for HuggingFace
|
|
use_mirror = False
|
|
if self.repo_type == "huggingface":
|
|
is_accessible, msg = check_huggingface_connectivity(timeout=5)
|
|
if not is_accessible:
|
|
use_mirror = True
|
|
status_widget.update("[yellow]HuggingFace not accessible, using mirror...[/yellow]")
|
|
|
|
# Verify repository exists
|
|
exists, msg = verify_repo_exists(repo_id, self.repo_type, use_mirror)
|
|
if not exists:
|
|
status_widget.update(f"[red]Repository not found: {msg}[/red]")
|
|
self.query_one("#btn-verify", Button).disabled = False
|
|
return
|
|
|
|
# List files
|
|
if self.repo_type == "huggingface":
|
|
all_files = list_remote_files_hf(repo_id, use_mirror)
|
|
else:
|
|
all_files = list_remote_files_ms(repo_id)
|
|
|
|
# Filter by pattern
|
|
filtered_files = filter_files_by_pattern(all_files, pattern)
|
|
|
|
if not filtered_files:
|
|
status_widget.update(f"[yellow]No files match pattern '{pattern}'[/yellow]")
|
|
self.query_one("#btn-verify", Button).disabled = False
|
|
return
|
|
|
|
# Calculate total size
|
|
self.total_size = calculate_total_size(filtered_files)
|
|
self.files = filtered_files
|
|
|
|
# Auto-fill model name (default: last part of repo_id)
|
|
model_name = repo_id.split("/")[-1]
|
|
model_name_input = self.query_one("#input-model-name", Input)
|
|
if not model_name_input.value:
|
|
model_name_input.value = model_name
|
|
|
|
# Update save path with model name
|
|
save_path_input = self.query_one("#input-save-path", Input)
|
|
from pathlib import Path
|
|
|
|
current_path = Path(save_path_input.value or "")
|
|
|
|
# Determine base directory
|
|
if current_path.is_dir():
|
|
# It's an existing directory, use it as base
|
|
base_dir = current_path
|
|
elif current_path.parent.exists():
|
|
# Parent exists, use parent as base (current path might be a full path with model name)
|
|
base_dir = current_path.parent
|
|
else:
|
|
# Use the path as-is (might be a default value)
|
|
base_dir = current_path if current_path.is_absolute() else Path(save_path_input.value or "").parent
|
|
|
|
# Build full path with model name
|
|
if base_dir:
|
|
full_path = base_dir / model_name
|
|
# Check if path already exists
|
|
if full_path.exists():
|
|
# Find a unique name
|
|
counter = 1
|
|
while (base_dir / f"{model_name}-{counter}").exists():
|
|
counter += 1
|
|
full_path = base_dir / f"{model_name}-{counter}"
|
|
|
|
save_path_input.value = str(full_path)
|
|
|
|
# Format size
|
|
size_gb = self.total_size / (1024**3)
|
|
|
|
# Show results
|
|
file_count = len(filtered_files)
|
|
status_text = f"[green]✓ Repository found![/green]\n\n"
|
|
status_text += f"Files to download: {file_count}\n"
|
|
status_text += f"Total size: {size_gb:.2f} GB\n\n"
|
|
|
|
# Show first few files
|
|
status_text += "Files:\n"
|
|
for i, f in enumerate(filtered_files[:10]):
|
|
file_size_mb = f["size"] / (1024**2)
|
|
status_text += f" • {f['path']} ({file_size_mb:.1f} MB)\n"
|
|
|
|
if len(filtered_files) > 10:
|
|
status_text += f" ... and {len(filtered_files) - 10} more files"
|
|
|
|
status_widget.update(status_text)
|
|
|
|
# Enable download button
|
|
self.verified = True
|
|
self.query_one("#btn-download", Button).disabled = False
|
|
self.query_one("#btn-verify", Button).disabled = False
|
|
|
|
except Exception as e:
|
|
status_widget.update(f"[red]Verification failed: {e}[/red]")
|
|
self.query_one("#btn-verify", Button).disabled = False
|
|
|
|
def _start_download(self) -> None:
|
|
"""Start download process"""
|
|
if not self.verified:
|
|
return
|
|
|
|
repo_input = self.query_one("#input-repo", Input)
|
|
pattern_input = self.query_one("#input-pattern", Input)
|
|
model_name_input = self.query_one("#input-model-name", Input)
|
|
save_path_input = self.query_one("#input-save-path", Input)
|
|
|
|
model_name = model_name_input.value.strip()
|
|
save_path = save_path_input.value.strip()
|
|
|
|
if not model_name:
|
|
model_name = repo_input.value.strip().split("/")[-1]
|
|
|
|
if not save_path:
|
|
from kt_kernel.cli.config.settings import get_settings
|
|
|
|
settings = get_settings()
|
|
model_paths = settings.get_model_paths()
|
|
from pathlib import Path
|
|
|
|
save_path = str(Path(model_paths[0]) / model_name) if model_paths else model_name
|
|
|
|
config = {
|
|
"repo_type": self.repo_type,
|
|
"repo_id": repo_input.value.strip(),
|
|
"pattern": pattern_input.value.strip() or "*",
|
|
"model_name": model_name,
|
|
"save_path": save_path,
|
|
"files": self.files,
|
|
"total_size": self.total_size,
|
|
}
|
|
|
|
self.dismiss(config)
|