gradio-zerogpu-page-vithplus

Dockerfile

@@ -1,24 +0,0 @@
-FROM python:3.11-slim
-
-ENV PYTHONUNBUFFERED=1 \
-    PIP_NO_CACHE_DIR=1 \
-    PORT=7860 \
-    SPACE_DEFAULT_HF_ONLY=1 \
-    DEFAULT_MODEL_NAME=page_vitb_screen \
-    PRELOAD_DEFAULT_MODEL=0
-
-WORKDIR /app
-
-RUN apt-get update \
-    && apt-get install -y --no-install-recommends git \
-    && rm -rf /var/lib/apt/lists/*
-
-COPY requirements.txt .
-RUN pip install --upgrade pip \
-    && pip install -r requirements.txt
-
-COPY . .
-
-EXPOSE 7860
-
-CMD ["python", "app.py"]

README.md

@@ -3,8 +3,9 @@ title: CrossGaze
 emoji: 👀
 colorFrom: blue
 colorTo: purple
-sdk: docker
-app_port: 7860
+sdk: gradio
+app_file: app_gradio.py
+short_description: CrossGaze / PaGE gaze target estimation demo on ZeroGPU.
 pinned: false
 ---
 
@@ -74,17 +75,32 @@ python app.py
 | `page_vitb_screen` | `Octopus1/page-vitb-screen` | PaGE ViT-B Screen |
 | `page_vitsplus` | `Octopus1/page-vitsplus` | PaGE ViT-S+ |
 | `page_vits` | `Octopus1/page-vits` | PaGE ViT-S |
+| `page_vithplus` | `Octopus1/page-vithplus` | PaGE ViT-H+ |
 
 这些模型会通过 `AutoModel.from_pretrained(..., trust_remote_code=True)` 加载,不需要本地 `checkpoints/*.pt`。
 
 ## Hugging Face Space 部署
 
-本项目在 Space 上使用 Docker SDK 部署。Space 环境默认设置：
+本项目提供两套入口：
+
+- `app.py`：原 FastAPI + 静态网页版本，适合 Docker Space/本地服务。
+- `app_gradio.py`：Gradio + ZeroGPU 版本，适合 Hugging Face ZeroGPU Space。
+
+Space YAML 已配置为 Gradio：
+
+```yaml
+sdk: gradio
+app_file: app_gradio.py
+```
+
+Gradio 版默认设置：
 
 - `SPACE_DEFAULT_HF_ONLY=1`：只暴露 Hugging Face 模型，避免依赖本地 checkpoint。
-- `DEFAULT_MODEL_NAME=page_vitb_screen`：启动时预加载 `Octopus1/page-vitb-screen`。
-- `PRELOAD_DEFAULT_MODEL=0`：Space 上不阻塞启动，模型在用户选择/推理时再加载。
-- `PORT=7860`：符合 Hugging Face Docker Space 默认入口。
+- `DEFAULT_MODEL_NAME=page_vithplus`：默认模型为 `Octopus1/page-vithplus`。
+- `PRELOAD_DEFAULT_MODEL=0`：不在启动时预加载模型，避免阻塞 Space 启动。
+- `@spaces.GPU(duration=180)`：只在推理函数执行期间申请 ZeroGPU。
+
+当前 Gradio 版使用文本输入 bbox：每行 `x1,y1,x2,y2`，支持 0~1 归一化坐标；可选第 5 列颜色，例如 `0.1,0.2,0.3,0.5,#ff7a00`。
 
 如果 PaGE 模型仓库是 private，需要在 Space Settings → Variables and secrets 中添加 secret：
 

app.py

@@ -67,10 +67,15 @@ ALL_MODEL_OPTIONS = {
         "repo_id": "Octopus1/page-vits",
         "display_name": "PaGE ViT-S",
     },
+    "page_vithplus": {
+        "source": "hf_page",
+        "repo_id": "Octopus1/page-vithplus",
+        "display_name": "PaGE ViT-H+",
+    },
 }
 
 HF_ONLY_MODE = os.environ.get("SPACE_DEFAULT_HF_ONLY", "").lower() in {"1", "true", "yes"} or bool(os.environ.get("SPACE_ID"))
-DEFAULT_MODEL_NAME = os.environ.get("DEFAULT_MODEL_NAME", "page_vitb_screen" if HF_ONLY_MODE else "cross_gaze_vitb_distill_screen")
+DEFAULT_MODEL_NAME = os.environ.get("DEFAULT_MODEL_NAME", "page_vithplus" if HF_ONLY_MODE else "cross_gaze_vitb_distill_screen")
 PRELOAD_DEFAULT_MODEL = os.environ.get("PRELOAD_DEFAULT_MODEL", "0" if HF_ONLY_MODE else "1").lower() in {"1", "true", "yes"}
 OUT_THRESHOLD = 0.5
 MAX_INFER_SIDE = 720
@@ -133,9 +138,28 @@ def _patch_factory_backbone_paths(model_name: str):
 class ModelManager:
     def __init__(self) -> None:
         self._cache: dict[str, tuple[torch.nn.Module, list, list | None]] = {}
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = self._current_device()
+
+    def _current_device(self) -> torch.device:
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    def refresh_device(self) -> torch.device:
+        next_device = self._current_device()
+        if next_device != self.device:
+            self.device = next_device
+            for model, _, _ in self._cache.values():
+                model.to(self.device)
+        return self.device
+
+    def move_to_cpu(self) -> None:
+        self.device = torch.device("cpu")
+        for model, _, _ in self._cache.values():
+            model.to(self.device)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
     def get(self, model_name: str):
+        self.refresh_device()
         cfg = MODEL_OPTIONS.get(model_name)
         if cfg is None:
             raise HTTPException(status_code=400, detail=f"Unsupported model: {model_name}")

app_gradio.py

@@ -0,0 +1,378 @@
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+os.environ.setdefault("SPACE_DEFAULT_HF_ONLY", "1")
+os.environ.setdefault("DEFAULT_MODEL_NAME", "page_vithplus")
+os.environ.setdefault("PRELOAD_DEFAULT_MODEL", "0")
+os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")
+os.environ.setdefault("NO_PROXY", "127.0.0.1,localhost")
+os.environ.setdefault("no_proxy", "127.0.0.1,localhost")
+os.environ.setdefault("RELEASE_GPU_AFTER_INFERENCE", "1")
+
+import gradio as gr
+import numpy as np
+import torch
+from PIL import Image
+
+try:
+    import spaces
+except Exception:
+    class _SpacesFallback:
+        @staticmethod
+        def GPU(*_args, **_kwargs):
+            def decorator(fn):
+                return fn
+
+            return decorator
+
+    spaces = _SpacesFallback()
+
+from app import (
+    BBox,
+    DEFAULT_MODEL_NAME,
+    MODEL_OPTIONS,
+    OUT_THRESHOLD,
+    PersonSpec,
+    _clamp_bbox,
+    _downsample_for_inference,
+    _draw_combined_visualization,
+    _draw_multi_visualization,
+    _find_peak,
+    _make_head_inputs,
+    model_manager,
+)
+
+
+ROOT = Path(__file__).resolve().parent
+EXAMPLES_DIR = ROOT / "static" / "examples"
+MANIFEST_PATH = EXAMPLES_DIR / "manifest.json"
+
+
+def _parse_bbox_line(line: str, width: int, height: int, default_color: str) -> PersonSpec:
+    fields = [field.strip() for field in line.replace("\t", ",").split(",") if field.strip()]
+    if len(fields) < 4:
+        raise ValueError(f"bbox needs at least 4 values: {line}")
+
+    values = [float(fields[i]) for i in range(4)]
+    color = fields[4] if len(fields) >= 5 else default_color
+    if all(0 <= value <= 1 for value in values):
+        x1, y1, x2, y2 = values
+        values = [x1 * width, y1 * height, x2 * width, y2 * height]
+
+    return PersonSpec(
+        bbox=BBox(x1=values[0], y1=values[1], x2=values[2], y2=values[3]),
+        color=color,
+    )
+
+
+def _parse_people(bbox_text: str, width: int, height: int) -> list[PersonSpec]:
+    colors = [
+        "#00ff00",
+        "#ff7a00",
+        "#00a7ff",
+        "#ff4fd8",
+        "#ffd400",
+        "#8b5cff",
+    ]
+    lines = [
+        line.split("# ", 1)[0].strip()
+        for line in bbox_text.replace(";", "\n").splitlines()
+    ]
+    lines = [line for line in lines if line]
+    if not lines:
+        raise ValueError("Please enter at least one bbox: x1,y1,x2,y2. Normalized coordinates in [0,1] are supported.")
+    return [_parse_bbox_line(line, width, height, colors[idx % len(colors)]) for idx, line in enumerate(lines)]
+
+
+def _to_pil_image(value) -> Image.Image | None:
+    if value is None:
+        return None
+    if isinstance(value, Image.Image):
+        return value.convert("RGB")
+    if isinstance(value, np.ndarray):
+        return Image.fromarray(value).convert("RGB")
+    if isinstance(value, (str, Path)):
+        return Image.open(value).convert("RGB")
+    return None
+
+
+def _layer_to_bbox(layer, width: int, height: int, color: str) -> PersonSpec | None:
+    if layer is None:
+        return None
+    if isinstance(layer, Image.Image):
+        arr = np.array(layer.convert("RGBA"))
+    elif isinstance(layer, np.ndarray):
+        arr = layer
+        if arr.ndim == 2:
+            mask = arr > 0
+            ys, xs = np.where(mask)
+            if len(xs) == 0:
+                return None
+            return PersonSpec(bbox=BBox(x1=float(xs.min()), y1=float(ys.min()), x2=float(xs.max() + 1), y2=float(ys.max() + 1)), color=color)
+        if arr.shape[-1] == 3:
+            mask = np.any(arr[..., :3] > 0, axis=-1)
+            ys, xs = np.where(mask)
+            if len(xs) == 0:
+                return None
+            return PersonSpec(bbox=BBox(x1=float(xs.min()), y1=float(ys.min()), x2=float(xs.max() + 1), y2=float(ys.max() + 1)), color=color)
+    else:
+        return None
+
+    mask = arr[..., 3] > 0 if arr.shape[-1] >= 4 else np.any(arr[..., :3] > 0, axis=-1)
+    ys, xs = np.where(mask)
+    if len(xs) == 0:
+        return None
+
+    x1 = max(0, min(width - 1, int(xs.min())))
+    y1 = max(0, min(height - 1, int(ys.min())))
+    x2 = max(x1 + 1, min(width, int(xs.max()) + 1))
+    y2 = max(y1 + 1, min(height, int(ys.max()) + 1))
+    return PersonSpec(bbox=BBox(x1=x1, y1=y1, x2=x2, y2=y2), color=color)
+
+
+def _image_and_people_from_editor(editor_value, bbox_text: str) -> tuple[Image.Image, list[PersonSpec], str]:
+    colors = ["#00ff00", "#ff7a00", "#00a7ff", "#ff4fd8", "#ffd400", "#8b5cff"]
+
+    if isinstance(editor_value, dict):
+        image = _to_pil_image(editor_value.get("background")) or _to_pil_image(editor_value.get("composite"))
+        if image is None:
+            raise gr.Error("Please upload or choose an image first.")
+
+        people: list[PersonSpec] = []
+        for idx, layer in enumerate(editor_value.get("layers") or []):
+            spec = _layer_to_bbox(layer, image.width, image.height, colors[idx % len(colors)])
+            if spec is not None:
+                people.append(spec)
+
+        if people:
+            return image, people, "editor"
+        return image, _parse_people(bbox_text, image.width, image.height), "text"
+
+    image = _to_pil_image(editor_value)
+    if image is None:
+        raise gr.Error("Please upload or choose an image first.")
+    return image, _parse_people(bbox_text, image.width, image.height), "text"
+
+
+def _run_one(
+    pil_image: Image.Image,
+    model_name: str,
+    spec: PersonSpec,
+    show_heatmap: bool,
+    show_gaze: bool,
+):
+    original_width, original_height = pil_image.size
+    pil_image, image_scale = _downsample_for_inference(pil_image)
+    width, height = pil_image.size
+
+    bbox_obj = spec.bbox
+    if image_scale != 1.0:
+        bbox_obj = BBox(
+            x1=bbox_obj.x1 * image_scale,
+            y1=bbox_obj.y1 * image_scale,
+            x2=bbox_obj.x2 * image_scale,
+            y2=bbox_obj.y2 * image_scale,
+        )
+    bbox_px = _clamp_bbox(bbox_obj, width, height)
+    bbox_norm = [bbox_px[0] / width, bbox_px[1] / height, bbox_px[2] / width, bbox_px[3] / height]
+
+    model, scene_transforms, head_transforms = model_manager.get(model_name)
+    scene_inputs = [transform(pil_image).unsqueeze(0).to(model_manager.device) for transform in scene_transforms]
+    head_inputs = _make_head_inputs(pil_image, bbox_px, head_transforms)
+    if head_inputs is not None:
+        head_inputs = [tensor.to(model_manager.device) for tensor in head_inputs]
+
+    with torch.inference_mode():
+        output = model({
+            "images": scene_inputs,
+            "head_images": head_inputs,
+            "bboxes": [[tuple(bbox_norm)]],
+        })
+
+    heatmap = output["heatmap"][0][0]
+    inout_score = float(output["inout"][0][0].detach().cpu().item()) if output.get("inout") is not None else 1.0
+    is_out = inout_score < OUT_THRESHOLD
+    target_x, target_y, peak_score = _find_peak(heatmap, width, height)
+    combined = _draw_combined_visualization(
+        pil_image,
+        bbox_px,
+        bbox_norm,
+        heatmap,
+        (target_x, target_y),
+        is_out,
+        show_heatmap=show_heatmap,
+        show_gaze=show_gaze,
+    )
+    return {
+        "bbox_px": bbox_px,
+        "bbox_norm": bbox_norm,
+        "heatmap": heatmap,
+        "target_xy": (target_x, target_y),
+        "is_out": is_out,
+        "result": {
+            "bbox": {"x1": bbox_px[0], "y1": bbox_px[1], "x2": bbox_px[2], "y2": bbox_px[3]},
+            "target": None if is_out else {"x": target_x, "y": target_y, "score": peak_score},
+            "inout_score": inout_score,
+            "is_out": is_out,
+            "image_size": {"width": width, "height": height},
+            "original_image_size": {"width": original_width, "height": original_height},
+        },
+        "combined": combined,
+    }
+
+
+@spaces.GPU(duration=180)
+def run_inference(
+    image_editor,
+    bbox_text: str,
+    model_name: str,
+    show_heatmap: bool,
+    show_gaze: bool,
+):
+    try:
+        pil_image, people, bbox_source = _image_and_people_from_editor(image_editor, bbox_text)
+        if model_name not in MODEL_OPTIONS:
+            raise gr.Error(f"Unsupported model: {model_name}")
+
+        if len(people) == 1:
+            item = _run_one(pil_image, model_name, people[0], show_heatmap, show_gaze)
+            return item["combined"], json.dumps({
+                "model_name": model_name,
+                "device": str(model_manager.device),
+                "bbox_source": bbox_source,
+                **item["result"],
+            }, ensure_ascii=False, indent=2)
+
+        downsampled, _ = _downsample_for_inference(pil_image)
+        width, height = downsampled.size
+
+        bbox_px_list = []
+        bbox_norm_list = []
+        heatmaps = []
+        target_xy_list = []
+        is_out_list = []
+        colors = []
+        results = []
+
+        for idx, spec in enumerate(people):
+            item = _run_one(pil_image, model_name, spec, False, show_gaze)
+            bbox_px_list.append(item["bbox_px"])
+            bbox_norm_list.append(item["bbox_norm"])
+            heatmaps.append(item["heatmap"])
+            target_xy_list.append(item["target_xy"])
+            is_out_list.append(item["is_out"])
+            colors.append(spec.color)
+            results.append({"index": idx, "color": spec.color, **item["result"]})
+
+        combined = _draw_multi_visualization(
+            downsampled,
+            bbox_px_list,
+            bbox_norm_list,
+            heatmaps,
+            target_xy_list,
+            is_out_list,
+            colors,
+            show_heatmap=show_heatmap,
+            show_gaze=show_gaze,
+        )
+        return combined, json.dumps({
+            "model_name": model_name,
+            "device": str(model_manager.device),
+            "bbox_source": bbox_source,
+            "people": results,
+            "image_size": {"width": width, "height": height},
+            "original_image_size": {"width": pil_image.width, "height": pil_image.height},
+        }, ensure_ascii=False, indent=2)
+    finally:
+        if os.environ.get("RELEASE_GPU_AFTER_INFERENCE", "1").lower() in {"1", "true", "yes"}:
+            model_manager.move_to_cpu()
+
+
+def _load_examples() -> list[list[str]]:
+    if not MANIFEST_PATH.exists():
+        return []
+    items = json.loads(MANIFEST_PATH.read_text())
+    examples = []
+    for item in items[:8]:
+        image_path = ROOT / item["image"].lstrip("/")
+        if image_path.exists():
+            examples.append([
+                str(image_path),
+                "0.35,0.15,0.60,0.55",
+                DEFAULT_MODEL_NAME,
+                True,
+                True,
+            ])
+    return examples
+
+
+def build_demo() -> gr.Blocks:
+    model_choices = [
+        (cfg.get("display_name", name), name)
+        for name, cfg in MODEL_OPTIONS.items()
+    ]
+
+    with gr.Blocks(title="CrossGaze ZeroGPU Demo") as demo:
+        gr.Markdown(
+            """
+            # CrossGaze / PaGE ZeroGPU Demo
+
+            Upload an image, enter one or more head/face bounding boxes, and run gaze target inference.
+            BBox format: `x1,y1,x2,y2` per line. If all four values are in `[0,1]`, they are treated as normalized coordinates.
+            Optional fifth value sets color, e.g. `0.1,0.2,0.3,0.5,#ff7a00`.
+            """
+        )
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                image_input = gr.ImageEditor(
+                    type="pil",
+                    label="Input image / draw head region",
+                    image_mode="RGBA",
+                    sources=["upload", "clipboard"],
+                    brush=gr.Brush(default_size=24),
+                    eraser=gr.Eraser(default_size=24),
+                    layers=True,
+                )
+                bbox_input = gr.Textbox(
+                    label="BBox coordinates (fallback / precise input)",
+                    value="0.35,0.15,0.60,0.55",
+                    lines=5,
+                    placeholder="0.35,0.15,0.60,0.55\n0.12,0.20,0.30,0.48,#ff7a00",
+                )
+                model_input = gr.Dropdown(
+                    choices=model_choices,
+                    value=DEFAULT_MODEL_NAME,
+                    label="Model",
+                )
+                show_heatmap = gr.Checkbox(value=True, label="Show heatmap")
+                show_gaze = gr.Checkbox(value=True, label="Show gaze point and line")
+                run_button = gr.Button("Run inference", variant="primary")
+
+            with gr.Column(scale=1):
+                output_image = gr.Image(type="pil", label="Result")
+                output_json = gr.Code(label="Result JSON", language="json")
+
+        gr.Examples(
+            examples=_load_examples(),
+            inputs=[image_input, bbox_input, model_input, show_heatmap, show_gaze],
+            label="Examples",
+        )
+
+        run_button.click(
+            run_inference,
+            inputs=[image_input, bbox_input, model_input, show_heatmap, show_gaze],
+            outputs=[output_image, output_json],
+        )
+
+    return demo
+
+
+demo = build_demo()
+
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -1,6 +1,10 @@
-fastapi
+fastapi==0.115.11
+starlette==0.46.1
 uvicorn
 python-multipart
+gradio==4.44.1
+spaces
+huggingface-hub>=0.34.0,<1.0
 # Model runtime deps (versions aligned with the gazelle env)
 torch==2.10.0
 torchvision==0.25.0