优化 Runtime 安装进度体验
- 状态栏与窗口强制刷新,进度更新无需移动鼠标即可实时显示 - 安装过程拆分为 12 个阶段,每段进度从 0 到 100% - 安装阶段解析 uv 下载事件推进进度,模型下载阶段实时反映字节进度
This commit is contained in:
@@ -32,10 +32,23 @@ def check_cancelled(cancel_check):
|
||||
raise TaskCancelled("Task cancelled")
|
||||
|
||||
|
||||
def write_status(path, progress, message, error=None, state=None):
|
||||
def write_status(
|
||||
path,
|
||||
progress,
|
||||
message,
|
||||
error=None,
|
||||
state=None,
|
||||
stage=None,
|
||||
stages=None,
|
||||
stage_label=None,
|
||||
):
|
||||
payload = {"progress": progress, "message": message}
|
||||
if error:
|
||||
payload["error"] = error
|
||||
if stage is not None:
|
||||
payload["stage"] = stage
|
||||
payload["stages"] = stages
|
||||
payload["stage_label"] = stage_label
|
||||
if state is None:
|
||||
if error:
|
||||
state = "failed"
|
||||
@@ -136,12 +149,27 @@ def prepare_input(input_path, video_frames, temporary_root, cancel_check=None):
|
||||
|
||||
|
||||
class DownloadProgressReporter:
|
||||
def __init__(self, status_path, repository, clock=time.monotonic):
|
||||
def __init__(
|
||||
self,
|
||||
status_path,
|
||||
repository,
|
||||
clock=time.monotonic,
|
||||
start=0.05,
|
||||
span=0.90,
|
||||
stage=None,
|
||||
stages=None,
|
||||
stage_label=None,
|
||||
):
|
||||
self.status_path = status_path
|
||||
self.repository = repository
|
||||
self.clock = clock
|
||||
self.lock = threading.Lock()
|
||||
self.progress = 0.05
|
||||
self.start = start
|
||||
self.span = span
|
||||
self.stage = stage
|
||||
self.stages = stages
|
||||
self.stage_label = stage_label
|
||||
self.progress = start
|
||||
self.last_write_time = 0.0
|
||||
|
||||
def report(self, current, total):
|
||||
@@ -149,7 +177,7 @@ class DownloadProgressReporter:
|
||||
return
|
||||
|
||||
factor = min(max(current / total, 0.0), 1.0)
|
||||
progress = 0.05 + factor * 0.90
|
||||
progress = self.start + factor * self.span
|
||||
now = self.clock()
|
||||
with self.lock:
|
||||
if progress <= self.progress:
|
||||
@@ -162,6 +190,9 @@ class DownloadProgressReporter:
|
||||
self.status_path,
|
||||
progress,
|
||||
f"Downloading {self.repository}",
|
||||
stage=self.stage,
|
||||
stages=self.stages,
|
||||
stage_label=self.stage_label,
|
||||
)
|
||||
|
||||
|
||||
@@ -170,11 +201,25 @@ def create_download_progress_class(
|
||||
repository,
|
||||
tqdm_type=None,
|
||||
clock=time.monotonic,
|
||||
start=0.05,
|
||||
span=0.90,
|
||||
stage=None,
|
||||
stages=None,
|
||||
stage_label=None,
|
||||
):
|
||||
if tqdm_type is None:
|
||||
from tqdm.auto import tqdm as tqdm_type
|
||||
|
||||
reporter = DownloadProgressReporter(status_path, repository, clock)
|
||||
reporter = DownloadProgressReporter(
|
||||
status_path,
|
||||
repository,
|
||||
clock,
|
||||
start,
|
||||
span,
|
||||
stage,
|
||||
stages,
|
||||
stage_label,
|
||||
)
|
||||
|
||||
class DownloadProgress(tqdm_type):
|
||||
def __init__(self, *args, **kwargs):
|
||||
@@ -205,10 +250,26 @@ def download_model(args):
|
||||
|
||||
model_dir = Path(args.model_dir)
|
||||
model_dir.mkdir(parents=True, exist_ok=True)
|
||||
write_status(args.status, 0.05, f"Downloading {args.repository}")
|
||||
staged = getattr(args, "stage", None) is not None
|
||||
start = 0.0 if staged else 0.05
|
||||
span = 1.0 if staged else 0.90
|
||||
stage_kwargs = {
|
||||
"stage": getattr(args, "stage", None),
|
||||
"stages": getattr(args, "stages", None),
|
||||
"stage_label": getattr(args, "stage_label", None),
|
||||
}
|
||||
write_status(
|
||||
args.status,
|
||||
start,
|
||||
f"Downloading {args.repository}",
|
||||
**stage_kwargs,
|
||||
)
|
||||
progress_class = create_download_progress_class(
|
||||
args.status,
|
||||
args.repository,
|
||||
start=start,
|
||||
span=span,
|
||||
**stage_kwargs,
|
||||
)
|
||||
snapshot_download(
|
||||
repo_id=args.repository,
|
||||
@@ -216,7 +277,12 @@ def download_model(args):
|
||||
tqdm_class=progress_class,
|
||||
allow_patterns=args.allow_pattern or None,
|
||||
)
|
||||
write_status(args.status, 1.0, "Model is ready")
|
||||
write_status(
|
||||
args.status,
|
||||
1.0,
|
||||
"Model is ready",
|
||||
**stage_kwargs,
|
||||
)
|
||||
|
||||
|
||||
def select_device(requested, torch):
|
||||
@@ -480,6 +546,9 @@ def parser():
|
||||
download.add_argument("--model-dir", required=True)
|
||||
download.add_argument("--status", required=True)
|
||||
download.add_argument("--allow-pattern", action="append", default=[])
|
||||
download.add_argument("--stage", type=int, default=None)
|
||||
download.add_argument("--stages", type=int, default=None)
|
||||
download.add_argument("--stage-label", default=None)
|
||||
|
||||
background = commands.add_parser("remove-background")
|
||||
background.add_argument("--input", required=True)
|
||||
|
||||
Reference in New Issue
Block a user