Spaces:
Running on Zero
Running on Zero
| """Ad-hoc check: the image downloader's own logic (paths, filters, pre-flight). | |
| Separate from `_selftest_transfer.py` (which exercises the shared engine) | |
| because these are the image-specific decisions: where the snapshot lands, | |
| which files are wanted, and whether this host can run the result. | |
| python scripts/_selftest_fetch_image.py | |
| """ | |
| import contextlib | |
| import io | |
| import shutil | |
| import sys | |
| import tempfile | |
| from pathlib import Path | |
| sys.path.insert(0, str(Path(__file__).resolve().parent)) | |
| sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) | |
| import fetch_image_model as f # noqa: E402 | |
| ok = True | |
| def check(label, condition): | |
| global ok | |
| ok = ok and bool(condition) | |
| print(f" [{'PASS' if condition else 'FAIL'}] {label}") | |
| print("1. the snapshot lands at the expected absolute path") | |
| target = f.resolve_target() | |
| expected = Path(__file__).resolve().parent.parent / "models" / "flux2-klein-4b" | |
| check(f"resolve_target() == {expected}", target == expected) | |
| check("target is absolute", target.is_absolute()) | |
| check("repo is the FLUX.2 Klein snapshot", f.resolve_repo() == f.DEFAULT_REPO) | |
| check( | |
| "repo is black-forest-labs/FLUX.2-klein-4B", | |
| f.DEFAULT_REPO == "black-forest-labs/FLUX.2-klein-4B", | |
| ) | |
| print("2. the Diffusers components are wanted, the extras are not") | |
| for rel in ("README.md", | |
| "LICENSE.md", | |
| "realism.jpg", | |
| "editing.jpg", | |
| "others.jpg", | |
| "flux-2-klein-4b.safetensors", | |
| "flux-2-klein-4b.workflow.json"): | |
| check(f"{rel} is ignored", not f.transfer._wanted(rel, f.ALLOW_PATTERNS, f.IGNORE_PATTERNS)) | |
| for rel in ("model_index.json", | |
| "scheduler/scheduler_config.json", | |
| "tokenizer/tokenizer.json", | |
| "tokenizer/chat_template.jinja", | |
| "text_encoder/model-00001-of-00002.safetensors", | |
| "text_encoder/model-00002-of-00002.safetensors", | |
| "text_encoder/model.safetensors.index.json", | |
| "transformer/diffusion_pytorch_model.safetensors", | |
| "vae/diffusion_pytorch_model.safetensors"): | |
| check(f"{rel} is still wanted", | |
| f.transfer._wanted(rel, f.ALLOW_PATTERNS, f.IGNORE_PATTERNS)) | |
| print("3. nothing REQUIRED_FILES needs is filtered out") | |
| for rel in f.REQUIRED_FILES: | |
| check(f"required {rel}", | |
| f.transfer._wanted(rel, f.ALLOW_PATTERNS, f.IGNORE_PATTERNS)) | |
| print("4. reclaim_unusable_partials only drops un-resumable state") | |
| work = Path(tempfile.mkdtemp()) | |
| try: | |
| orphan = work / "unet" / "diffusion_pytorch_model.fp16.safetensors.part" | |
| orphan.parent.mkdir(parents=True) | |
| orphan.write_bytes(b"x" * 1000) # wrong size, no ledger -> dead weight | |
| resumable = work / "vae" / "diffusion_pytorch_model.fp16.safetensors.part" | |
| resumable.parent.mkdir(parents=True) | |
| resumable.write_bytes(b"y" * 500) | |
| ledger = resumable.with_name(resumable.name[: -len(".part")] + ".parts.json") | |
| ledger.write_text('{"repo":"x","revision":"y","size":500,"done":[[0,500]]}') | |
| manifest = {"files": { | |
| "unet/diffusion_pytorch_model.fp16.safetensors": 9999, | |
| "vae/diffusion_pytorch_model.fp16.safetensors": 500, | |
| }} | |
| freed = f.reclaim_unusable_partials(work, manifest) | |
| check("ledger-less orphan .part reclaimed", not orphan.exists()) | |
| check(f"its bytes were freed ({freed})", freed == 1000) | |
| check("resumable .part kept", resumable.exists()) | |
| check("its resume ledger kept", ledger.exists()) | |
| finally: | |
| shutil.rmtree(work, ignore_errors=True) | |
| print("5. pre-flight reports this host honestly (no download is started)") | |
| print(f" HF token found: {'yes' if f._hub_token() else 'no'}") | |
| problems = f.check_hardware() | |
| print(f" hardware problems on this machine: {len(problems)}") | |
| for p in problems: | |
| print(f" - {p.splitlines()[0]}") | |
| print("6. an incomplete-snapshot report names each missing file once") | |
| labelled = f.dedupe([ | |
| "missing: model_index.json", | |
| "missing: vae/config.json", | |
| "missing: model_index.json", | |
| "empty (0 bytes): vae/config.json", | |
| ]) | |
| check("duplicates dropped, order kept", labelled == [ | |
| "missing: model_index.json", | |
| "missing: vae/config.json", | |
| "empty (0 bytes): vae/config.json", | |
| ]) | |
| report_buffer = io.StringIO() | |
| with contextlib.redirect_stdout(report_buffer): | |
| code = f.reject(Path("models/flux2-klein-4b"), ["missing: model_index.json"]) | |
| report = report_buffer.getvalue() | |
| check("reject() fails the run", code == 1) | |
| check("the report names the folder", "flux2-klein-4b" in report) | |
| check("the report names the missing file", "missing: model_index.json" in report) | |
| check("the report names the fix", "fetch_image_model.py" in report) | |
| print("\nRESULT:", "ALL CHECKS PASSED" if ok else "FAILURES PRESENT") | |
| raise SystemExit(0 if ok else 1) | |