Spaces:
Running on Zero
Running on Zero
| """Checkpoint catalogue and cached downloads for the ZeroGPU Space.""" | |
| import hashlib | |
| import os | |
| import tempfile | |
| from pathlib import Path | |
| import requests | |
| SDXL_BASE = "SDXL Base 1.0" | |
| RAGNAROK = "Juggernaut XL - Ragnarok" | |
| MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0" | |
| RAGNAROK_FILENAME = "juggernautXL_ragnarokBy.safetensors" | |
| RAGNAROK_URL = "https://civitai.com/api/download/models/1759168" | |
| RAGNAROK_SHA256 = "dd08fa32f98d05a2443ca1419e46df1575a0811f6e3b246d9dd47ff20f5eb66a" | |
| RAGNAROK_SIZE = 7105350162 | |
| MODEL_CHOICES = (RAGNAROK,) | |
| def model_cache_dir(): | |
| """Use an attached persistent volume when available, otherwise local disk.""" | |
| configured = os.getenv("RF_MODEL_CACHE_DIR") | |
| if configured: | |
| return Path(configured) | |
| if Path("/data").is_dir(): | |
| return Path("/data/ruinedfooocus/checkpoints") | |
| return Path(__file__).resolve().parent / "model_cache" / "checkpoints" | |
| def file_sha256(path): | |
| digest = hashlib.sha256() | |
| with Path(path).open("rb") as stream: | |
| for chunk in iter(lambda: stream.read(4 * 1024 * 1024), b""): | |
| digest.update(chunk) | |
| return digest.hexdigest() | |
| _verified_paths = set() | |
| def download_verified_file(target, url, expected_sha256, progress=None, cached_size=None): | |
| """Cache one remote file, accepting it only when its checksum matches.""" | |
| target = Path(target) | |
| target.parent.mkdir(parents=True, exist_ok=True) | |
| if target.is_file(): | |
| if (str(target) in _verified_paths | |
| or (cached_size is not None and target.stat().st_size == cached_size) | |
| or file_sha256(target) == expected_sha256): | |
| _verified_paths.add(str(target)) | |
| return str(target) | |
| target.unlink() | |
| partial = target.with_suffix(target.suffix + ".part") | |
| headers = {} | |
| if os.getenv("CIVITAI_TOKEN"): | |
| headers["Authorization"] = f"Bearer {os.environ['CIVITAI_TOKEN']}" | |
| try: | |
| with requests.get(url, headers=headers, stream=True, timeout=(15, 120)) as response: | |
| response.raise_for_status() | |
| total = int(response.headers.get("content-length", 0)) | |
| digest = hashlib.sha256() | |
| downloaded = 0 | |
| with partial.open("wb") as stream: | |
| for chunk in response.iter_content(chunk_size=4 * 1024 * 1024): | |
| if not chunk: | |
| continue | |
| stream.write(chunk) | |
| digest.update(chunk) | |
| downloaded += len(chunk) | |
| if progress and total: | |
| progress(min(downloaded / total, 1), desc=f"Downloading {target.name}") | |
| if digest.hexdigest() != expected_sha256: | |
| raise ValueError(f"{target.name} download failed SHA-256 verification") | |
| partial.replace(target) | |
| _verified_paths.add(str(target)) | |
| return str(target) | |
| finally: | |
| partial.unlink(missing_ok=True) | |
| def stage_checkpoint(source, target, expected_sha256): | |
| """Copy a bucket checkpoint to local disk and verify it in one sequential read.""" | |
| source = Path(source) | |
| target = Path(target) | |
| target.parent.mkdir(parents=True, exist_ok=True) | |
| if target.is_file() and target.stat().st_size == source.stat().st_size: | |
| if file_sha256(target) == expected_sha256: | |
| return str(target) | |
| partial = target.with_suffix(target.suffix + ".part") | |
| digest = hashlib.sha256() | |
| try: | |
| with source.open("rb") as input_file, partial.open("wb") as output_file: | |
| for chunk in iter(lambda: input_file.read(16 * 1024 * 1024), b""): | |
| output_file.write(chunk) | |
| digest.update(chunk) | |
| if digest.hexdigest() != expected_sha256: | |
| raise ValueError(f"{source.name} cache failed SHA-256 verification") | |
| partial.replace(target) | |
| return str(target) | |
| finally: | |
| partial.unlink(missing_ok=True) | |
| def prepare_model(model_name, progress=None): | |
| """Return a local checkpoint path, downloading Ragnarok only if needed.""" | |
| if model_name == SDXL_BASE: | |
| return "" | |
| if model_name != RAGNAROK: | |
| raise ValueError(f"Unknown model: {model_name}") | |
| cached_path = download_verified_file( | |
| model_cache_dir() / RAGNAROK_FILENAME, | |
| RAGNAROK_URL, | |
| RAGNAROK_SHA256, | |
| progress=progress, | |
| cached_size=RAGNAROK_SIZE if Path("/data").is_dir() else None, | |
| ) | |
| if Path("/data").is_dir() and not os.getenv("RF_MODEL_CACHE_DIR"): | |
| local_path = Path(tempfile.gettempdir()) / "ruinedfooocus" / RAGNAROK_FILENAME | |
| try: | |
| return stage_checkpoint(cached_path, local_path, RAGNAROK_SHA256) | |
| except ValueError: | |
| Path(cached_path).unlink() | |
| _verified_paths.discard(cached_path) | |
| cached_path = download_verified_file( | |
| model_cache_dir() / RAGNAROK_FILENAME, | |
| RAGNAROK_URL, | |
| RAGNAROK_SHA256, | |
| progress=progress, | |
| ) | |
| return stage_checkpoint(cached_path, local_path, RAGNAROK_SHA256) | |
| return cached_path | |