Spaces:
Running on Zero
Running on Zero
| """Single-process PhotoGuider runtime for the Hugging Face Space. | |
| The release runs each extractor family in its own environment and process (envs/, workers.py). | |
| Here every model lives in one environment and one process, loaded once and kept resident: the | |
| extractor ops run in-process through the same load()/run(args) interface worker_serve.py uses, and | |
| the explainer VLM is loaded once instead of once per call. On a Space (SPACE_ID set) the explainer | |
| runs in bf16 instead of NF4, which the 48 GB ZeroGPU slice holds without bitsandbytes. | |
| """ | |
| import functools | |
| import importlib | |
| import os | |
| import sys | |
| import types | |
| from pathlib import Path | |
| ROOT = Path(os.environ.get("PHOTOGUIDER_ROOT", Path(__file__).resolve().parent)).resolve() | |
| INF = ROOT / "inference" | |
| LIB = INF / "_lib" | |
| # one environment serves every role; require_env() only checks that the interpreter exists | |
| for _role in ("MAIN", "LAION", "FGAESQ", "SAM", "EXPLAINER"): | |
| os.environ.setdefault(f"CROPPER_ENV_{_role}", sys.executable) | |
| sys.path[:0] = [str(INF), str(LIB), str(LIB / "extractors")] | |
| # FG-IAA ships `utils/` as a namespace package; transparent-background appends a directory holding | |
| # a regular `utils.py` to sys.path, which would shadow it. Bind the name to FG-IAA's package first. | |
| from env_config import asset # noqa: E402 | |
| if "utils" not in sys.modules: | |
| _u = types.ModuleType("utils") | |
| _u.__path__ = [str(Path(asset("fgaesq_source", "envs/fgaesq/FGAesQ_Inference")) / "utils")] | |
| sys.modules["utils"] = _u | |
| import pandas as pd # noqa: E402 | |
| if os.environ.get("SPACE_ID"): | |
| # ZeroGPU: the Gradio process has no GPU and forks a worker per GPU call. transformers 5 asks | |
| # torch.cuda.is_current_stream_capturing() during a CPU forward, which initialises the CUDA | |
| # runtime in that GPU-less parent, and every worker forked afterwards then sees no GPU | |
| # ("No CUDA GPUs are available"). We never capture CUDA graphs, so answer without asking. | |
| import transformers.utils.import_utils as _tiu | |
| _tiu.is_cuda_stream_capturing = lambda: False | |
| import _shared # noqa: E402 | |
| import feature_extractors # noqa: E402 | |
| import feat_saliency as FS # noqa: E402 | |
| import explain_vlm # noqa: E402 | |
| import explain as EX # noqa: E402 | |
| import photoguider # noqa: E402 | |
| class InProcessPool: | |
| """ExtractorPool interface; ops run in this process with their models resident.""" | |
| def __init__(self, verbose=False): | |
| self.mods = {} | |
| def module(self, op): | |
| if op not in self.mods: | |
| if op not in feature_extractors.OP_ENV: | |
| raise ValueError(f"unknown extractor op '{op}'") | |
| m = importlib.import_module(op) | |
| m.load() | |
| self.mods[op] = m | |
| return self.mods[op] | |
| def run_op(self, op, args=None): | |
| self.module(op).run(args) | |
| return True | |
| def close(self): # models stay resident across requests | |
| pass | |
| def load_vlm_bf16(adapter_dir): | |
| """explain_vlm.load_vlm without the NF4 quantization.""" | |
| import torch | |
| from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration | |
| from peft import PeftModel | |
| adapter_dir = Path(adapter_dir) | |
| torch.manual_seed(1234) | |
| proc = AutoProcessor.from_pretrained(explain_vlm.BASE_VLM, revision=explain_vlm.revision("qwen")) | |
| model = Qwen2_5_VLForConditionalGeneration.from_pretrained( | |
| explain_vlm.BASE_VLM, revision=explain_vlm.revision("qwen"), torch_dtype=torch.bfloat16).to("cuda") | |
| model = PeftModel.from_pretrained(model, str(adapter_dir), is_trainable=False).eval() | |
| return {"proc": proc, "model": model, | |
| "bad": explain_vlm.cjk_bad_words_ids(proc.tokenizer, adapter_dir / "cjk_ids_cache.json")} | |
| # ── two speed-ups that leave every number unchanged ── | |
| # 1. The ~1,500 candidate boxes get their hand-crafted features in a per-row Python loop | |
| # (consultant_features.stage_c_features); the rows are independent, so large pools are split | |
| # across forked CPU workers and concatenated back in order. | |
| import multiprocessing # noqa: E402 | |
| import infer # noqa: E402 | |
| _WORKERS = int(os.environ.get("PHOTOGUIDER_WORKERS") or max(1, min(8, (os.cpu_count() or 2) - 1))) # 1 = serial | |
| _serial_features = infer.af_chain_features | |
| def _chunk_features(args): | |
| dfp, maps = args | |
| return infer.M02.stage_c_features(dfp.reset_index(drop=True), maps) | |
| def parallel_features(dfp, vmix_2d, v1_2d, mask_np, min_rows=200): | |
| if len(dfp) < min_rows or _WORKERS == 1: | |
| return _serial_features(dfp, vmix_2d, v1_2d, mask_np) | |
| maps = infer.maps_for(dfp.photo_id.iloc[0], vmix_2d, v1_2d, mask_np) | |
| dfp = dfp.reset_index(drop=True) | |
| bounds = [round(i * len(dfp) / _WORKERS) for i in range(_WORKERS + 1)] | |
| # The ZeroGPU worker is a daemonic process, which may not have children; it is forked fresh for | |
| # every call, so lift that flag here and fork a pool for this call only (numpy only, never CUDA). | |
| multiprocessing.current_process()._config["daemon"] = False | |
| with multiprocessing.get_context("fork").Pool(_WORKERS) as pool: | |
| parts = pool.map(_chunk_features, [(dfp.iloc[a:b], maps) for a, b in zip(bounds, bounds[1:]) if b > a]) | |
| out = pd.concat(parts, ignore_index=True) | |
| assert len(out) == len(dfp), "parallel candidate features lost rows" | |
| return out | |
| infer.af_chain_features = parallel_features | |
| # 2. The SAM image encoder sees the same full image several times per request (mean embedding, | |
| # patch norms, the decoder's set_image, both extractor passes) and the same candidate batches | |
| # in both passes. Memoise its forward on the exact input tensor; a hit returns the stored output. | |
| def memoise_sam_encoder(sam, max_entries=48): | |
| import hashlib | |
| enc, cache = sam.image_encoder, {} | |
| forward = enc.forward | |
| def cached_forward(x): | |
| key = (tuple(x.shape), str(x.dtype), hashlib.sha1(x.detach().cpu().numpy().tobytes()).hexdigest()) | |
| if key not in cache: | |
| if len(cache) >= max_entries: | |
| cache.pop(next(iter(cache))) | |
| cache[key] = forward(x) | |
| return cache[key] | |
| enc.forward = cached_forward | |
| return sam | |
| _get_sam = _shared.get_sam | |
| def get_sam_cached(): | |
| if "sam" not in _shared._CACHE: | |
| memoise_sam_encoder(_get_sam()) | |
| return _shared._CACHE["sam"] | |
| _shared.get_sam = get_sam_cached | |
| explain_vlm.load_vlm = functools.cache(load_vlm_bf16 if os.environ.get("SPACE_ID") else explain_vlm.load_vlm) | |
| EX.write_explanation = lambda crop_path, prompt: explain_vlm.generate_explanation(crop_path, prompt, str(EX.ADAPTER_DIR)) | |
| photoguider.ExtractorPool = InProcessPool | |
| def preload(session, explain=True): | |
| """Load every model once (extractors, saliency, explainer) so requests only run inference.""" | |
| for op in feature_extractors.OP_ENV: | |
| session.pool.module(op) | |
| FS._model() | |
| import feat_clip_refine | |
| feat_clip_refine._load() | |
| if "rem" not in FS._STATE: | |
| from transparent_background import Remover | |
| FS._STATE["rem"] = Remover() | |
| if explain: | |
| explain_vlm.load_vlm(str(EX.ADAPTER_DIR)) | |