diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 6e4eb14054..22188975fb 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -92,6 +92,7 @@ from .video_families import ( supported_video_family_names, ) from utils.hardware import clear_gpu_cache + # Shared with the image backend so both pin every loader call to the same live cache root. from core.inference.diffusion import hub_cache_dir @@ -751,9 +752,7 @@ class VideoBackend: if gguf_filename and not Path(repo_id).expanduser().exists(): info = api.model_info(repo_id, files_metadata = True) size = sum( - int(s.size or 0) - for s in (info.siblings or []) - if s.rfilename == gguf_filename + int(s.size or 0) for s in (info.siblings or []) if s.rfilename == gguf_filename ) total += size entries.append( diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index e5c6bf37a6..ed40e9c3fe 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -1682,10 +1682,7 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto # on embeddings and latent stats produced by the other model -- incompatible # shapes at best, silently wrong conditioning at worst. from .diffusion_train_extras import source_revision # noqa: PLC0415 - - namespace = ( - f"{spec.family}_{cfg.base_model}_{source_revision(cfg.base_model)}" - ) + namespace = f"{spec.family}_{cfg.base_model}_{source_revision(cfg.base_model)}" pcache = PersistentConditioningCache(cfg.cond_cache_dir, namespace, cfg.resolution) except Exception as exc: # noqa: BLE001 -- the cache is an optimisation, never fatal _emit(on_event, "warning", message = f"conditioning cache disabled: {exc}") diff --git a/studio/backend/core/training/diffusion_train_extras.py b/studio/backend/core/training/diffusion_train_extras.py index 501167296b..a25e45198d 100644 --- a/studio/backend/core/training/diffusion_train_extras.py +++ b/studio/backend/core/training/diffusion_train_extras.py @@ -183,7 +183,6 @@ def source_revision(ref: Any) -> str: not load them. """ import os # noqa: PLC0415 — keep the module import list light for the subprocess - try: name = str(ref or "").strip() if not name: @@ -193,7 +192,8 @@ def source_revision(ref: Any) -> str: roots = [name] with os.scandir(name) as it: roots += [ - e.path for e in it + e.path + for e in it if e.is_dir() and e.name.startswith(("text_encoder", "tokenizer")) ] for root in roots: diff --git a/studio/backend/hub/services/models/downloads.py b/studio/backend/hub/services/models/downloads.py index 0f7dd89499..4439c2337b 100644 --- a/studio/backend/hub/services/models/downloads.py +++ b/studio/backend/hub/services/models/downloads.py @@ -164,9 +164,7 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional detail = "scope_id and gguf_variant are mutually exclusive.", ) if not scoped_files: - raise HTTPException( - status_code = 400, detail = "scope_id requires a non-empty files list." - ) + raise HTTPException(status_code = 400, detail = "scope_id requires a non-empty files list.") if not _is_valid_gguf_variant(scope_variant): raise HTTPException(status_code = 400, detail = f"Invalid scope_id: {body.scope_id!r}") variant = scope_variant diff --git a/studio/backend/hub/workers/hf_download.py b/studio/backend/hub/workers/hf_download.py index 994eb68707..debc552351 100644 --- a/studio/backend/hub/workers/hf_download.py +++ b/studio/backend/hub/workers/hf_download.py @@ -690,8 +690,7 @@ def _download_scoped_snapshot( info = _model_info_with_retry(repo_id, hf_token) except Exception as e: print( - f"metadata unavailable for scoped download of {repo_id} " - f"({type(e).__name__}: {e})", + f"metadata unavailable for scoped download of {repo_id} " f"({type(e).__name__}: {e})", file = sys.stderr, ) info = None @@ -709,6 +708,7 @@ def _download_scoped_snapshot( for s in siblings ] from hub.utils.snapshot_filters import blob_hashes_for_siblings + blob_hashes = blob_hashes_for_siblings(siblings) download_manifest.write_manifest("model", repo_id, scope, expected_files, mode) @@ -821,7 +821,6 @@ def main() -> None: scoped_files: list[str] = [] if args.files_json: import json - try: with open(args.files_json, encoding = "utf-8") as handle: scoped_files = [str(f) for f in json.load(handle)] diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 5866e4472e..57f7783fd2 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -3755,11 +3755,20 @@ def test_generate_keeps_a_scalar_negative_prompt_off_the_list_paths(fake_runtime class _TracingPipe(_CountingPipe): """Appends ``("call", n)`` to a shared trace so resets can be interleaved with forwards.""" - def __init__(self, trace, max_images = None): + def __init__( + self, + trace, + max_images = None, + ): super().__init__(max_images = max_images) self.trace = trace - def __call__(self, *, prompt = None, **kwargs): + def __call__( + self, + *, + prompt = None, + **kwargs, + ): n = kwargs.get("num_images_per_prompt", 1) if isinstance(prompt, list): n *= len(prompt) @@ -3776,9 +3785,7 @@ def test_generate_resets_the_step_cache_before_an_oom_retry(fake_runtime, tmp_pa backend = _load_zimage_backend(tmp_path) trace: list = [] pipe = _TracingPipe(trace, max_images = 2) - pipe.transformer = types.SimpleNamespace( - _reset_stateful_cache = lambda: trace.append(("reset",)) - ) + pipe.transformer = types.SimpleNamespace(_reset_stateful_cache = lambda: trace.append(("reset",))) object.__setattr__(backend._state, "pipe", pipe) object.__setattr__(backend._state, "transformer_cache", "fbcache") out = backend.generate(prompt = "p", seeds = [1, 2, 3, 4]) @@ -3799,13 +3806,13 @@ def test_generate_resets_the_step_cache_before_every_chunk(fake_runtime, tmp_pat backend = _load_zimage_backend(tmp_path) trace: list = [] pipe = _TracingPipe(trace) - pipe.transformer = types.SimpleNamespace( - _reset_stateful_cache = lambda: trace.append(("reset",)) - ) + pipe.transformer = types.SimpleNamespace(_reset_stateful_cache = lambda: trace.append(("reset",))) object.__setattr__(backend._state, "pipe", pipe) object.__setattr__(backend._state, "transformer_cache", "fbcache") backend.generate(prompt = "p", seeds = [1, 2, 3], batch_size = 2) assert trace == [("reset",), ("call", 2), ("reset",), ("call", 1)] + + class _FakeSibling: def __init__(self, rfilename, size): self.rfilename = rfilename @@ -3817,7 +3824,7 @@ class _FakeInfo: self.siblings = siblings -GB = 1024 ** 3 +GB = 1024**3 # A FLUX-shaped base repo: the packaged root single and the transformer shards are what a # plain snapshot_download would drag in and the loader never opens. _FLUX_BASE_SIBLINGS = [ @@ -3834,8 +3841,14 @@ _FLUX_BASE_SIBLINGS = [ def _fake_hf_api(monkeypatch, repos): """Point HfApi.model_info at a canned sibling list per repo id.""" + class _Api: - def model_info(self, repo_id, files_metadata = False, token = None): + def model_info( + self, + repo_id, + files_metadata = False, + token = None, + ): return _FakeInfo(repos[repo_id]) monkeypatch.setattr("huggingface_hub.HfApi", lambda *a, **k: _Api()) diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index 53844f4ea2..0eebe4e5a6 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -170,15 +170,14 @@ def test_every_train_base_is_deployable_as_an_inference_pipeline(): # it -- which is what happened to both FLUX.2 families, trusted for training only. from core.inference.diffusion import _is_trusted_diffusion_repo from core.inference.diffusion_families import _FAMILIES - for fam in _FAMILIES: if not fam.trainable: continue for base in fam.train_base_repos: deploy_base = fam.deploy_base_repo or base - assert _is_trusted_diffusion_repo(deploy_base), ( - f"{fam.name}: deploy base {deploy_base!r} is not loadable for inference" - ) + assert _is_trusted_diffusion_repo( + deploy_base + ), f"{fam.name}: deploy base {deploy_base!r} is not loadable for inference" def test_gated_access_requires_token(): diff --git a/studio/backend/tests/test_diffusion_prequant.py b/studio/backend/tests/test_diffusion_prequant.py index 52a00070b0..aa67756e94 100644 --- a/studio/backend/tests/test_diffusion_prequant.py +++ b/studio/backend/tests/test_diffusion_prequant.py @@ -413,7 +413,6 @@ def test_load_exclude_tokens_need_the_recorded_family(monkeypatch, tmp_path): # and accept only the family-aware set. Pins the offline builder # (scripts/build_prequant_checkpoint.py) to exclude_tokens_for_scheme(scheme, fam.name). from core.inference.diffusion_transformer_quant import exclude_tokens_for_scheme - for family in ("qwen-image", "qwen-image-edit"): family_less = _good_ckpt(scheme = "int8") family_less["metadata"]["family"] = family diff --git a/studio/backend/tests/test_hf_cache_settings.py b/studio/backend/tests/test_hf_cache_settings.py index 4997c6609f..a27cea00f5 100644 --- a/studio/backend/tests/test_hf_cache_settings.py +++ b/studio/backend/tests/test_hf_cache_settings.py @@ -319,6 +319,6 @@ def test_diffusion_loader_calls_pin_the_cache_dir(): if call not in line: continue window = "\n".join(source.splitlines()[index - 1 : index + 8]) - assert "cache_dir" in window or "kwargs" in window, ( - f"{rel}:{index} calls {call} without a pinned cache_dir" - ) + assert ( + "cache_dir" in window or "kwargs" in window + ), f"{rel}:{index} calls {call} without a pinned cache_dir" diff --git a/studio/backend/tests/test_scoped_download_job.py b/studio/backend/tests/test_scoped_download_job.py index 0644e7ed50..0951f02723 100644 --- a/studio/backend/tests/test_scoped_download_job.py +++ b/studio/backend/tests/test_scoped_download_job.py @@ -115,9 +115,7 @@ def test_scoped_files_survive_into_the_registry(monkeypatch): monkeypatch.setattr(dl, "resolve_cached_repo_id_case", lambda repo, **k: repo) monkeypatch.setattr(dl, "scoped_file_blob_hashes", lambda *a, **k: frozenset()) monkeypatch.setattr(dl._registry, "claim", _spy_claim) - monkeypatch.setattr( - download_lifecycle, "launch_worker", lambda *a, **k: "running" - ) + monkeypatch.setattr(download_lifecycle, "launch_worker", lambda *a, **k: "running") asyncio.run(dl.download_model_response(_request())) assert captured["scoped_files"] == FILES diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py index 9285b72ff7..c13e0851e7 100644 --- a/tests/studio/test_model_picker_contracts.py +++ b/tests/studio/test_model_picker_contracts.py @@ -585,9 +585,7 @@ def test_variant_expander_forwards_the_gguf_filename(): by filename and cannot map a quant label back to one, so without it every hub GGUF pick on Images/Video fell through to a silent return and nothing loaded.""" src = _read("features/model-picker/components/model-selector/pickers.tsx") - handler = re.search( - r"const handleVariantClick = useCallback\(.*?\n \);", src, re.S - ) + handler = re.search(r"const handleVariantClick = useCallback\(.*?\n \);", src, re.S) assert handler, "handleVariantClick not found" assert "ggufFilename: filename," in handler.group(0) # The call site has to actually pass it through. @@ -648,9 +646,7 @@ def test_local_model_sections_respect_the_task_filter(): for memo in ("sortedLmStudio", "sortedLocalDir", "sortedCustomFolderModels"): block = re.search(rf"const {memo} = useMemo\(.*?\n \);", src, re.S) assert block, f"{memo} not found" - assert "passesTaskGate(m.task" in block.group(0), ( - f"{memo} does not apply the task gate" - ) + assert "passesTaskGate(m.task" in block.group(0), f"{memo} does not apply the task gate" def test_chat_picker_routes_diffusion_picks_to_their_page():