Keep CPU inference safe across ZeroGPU forks and size GPU requests per task

app.py

@@ -212,7 +212,9 @@ def detect(session, jid, request: DetectRequest):
     import torch
     from pixtext_ocr import detect_words
 
-    torch.set_num_threads(4)
+    threads = int(os.environ.get("PIXTEXT_CPU_THREADS", "4"))
+    if torch.get_num_threads() != threads:
+        torch.set_num_threads(threads)
     root = ocr_root(jid)
     update_job(jid, message="Detecting words and reading text", progress=0.3)
     with Image.open(session["directory"] / "source.png") as image:
@@ -229,7 +231,9 @@ def refine(session, jid, request: RefineRequest, progress_report=None):
 
     if not torch.cuda.is_available():
         raise RuntimeError("Refinement needs a CUDA GPU. Detection and box editing are ready to use.")
-    torch.set_num_threads(4)
+    threads = int(os.environ.get("PIXTEXT_CPU_THREADS", "4"))
+    if torch.get_num_threads() != threads:
+        torch.set_num_threads(threads)
     if os.environ.get("PIXTEXT_GPU_MEMORY_FRACTION"):
         torch.cuda.set_per_process_memory_fraction(float(os.environ["PIXTEXT_GPU_MEMORY_FRACTION"]))
     if PIPELINE_VARIANT != request.variant:

pixtext_ocr.py

@@ -8,6 +8,7 @@ from __future__ import annotations
 import hashlib
 import json
 import math
+import os
 from functools import lru_cache
 from pathlib import Path
 
@@ -82,7 +83,7 @@ class WordOCR:
         )
         self.detector.eval()
         options = ort.SessionOptions()
-        options.intra_op_num_threads = 4
+        options.intra_op_num_threads = int(os.environ.get("PIXTEXT_CPU_THREADS", "4"))
         options.inter_op_num_threads = 1
         self.recognizer = ort.InferenceSession(
             str(rec_path), sess_options=options, providers=["CPUExecutionProvider"]

space.py

@@ -3,6 +3,9 @@
 import os
 
 os.environ["PIXTEXT_GRADIO"] = "1"
+os.environ.setdefault("PIXTEXT_CPU_THREADS", "1")
+os.environ.setdefault("OMP_NUM_THREADS", "1")
+os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
 os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
 
 # Import before torch so ZeroGPU can enable its CUDA staging mode.
@@ -26,7 +29,7 @@ def load_models():
     from pixtext import PixTextPipeline
     from huggingface_hub import snapshot_download
 
-    torch.set_num_threads(4)
+    torch.set_num_threads(int(os.environ["PIXTEXT_CPU_THREADS"]))
     if not torch.cuda.is_available() and not spaces.config.Config.zero_gpu:
         return {}
     pipelines = {}
@@ -47,8 +50,16 @@ def load_models():
 PIPELINES = load_models()
 
 
-@spaces.GPU(duration=120)
+def estimate_duration(directory: str, size: list[int], jid: str, payload: str):
+    import math
+    request = json.loads(payload)
+    count = sum(row.get("enabled", True) for row in request["regions"])
+    return max(20, min(90, 10 + math.ceil(request["steps"] * max(1, count / 8))))
+
+
+@spaces.GPU(duration=estimate_duration)
 def gpu_refine(directory: str, size: list[int], jid: str, payload: str):
+    print("PixText GPU worker: starting", flush=True)
     started = time.perf_counter()
     request = studio.RefineRequest.model_validate_json(payload)
     studio.PIPELINE = PIPELINES.get(request.variant)
@@ -56,10 +67,13 @@ def gpu_refine(directory: str, size: list[int], jid: str, payload: str):
     progress = gr.Progress()
 
     def report(**values):
-        progress(values.get("progress", 0), desc=values.get("message", "Refining"))
+        message = values.get("message", "Refining")
+        print("PixText GPU:", message, flush=True)
+        progress(values.get("progress", 0), desc=message)
 
     result = studio.refine({"directory": Path(directory), "size": size}, jid, request, report)
     result["elapsed"] = round(time.perf_counter() - started, 2)
+    print("PixText GPU worker: complete", result["elapsed"], flush=True)
     return result