Spaces:
Running on Zero
Running on Zero
File size: 4,529 Bytes
63f2fb8 dbc3379 63f2fb8 df87cf3 63f2fb8 dbc3379 63f2fb8 b9346a5 63f2fb8 dbc3379 63f2fb8 dbc3379 63f2fb8 dbc3379 63f2fb8 dbc3379 63f2fb8 87ade98 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """Hugging Face entry point: custom studio UI, Gradio queue, free ZeroGPU."""
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.
import spaces
# Only gpu_refine needs a GPU; keep validation and session state in the web process.
spaces.disable_gradio_auto_wrap()
import gradio as gr
import hashlib
import json
import time
import uuid
from pathlib import Path
import app as studio
from zerogpu_attention import enable_sdpa
enable_sdpa()
def load_models():
import torch
from pixtext import PixTextPipeline
from huggingface_hub import snapshot_download
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 = {}
for variant, name in (("meanflow", "PixText-MeanFlow"), ("teacher", "PixText-Teacher")):
if os.environ.get("PIXTEXT_MODEL_ROOT"):
package = Path(os.environ["PIXTEXT_MODEL_ROOT"]) / name
else:
package = Path(snapshot_download("ATH-MaaS/PixText", allow_patterns=[f"models/{name}/*"],
local_dir=studio.RUNTIME / "model-packages")) / "models" / name
pipeline = PixTextPipeline.from_pretrained(package, device="cpu", dtype=torch.bfloat16)
pipeline.model = pipeline.model.to("cuda")
pipeline.device = torch.device("cuda")
pipelines[variant] = pipeline
return pipelines
# Stage both models before serving requests, as required by ZeroGPU.
PIPELINES = load_models()
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)
studio.PIPELINE_VARIANT = request.variant if studio.PIPELINE is not None else None
progress = gr.Progress()
def report(**values):
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
@studio.app.api(name="refine_image", concurrency_limit=1, concurrency_id="pixtext-gpu")
def refine_image(sid: str, payload: str) -> dict:
request = studio.RefineRequest.model_validate_json(payload)
session = studio.validate_refine(sid, request)
canonical = request.model_dump_json()
key = hashlib.sha256(canonical.encode()).hexdigest()
with studio.LOCK:
previous = session.get("last_gradio_result")
if previous and previous["key"] == key:
job = studio.JOBS.get(previous["jid"])
if job and job["status"] == "done":
return {"job_id": job["id"], "result": job["result"]}
jid = uuid.uuid4().hex
studio.JOBS[jid] = {"id": jid, "session": sid, "kind": "refine", "status": "running",
"progress": 0, "message": "Waiting for GPU", "created": time.time()}
session["job"] = jid
try:
result = gpu_refine(str(session["directory"]), list(session["size"]), jid, canonical)
studio.update_job(jid, status="done", progress=1, message="Complete", result=result)
session["last_gradio_result"] = {"key": key, "jid": jid}
return {"job_id": jid, "result": result}
except Exception:
studio.update_job(jid, status="error", message="Refinement could not finish. Please retry.")
raise
if __name__ == "__main__":
studio.app.launch(server_name=os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0"),
server_port=int(os.environ.get("PIXTEXT_PORT", "7860")), show_error=False, ssr_mode=False)
|