Spaces:
Paused
Paused
| import hashlib | |
| import json | |
| import os | |
| import shutil | |
| import signal | |
| import subprocess | |
| import time | |
| from pathlib import Path | |
| from huggingface_hub import hf_hub_download | |
| CONFIG = json.loads(Path(__file__).with_name("models.json").read_text()) | |
| BINARY = os.getenv("SD_CLI", "/usr/local/bin/sd-cli") | |
| def download(spec): | |
| return hf_hub_download(repo_id=spec["repo"], filename=spec["file"], | |
| revision=spec.get("revision", "main"), | |
| token=os.getenv("HF_TOKEN") or None) | |
| def diagnostics(): | |
| pieces = ["CPU-only backend; model weights load on first request."] | |
| commit = Path("/engine-commit.txt") | |
| if commit.exists(): | |
| pieces.append("Engine commit: " + commit.read_text().strip()) | |
| for name in ("/sys/fs/cgroup/memory.max", "/sys/fs/cgroup/memory/memory.limit_in_bytes"): | |
| p = Path(name) | |
| if p.exists(): | |
| pieces.append("Container memory limit (bytes): " + p.read_text().strip()) | |
| break | |
| if shutil.which(BINARY): | |
| result = subprocess.run([BINARY, "--help"], capture_output=True, text=True, timeout=30) | |
| pieces.append((result.stdout + result.stderr)[-18000:]) | |
| else: | |
| pieces.append("sd-cli missing: build the Docker image before running inference.") | |
| return "\n".join(pieces) | |
| def run_one(image, output, prompt, negative, seed, steps, cfg, width, height, | |
| adapters, jobdir, status): | |
| help_result = subprocess.run([BINARY, "--help"], capture_output=True, text=True, timeout=30) | |
| help_text = help_result.stdout + help_result.stderr | |
| models = {} | |
| for name, spec in CONFIG.items(): | |
| status(f"Downloading/checking {name}: {spec['file']}") | |
| models[name] = download(spec) | |
| lora_dir = jobdir / "loras" | |
| lora_dir.mkdir(exist_ok=True) | |
| tags = [] | |
| for index, (spec, strength) in enumerate(adapters): | |
| path = Path(download(spec)) | |
| name = f"adapter_{index}" | |
| target = lora_dir / (name + path.suffix) | |
| if not target.exists(): | |
| target.symlink_to(path.resolve()) | |
| tags.append(f"<lora:{name}:{strength:.4f}>") | |
| command = [BINARY, "--diffusion-model", models["diffusion"], "--vae", models["vae"], | |
| "--llm", models["llm"], "--llm_vision", models["vision"], | |
| "-r", str(image), "-p", prompt + " " + " ".join(tags), | |
| "-n", negative, "--seed", str(seed), "--steps", str(steps), | |
| "--cfg-scale", str(cfg), "--sampling-method", "euler", "--flow-shift", "3", | |
| "-W", str(width), "-H", str(height), "--threads", os.getenv("CPU_THREADS", "2"), | |
| "-o", str(output)] | |
| for flag in ("--mmap", "--vae-tiling", "--diffusion-fa"): | |
| if flag in help_text: | |
| command.append(flag) | |
| if tags: | |
| if "--lora-model-dir" not in help_text: | |
| raise RuntimeError("This engine version does not expose --lora-model-dir") | |
| command += ["--lora-model-dir", str(lora_dir)] | |
| log = output.with_suffix(".log") | |
| timeout = int(os.getenv("INFERENCE_TIMEOUT", "21600")) | |
| status("CPU inference running. This can take a long time; see container logs.") | |
| with log.open("w") as stream: | |
| process = subprocess.Popen(command, stdout=stream, stderr=subprocess.STDOUT, start_new_session=True) | |
| start = time.monotonic() | |
| try: | |
| while process.poll() is None: | |
| if time.monotonic() - start > timeout: | |
| raise TimeoutError(f"CPU inference exceeded {timeout} seconds") | |
| time.sleep(2) | |
| if process.returncode or not output.exists(): | |
| tail = log.read_text(errors="replace")[-7000:] | |
| if process.returncode == -9: | |
| tail = "Process killed (possible container RAM exhaustion). " + tail | |
| raise RuntimeError(f"sd-cli exit={process.returncode}\n{tail}") | |
| finally: | |
| if process.poll() is None: | |
| os.killpg(process.pid, signal.SIGTERM) | |
| try: | |
| process.wait(timeout=10) | |
| except subprocess.TimeoutExpired: | |
| os.killpg(process.pid, signal.SIGKILL) | |
| process.wait() | |
| return log | |