Spaces:
Running on Zero
Running on Zero
| """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)))) | |
| 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 | |
| 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) | |