diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 95a84de10c..e66216a7c5 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -985,18 +985,36 @@ class VideoBackend: return None def _cache_bytes(self, repo_id: Optional[str]) -> int: - """Bytes of ``repo_id`` currently in the HF blob cache (progress polling).""" + """Bytes of ``repo_id`` currently in the HF blob cache (progress polling). + + Walks the repo's cache directory directly instead of ``scan_cache_dir``: + the scanner skips in-flight ``*.incomplete`` blobs, so during a multi-GB + shard pull the counter would freeze at the last completed blob for minutes + while the disk keeps filling (the bar sat stuck mid-download).""" if not repo_id: return 0 try: - from huggingface_hub import scan_cache_dir - rid = repo_id.strip() - for repo in scan_cache_dir().repos: - if repo.repo_id == rid: - return int(repo.size_on_disk) + import os + + from huggingface_hub.constants import HF_HUB_CACHE + + folder = Path(HF_HUB_CACHE) / ("models--" + repo_id.strip().replace("/", "--")) + if not folder.is_dir(): + return 0 + total = 0 + for root, _dirs, files in os.walk(folder): + for name in files: + try: + path = os.path.join(root, name) + # Snapshot entries are symlinks into blobs/; skip them so a + # blob is not counted twice. + if not os.path.islink(path): + total += os.path.getsize(path) + except OSError: + continue + return int(total) except Exception: # noqa: BLE001 -- cache scan is best-effort return 0 - return 0 def load_progress(self) -> dict[str, Any]: """Phase + downloaded/total bytes for the in-flight load (cache-scan based).""" @@ -2436,6 +2454,14 @@ class VideoBackend: if self._generate_job_active: gen["active"] = True gen.setdefault("active", False) + # Mirror the image endpoint's field names (total_steps / fraction) alongside the + # native "total": the two generate-progress APIs used to disagree, so a client + # polling the image shape against video read total_steps=null / fraction=0 while + # the step counter advanced. + total = int(gen.get("total") or 0) + step = int(gen.get("step") or 0) + gen["total_steps"] = total + gen["fraction"] = min(1.0, step / total) if total > 0 else 0.0 return gen def cancel_generate(self) -> bool: diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 1cb2a1a9a1..692f9cf1fa 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -2620,6 +2620,9 @@ class VideoGenerateProgressResponse(BaseModel): ) step: int = Field(0, description = "Denoising steps completed so far") total: int = Field(0, description = "Total denoising steps for this run") + # Image-endpoint-compatible aliases so one poller works against both APIs. + total_steps: int = Field(0, description = "Total denoising steps (alias of total)") + fraction: float = Field(0.0, description = "step / total, clamped to [0,1]") eta_seconds: Optional[float] = Field(None, description = "Estimated seconds remaining") video: Optional[GalleryVideo] = Field( None, description = "Saved gallery record when phase is 'completed'" diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index 5782ad8e1b..3dc98bf800 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -1063,10 +1063,49 @@ def test_generate_without_load_raises(fake_runtime): def test_generate_progress_and_cancel_idle(fake_runtime): backend = VideoBackend() - assert backend.generate_progress() == {"active": False} + # Idle shape carries the image-endpoint-compatible aliases (total_steps / fraction) + # so one poller works against both generate-progress APIs. + assert backend.generate_progress() == { + "active": False, + "total_steps": 0, + "fraction": 0.0, + } assert backend.cancel_generate() is False +def test_generate_progress_derives_total_steps_and_fraction(fake_runtime): + # A mid-denoise poll must report fraction = step / total under BOTH field names: + # a client polling the image API's shape against video used to read + # total_steps=null / fraction=0 while step advanced. + backend = VideoBackend() + backend._gen = {"active": True, "phase": "denoise", "step": 5, "total": 20} + gen = backend.generate_progress() + assert gen["total"] == 20 and gen["total_steps"] == 20 + assert gen["step"] == 5 and gen["fraction"] == 0.25 + + +def test_cache_bytes_counts_incomplete_blobs(fake_runtime, tmp_path, monkeypatch): + # scan_cache_dir skips in-flight *.incomplete blobs, so the old counter froze at the + # last completed blob for the whole multi-GB shard pull. The walk must count both, + # without double-counting snapshot symlinks. + import huggingface_hub.constants as hub_constants + + repo_dir = tmp_path / "models--Wan-AI--Wan2.2-TI2V-5B-Diffusers" + blobs = repo_dir / "blobs" + blobs.mkdir(parents = True) + (blobs / "aa11").write_bytes(b"x" * 1000) # completed blob + (blobs / "bb22.incomplete").write_bytes(b"y" * 500) # in-flight shard + snap = repo_dir / "snapshots" / "deadbeef" + snap.mkdir(parents = True) + (snap / "model_index.json").symlink_to(blobs / "aa11") # must not double-count + monkeypatch.setattr(hub_constants, "HF_HUB_CACHE", str(tmp_path)) + + backend = VideoBackend() + assert backend._cache_bytes("Wan-AI/Wan2.2-TI2V-5B-Diffusers") == 1500 + assert backend._cache_bytes("Wan-AI/absent-repo") == 0 + assert backend._cache_bytes(None) == 0 + + def test_hv15_guider_and_scheduler_progress(fake_runtime): # HunyuanVideo-1.5: no guidance kwarg (CFG set on the guider), no step # callback (progress via the scheduler.step wrapper, restored afterwards).