Space app: SLaT session cache, tone-mapper safety, CUDA/inference fixes

- HDRI preview uses standalone ToneMapper(view=); dropdown calls setup_tone_mapper
- Image workflow stores SLaT in _SESSION_SLAT; file path still for uploaded npz
- Move NeAR and Hunyuan to CUDA when available on __main__
- export_glb and export_glb_from_slat run under inference_mode for slat_decoder_mesh
- setup_tone_mapper uses ToneMapper(view=) to avoid missing .cpu after OCIO errors
- spconv SparseTensor.replace: align feat buffer with spconv tensor storage

Made-with: Cursor

app.py

@@ -1,28 +1,6 @@
 import os
 import sys
-
-if not os.environ.get("HF_TOKEN") and not os.environ.get("HUGGING_FACE_HUB_TOKEN"):
-    _hub_tok = (os.environ.get("near") or os.environ.get("NEAR") or "").strip()
-    if _hub_tok:
-        os.environ["HF_TOKEN"] = _hub_tok
-        print("[NeAR] HF_TOKEN from Space secret 'near'.", flush=True)
-
-try:
-    _raw_zerogpu_cap = int(os.environ.get("NEAR_ZEROGPU_HF_CEILING_S", "90"))
-except ValueError:
-    _raw_zerogpu_cap = 90
-_ZEROGPU_ENV_CAP_S = min(max(15, _raw_zerogpu_cap), 120)
-for _ek in ("NEAR_ZEROGPU_MAX_SECONDS", "NEAR_ZEROGPU_DURATION_CAP"):
-    if _ek in os.environ:
-        try:
-            if int(os.environ[_ek]) > _ZEROGPU_ENV_CAP_S:
-                os.environ[_ek] = str(_ZEROGPU_ENV_CAP_S)
-        except ValueError:
-            pass
-print(f"[NeAR] ZeroGPU cap {_ZEROGPU_ENV_CAP_S}s (NEAR_ZEROGPU_HF_CEILING_S).", flush=True)
-
 import shutil
-import subprocess
 import threading
 import time
 from pathlib import Path
@@ -31,7 +9,7 @@ from typing import Any, Dict, Optional
 import gradio as gr
 
 try:
-    import spaces
+    import spaces  # pyright: ignore[reportMissingImports]
 except ImportError:
     spaces = None
 import imageio
@@ -41,20 +19,6 @@ import trimesh
 from PIL import Image
 from simple_ocio import ToneMapper  # pyright: ignore[reportMissingImports]
 
-try:
-    import gradio_client.utils as client_utils
-
-    _get_type_orig = client_utils.get_type
-
-    def _get_type_patched(schema):
-        if isinstance(schema, bool):
-            return "boolean"
-        return _get_type_orig(schema)
-
-    client_utils.get_type = _get_type_patched
-except Exception:
-    pass
-
 sys.path.insert(0, "./hy3dshape")
 os.environ.setdefault("ATTN_BACKEND", "xformers")
 os.environ.setdefault("SPCONV_ALGO", "native")
@@ -103,126 +67,19 @@ def _warn_example_assets() -> None:
 _warn_example_assets()
 
 DEFAULT_IMAGE = APP_DIR / "assets/example_image/T.png"
-DEFAULT_SLAT = APP_DIR / "assets/example_slats/2a0d671ce308adb93323eae7141953fc1a5ba68f38cc69f476d5e904c634864d.npz"
 DEFAULT_HDRI = APP_DIR / "assets/hdris/studio_small_03_1k.exr"
-DEFAULT_PORT = 7860
 MAX_SEED = np.iinfo(np.int32).max
 
 
-_SESSION_LAST_TOUCH: Dict[str, float] = {}
-_SESSION_TOUCH_LOCK = threading.Lock()
-
-
-def _session_touch(session_id: str) -> None:
-    with _SESSION_TOUCH_LOCK:
-        _SESSION_LAST_TOUCH[session_id] = time.time()
-
-
-def _session_forget(session_id: str) -> None:
-    with _SESSION_TOUCH_LOCK:
-        _SESSION_LAST_TOUCH.pop(session_id, None)
-
-
-def ensure_session_dir(req: Optional[gr.Request]) -> Path:
-    session_id = getattr(req, "session_hash", None) or "shared"
-    d = CACHE_DIR / str(session_id)
-    d.mkdir(parents=True, exist_ok=True)
-    _session_touch(session_id)
-    return d
-
-
-def clear_session_dir(req: Optional[gr.Request]) -> str:
-    d = ensure_session_dir(req)
-    shutil.rmtree(d, ignore_errors=True)
-    d.mkdir(parents=True, exist_ok=True)
-    # Do not touch torch.cuda here: on HF Stateless / ZeroGPU Spaces, CUDA must
-    # not be initialized in the main Gradio process (only inside @spaces.GPU).
-    return "Session cache cleared."
-
-
+def start_session(req: gr.Request):
+    user_dir = CACHE_DIR / str(req.session_hash)
+    os.makedirs(user_dir, exist_ok=True)
+    
+    
 def end_session(req: gr.Request):
-    session_id = getattr(req, "session_hash", None) or "shared"
-    d = CACHE_DIR / str(session_id)
-    shutil.rmtree(d, ignore_errors=True)
-    _session_forget(session_id)
-
-
-def _session_dir_latest_mtime(path: Path) -> float:
-    try:
-        latest = path.stat().st_mtime
-    except OSError:
-        return 0.0
-    try:
-        for child in path.rglob("*"):
-            if child.is_file():
-                try:
-                    latest = max(latest, child.stat().st_mtime)
-                except OSError:
-                    pass
-    except OSError:
-        pass
-    return latest
-
-
-def prune_tmp_gradio_sessions(max_age_s: int) -> int:
-    """Delete CACHE_DIR/<session> when both API idle and on-disk files are older than max_age_s."""
-
-    if max_age_s <= 0 or not CACHE_DIR.is_dir():
-        return 0
-    now = time.time()
-    removed = 0
-    with _SESSION_TOUCH_LOCK:
-        touch_snap = dict(_SESSION_LAST_TOUCH)
-    try:
-        entries = list(CACHE_DIR.iterdir())
-    except OSError:
-        return 0
-    for p in entries:
-        if not p.is_dir():
-            continue
-        sid = p.name
-        try:
-            last_api = touch_snap.get(sid)
-            latest_file = _session_dir_latest_mtime(p)
-            last_seen = max(last_api or 0.0, latest_file)
-            if now - last_seen < max_age_s:
-                continue
-            shutil.rmtree(p, ignore_errors=True)
-            _session_forget(sid)
-            removed += 1
-        except OSError:
-            continue
-    return removed
-
-
-_tmp_pruner_started = False
-_tmp_pruner_lock = threading.Lock()
-
-
-def start_tmp_gradio_pruner() -> None:
-    """Background prune of tmp_gradio session dirs (disk, not RAM). Configurable via env."""
-
-    global _tmp_pruner_started
-    interval = int(os.environ.get("NEAR_TMP_PRUNE_INTERVAL_S", "900"))
-    if interval <= 0:
-        return
-    with _tmp_pruner_lock:
-        if _tmp_pruner_started:
-            return
-        _tmp_pruner_started = True
-    max_age = int(os.environ.get("NEAR_TMP_GRADIO_MAX_AGE_S", str(48 * 3600)))
-
-    def _loop() -> None:
-        while True:
-            try:
-                n = prune_tmp_gradio_sessions(max_age)
-                if n:
-                    print(f"[NeAR] tmp_gradio prune: removed {n} stale session dir(s)", flush=True)
-            except Exception as exc:
-                print(f"[NeAR] tmp_gradio prune error: {exc}", flush=True)
-            time.sleep(interval)
-
-    threading.Thread(target=_loop, daemon=True, name="near-tmp-prune").start()
+    user_dir = CACHE_DIR / str(req.session_hash)
+    shutil.rmtree(user_dir)
+    _SESSION_SLAT.pop(str(req.session_hash), None)
 
 
 def get_file_path(file_obj: Any) -> Optional[str]:
@@ -237,62 +94,20 @@ def get_file_path(file_obj: Any) -> Optional[str]:
     return None
 
 
-_model_lock = threading.Lock()
 PIPELINE: Optional[NeARImageToRelightable3DPipeline] = None
 GEOMETRY_PIPELINE: Optional[Hunyuan3DDiTFlowMatchingPipeline] = None
-_light_preprocess_lock = threading.Lock()
-_light_preprocessor: Any | None = None
-_geometry_on_cuda = False
-_near_on_cuda = False
-
-_FALLBACK_TONE_MAPPER_CHOICES = ["AgX", "False", "Khronos neutrals", "Filmic", "Khronos glTF PBR"]
-
-
-def _truthy_env(name: str, default: str) -> bool:
-    v = (os.environ.get(name) if name in os.environ else default).strip().lower()
-    return v in ("1", "true", "yes", "on")
+tone_mapper = ToneMapper()
+AVAILABLE_TONE_MAPPERS = getattr(tone_mapper, "available_views", ["AgX"])
 
+# In-process SLaT for the image workflow (not serialized through Gradio State).
+_SESSION_SLAT: Dict[str, Any] = {}
 
-# When enabled, Hunyuan + NeAR weights are loaded on CPU inside demo.launch() *before* the
-# HTTP server binds, so the main UI only becomes reachable after CPU load finishes (avoids
-# clicks while models are missing → ZeroGPU timeout). Set NEAR_MODEL_CPU_PRELOAD_AT_START=0
-# to bind immediately and load on first @spaces.GPU click instead (faster "page up", riskier UX).
-_CPU_PRELOAD_DEFAULT = "1"
-_CPU_PRELOAD_AT_START = _truthy_env(
-    "NEAR_MODEL_CPU_PRELOAD_AT_START",
-    _CPU_PRELOAD_DEFAULT,
-)
-print(
-    f"[NeAR] NEAR_MODEL_CPU_PRELOAD_AT_START={'1' if _CPU_PRELOAD_AT_START else '0'} "
-    "(1 = block server start until CPU weights ready; 0 = lazy load on first GPU action).",
-    flush=True,
-)
-
-
-def _default_tone_mapper_choices() -> list[str]:
-    try:
-        views = getattr(ToneMapper(), "available_views", None)
-        if isinstance(views, (list, tuple)) and views:
-            return [str(v) for v in views]
-    except Exception as exc:
-        print(f"[NeAR] ToneMapper view discovery failed, using fallback choices: {exc}", flush=True)
-    return list(_FALLBACK_TONE_MAPPER_CHOICES)
-
-
-TONE_MAPPER_CHOICES = _default_tone_mapper_choices()
-
-
-def _get_light_image_preprocessor():
-    global _light_preprocessor
-    if _light_preprocessor is not None:
-        return _light_preprocessor
-    with _light_preprocess_lock:
-        if _light_preprocessor is None:
-            from hy3dshape.rembg import BackgroundRemover  # pyright: ignore[reportMissingImports]
+def set_tone_mapper(view_name: str):
+    if view_name and PIPELINE is not None:
+        PIPELINE.setup_tone_mapper(view_name)
 
-            _light_preprocessor = BackgroundRemover()
-            print("[NeAR] BackgroundRemover ready.", flush=True)
-    return _light_preprocessor
+from hy3dshape.rembg import BackgroundRemover  # pyright: ignore[reportMissingImports]
+LIGHT_PREPROCESSOR = BackgroundRemover()
 
 
 def _preprocess_image_rgba_light(input_image: Image.Image) -> Image.Image:
@@ -313,7 +128,7 @@ def _preprocess_image_rgba_light(input_image: Image.Image) -> Image.Image:
                 (int(rgb.width * scale), int(rgb.height * scale)),
                 Image.Resampling.LANCZOS,
             )
-        output = _get_light_image_preprocessor()(rgb)
+        output = LIGHT_PREPROCESSOR(rgb)
 
     if output.mode != "RGBA":
         output = output.convert("RGBA")
@@ -344,96 +159,15 @@ def _flatten_rgba_on_matte(image: Image.Image, matte_rgb: tuple[float, float, fl
     return NeARImageToRelightable3DPipeline.flatten_rgba_on_matte(image, matte_rgb)
 
 
-def _update_tone_mapper_choices(tone_mapper: Any) -> None:
-    global TONE_MAPPER_CHOICES
-    views = getattr(tone_mapper, "available_views", None)
-    if isinstance(views, (list, tuple)) and views:
-        TONE_MAPPER_CHOICES = [str(v) for v in views]
-
-
-def _ensure_geometry_cpu_locked() -> None:
-    global GEOMETRY_PIPELINE
-    if GEOMETRY_PIPELINE is not None:
-        return
-    hy_id = os.environ.get("NEAR_HUNYUAN_PRETRAINED", "tencent/Hunyuan3D-2.1")
-    t0 = time.time()
-    print(f"[NeAR] Hunyuan geometry on CPU from {hy_id!r}...", flush=True)
-    GEOMETRY_PIPELINE = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(hy_id, device="cpu")
-    print(f"[NeAR] Hunyuan CPU load {time.time() - t0:.1f}s", flush=True)
-
-
-def _ensure_near_cpu_locked() -> None:
-    global PIPELINE
-    if PIPELINE is not None:
-        return
-    near_id = os.environ.get("NEAR_PRETRAINED", "luh0502/NeAR")
-    t0 = time.time()
-    print(f"[NeAR] NeAR on CPU from {near_id!r}...", flush=True)
-    p = NeARImageToRelightable3DPipeline.from_pretrained(near_id)
-    p.to("cpu")
-    _update_tone_mapper_choices(p.tone_mapper)
-    PIPELINE = p
-    print(f"[NeAR] NeAR CPU load {time.time() - t0:.1f}s", flush=True)
-
-
-def ensure_geometry_on_cuda() -> Hunyuan3DDiTFlowMatchingPipeline:
-    global _geometry_on_cuda
-    with _model_lock:
-        _ensure_geometry_cpu_locked()
-        assert GEOMETRY_PIPELINE is not None
-        if torch.cuda.is_available() and not _geometry_on_cuda:
-            t0 = time.time()
-            GEOMETRY_PIPELINE.to("cuda")
-            _geometry_on_cuda = True
-            print(f"[NeAR] Hunyuan -> cuda {time.time() - t0:.1f}s", flush=True)
-        return GEOMETRY_PIPELINE
-
-
-def ensure_near_on_cuda() -> NeARImageToRelightable3DPipeline:
-    global _near_on_cuda
-    with _model_lock:
-        _ensure_near_cpu_locked()
-        assert PIPELINE is not None
-        if torch.cuda.is_available() and not _near_on_cuda:
-            t0 = time.time()
-            PIPELINE.to("cuda")
-            _near_on_cuda = True
-            print(f"[NeAR] NeAR -> cuda {time.time() - t0:.1f}s", flush=True)
-        return PIPELINE
-
-
-def run_model_cpu_preload_blocking() -> None:
-    """Load Hunyuan + NeAR on CPU before Gradio binds (main UI appears only after this)."""
-
-    t0 = time.time()
-    print("[NeAR] blocking CPU preload before server bind ...", flush=True)
-    with _model_lock:
-        _ensure_geometry_cpu_locked()
-        _ensure_near_cpu_locked()
-    print(
-        f"[NeAR] CPU preload done {time.time() - t0:.1f}s — Gradio will accept traffic now.",
-        flush=True,
-    )
-
-
-def set_tone_mapper(view_name: str):
-    pipeline = ensure_near_on_cuda()
-    if view_name:
-        pipeline.setup_tone_mapper(view_name)
-    return pipeline
-
-
-def preview_hdri(hdri_file_obj: Any, tone_mapper_name: str):
+def preview_hdri(hdri_file_obj: Any):
     hdri_path = get_file_path(hdri_file_obj)
     if not hdri_path:
         return None, "Upload an HDRI `.exr` (left column)."
     import pyexr  # pyright: ignore[reportMissingImports]
 
-    tone_mapper = ToneMapper()
-    if tone_mapper_name:
-        tone_mapper.view = tone_mapper_name
     hdri_np = pyexr.read(hdri_path)[..., :3]
-    preview = tone_mapper.hdr_to_ldr(hdri_np)
+    tm = ToneMapper(view="Khronos PBR Neutral")
+    preview = tm.hdr_to_ldr(hdri_np)
     preview = (np.clip(preview, 0, 1) * 255).astype(np.uint8)
     name = Path(hdri_path).name
     return preview, f"HDRI **{name}** — preview updated."
@@ -446,10 +180,6 @@ def switch_asset_source(mode: str):
 def _ensure_rgba(img: Image.Image) -> Image.Image:
     if img.mode == "RGBA":
         return img
-    if img.mode == "RGB":
-        r, g, b = img.split()
-        a = Image.new("L", img.size, 255)
-        return Image.merge("RGBA", (r, g, b, a))
     return img.convert("RGBA")
 
 
@@ -460,14 +190,6 @@ def preprocess_image_only(image_input: Optional[Image.Image]):
     return _preprocess_image_rgba_light(image_input)
 
 
-def save_slat_npz(slat, save_path: Path):
-    np.savez(
-        save_path,
-        feats=slat.feats.detach().cpu().numpy(),
-        coords=slat.coords.detach().cpu().numpy(),
-    )
-
-
 @GPU
 @torch.inference_mode()
 def generate_mesh(
@@ -475,8 +197,7 @@ def generate_mesh(
     req: gr.Request,
     progress=gr.Progress(track_tqdm=True),
 ):
-    geometry_pipeline = ensure_geometry_on_cuda()
-    session_dir = ensure_session_dir(req)
+    session_dir = CACHE_DIR / str(req.session_hash)
 
     if image_input is None:
         raise gr.Error("Please upload an input image.")
@@ -490,15 +211,17 @@ def generate_mesh(
     mesh_rgb.save(session_dir / "input_processed.png")
 
     progress(0.6, desc="Generating geometry")
-    mesh = geometry_pipeline(image=mesh_rgb)[0]
+    mesh = GEOMETRY_PIPELINE(image=mesh_rgb)[0]
     mesh_path = session_dir / "initial_3d_shape.glb"
     mesh.export(mesh_path)
 
+    _SESSION_SLAT.pop(str(req.session_hash), None)
     state = {
         "mode": "image",
         "mesh_path": str(mesh_path),
         "processed_image_path": str(session_dir / "input_processed.png"),
         "slat_path": None,
+        "slat_in_memory": False,
     }
     return (
         state,
@@ -509,16 +232,14 @@ def generate_mesh(
 
 @GPU
 @torch.inference_mode()
-def generate_slat(
+def _generate_slat_inner(
     asset_state: Dict[str, Any],
     image_input: Optional[Image.Image],
     seed: int,
     req: gr.Request,
     progress=gr.Progress(track_tqdm=True),
 ):
-    pipeline = ensure_near_on_cuda()
-    session_dir = ensure_session_dir(req)
-
+    """GPU body for SLaT generation — must be called from within a @GPU context."""
     if not asset_state or not asset_state.get("mesh_path"):
         raise gr.Error("Please run ① Generate Mesh first.")
     mesh_path = asset_state["mesh_path"]
@@ -536,28 +257,35 @@ def generate_slat(
     slat_rgb = _flatten_rgba_on_matte(rgba, (0.0, 0.0, 0.0))
 
     progress(0.3, desc="Computing SLaT coordinates")
-    coords = pipeline.shape_to_coords(mesh)
+    coords = PIPELINE.shape_to_coords(mesh)
 
     progress(0.6, desc="Generating SLaT")
-    slat = pipeline.run_with_coords([slat_rgb], coords, seed=int(seed), preprocess_image=False)
-
-    slat_path = session_dir / "generated_slat.npz"
-    save_slat_npz(slat, slat_path)
+    slat = PIPELINE.run_with_coords([slat_rgb], coords, seed=int(seed), preprocess_image=False)
 
-    new_state = {**asset_state, "slat_path": str(slat_path)}
+    _SESSION_SLAT[str(req.session_hash)] = slat
+    new_state = {**asset_state, "slat_path": None, "slat_in_memory": True}
     return new_state, f"**Asset ready** — SLaT generated (seed `{seed}`)."
 
 
-def load_slat_file(slat_upload: Any, slat_path_text: str, req: gr.Request):
+def _load_slat_file_inner(slat_upload: Any, slat_path_text: str, req: gr.Request):
     resolved = get_file_path(slat_upload) or (slat_path_text.strip() if slat_path_text else "")
     if not resolved:
         raise gr.Error("Please provide a SLaT `.npz` path or upload one.")
     if not os.path.exists(resolved):
         raise gr.Error(f"SLaT file not found: `{resolved}`")
-    state = {"mode": "slat", "slat_path": resolved, "mesh_path": None, "processed_image_path": None}
+    _SESSION_SLAT.pop(str(req.session_hash), None)
+    state = {
+        "mode": "slat",
+        "slat_path": resolved,
+        "mesh_path": None,
+        "processed_image_path": None,
+        "slat_in_memory": False,
+    }
     return state, f"SLaT **{Path(resolved).name}** loaded."
 
 
+@GPU
+@torch.inference_mode()
 def prepare_slat(
     source_mode: str,
     asset_state: Dict[str, Any],
@@ -569,25 +297,34 @@ def prepare_slat(
     progress=gr.Progress(track_tqdm=True),
 ):
     if source_mode == "From Image":
-        return generate_slat(asset_state, image_input, seed, req, progress)
-    return load_slat_file(slat_upload, slat_path_text, req)
+        return _generate_slat_inner(asset_state, image_input, seed, req, progress)
+    return _load_slat_file_inner(slat_upload, slat_path_text, req)
 
 
 def require_asset_state(asset_state: Optional[Dict[str, Any]]) -> Dict[str, Any]:
-    if not asset_state or not asset_state.get("slat_path"):
+    if not asset_state:
         raise gr.Error("Please generate or load a SLaT first.")
-    return asset_state
+    if asset_state.get("slat_in_memory") or asset_state.get("slat_path"):
+        return asset_state
+    raise gr.Error("Please generate or load a SLaT first.")
 
 
-def load_asset_and_hdri(asset_state: Dict[str, Any], hdri_file_obj: Any, tone_mapper_name: str):
+def load_asset_and_hdri(asset_state: Dict[str, Any], hdri_file_obj: Any, req: gr.Request):
     asset_state = require_asset_state(asset_state)
     hdri_path = get_file_path(hdri_file_obj)
     if not hdri_path:
         raise gr.Error("Please upload an HDRI `.exr` file.")
-    pipeline = set_tone_mapper(tone_mapper_name)
-    slat = pipeline.load_slat(asset_state["slat_path"])
-    hdri_np = pipeline.load_hdri(hdri_path)
-    return pipeline, slat, hdri_np
+    if asset_state.get("slat_in_memory"):
+        slat = _SESSION_SLAT.get(str(req.session_hash))
+        if slat is None:
+            raise gr.Error("SLaT session expired — run **② Generate / Load SLaT** again.")
+    else:
+        slat_path = asset_state.get("slat_path")
+        if not slat_path:
+            raise gr.Error("Please generate or load a SLaT first.")
+        slat = PIPELINE.load_slat(slat_path)
+    hdri_np = PIPELINE.load_hdri(hdri_path)
+    return slat, hdri_np
 
 
 @GPU
@@ -595,7 +332,6 @@ def load_asset_and_hdri(asset_state: Dict[str, Any], hdri_file_obj: Any, tone_ma
 def render_preview(
     asset_state: Dict[str, Any],
     hdri_file_obj: Any,
-    tone_mapper_name: str,
     hdri_rot: float,
     yaw: float,
     pitch: float,
@@ -606,12 +342,12 @@ def render_preview(
     progress=gr.Progress(track_tqdm=True),
 ):
     t0 = time.time()
-    session_dir = ensure_session_dir(req)
+    session_dir = CACHE_DIR / str(req.session_hash)
     progress(0.1, desc="Loading SLaT and HDRI")
-    pipeline, slat, hdri_np = load_asset_and_hdri(asset_state, hdri_file_obj, tone_mapper_name)
+    slat, hdri_np = load_asset_and_hdri(asset_state, hdri_file_obj, req)
 
     progress(0.5, desc="Rendering")
-    views = pipeline.render_view(
+    views = PIPELINE.render_view(
         slat, hdri_np,
         yaw_deg=yaw, pitch_deg=pitch, fov=fov, radius=radius,
         hdri_rot_deg=hdri_rot, resolution=int(resolution),
@@ -640,7 +376,6 @@ def render_preview(
 def render_camera_video(
     asset_state: Dict[str, Any],
     hdri_file_obj: Any,
-    tone_mapper_name: str,
     hdri_rot: float,
     fps: int,
     num_views: int,
@@ -652,12 +387,12 @@ def render_camera_video(
     progress=gr.Progress(track_tqdm=True),
 ):
     t0 = time.time()
-    session_dir = ensure_session_dir(req)
+    session_dir = CACHE_DIR / str(req.session_hash)
     progress(0.1, desc="Loading SLaT and HDRI")
-    pipeline, slat, hdri_np = load_asset_and_hdri(asset_state, hdri_file_obj, tone_mapper_name)
+    slat, hdri_np = load_asset_and_hdri(asset_state, hdri_file_obj, req)
 
     progress(0.4, desc="Rendering camera path")
-    frames = pipeline.render_camera_path_video(
+    frames = PIPELINE.render_camera_path_video(
         slat, hdri_np,
         num_views=int(num_views), fov=fov, radius=radius,
         hdri_rot_deg=hdri_rot, full_video=full_video, shadow_video=shadow_video,
@@ -674,7 +409,6 @@ def render_camera_video(
 def render_hdri_video(
     asset_state: Dict[str, Any],
     hdri_file_obj: Any,
-    tone_mapper_name: str,
     fps: int,
     num_frames: int,
     yaw: float,
@@ -687,12 +421,12 @@ def render_hdri_video(
     progress=gr.Progress(track_tqdm=True),
 ):
     t0 = time.time()
-    session_dir = ensure_session_dir(req)
+    session_dir = CACHE_DIR / str(req.session_hash)
     progress(0.1, desc="Loading SLaT and HDRI")
-    pipeline, slat, hdri_np = load_asset_and_hdri(asset_state, hdri_file_obj, tone_mapper_name)
+    slat, hdri_np = load_asset_and_hdri(asset_state, hdri_file_obj, req)
 
     progress(0.4, desc="Rendering HDRI rotation")
-    hdri_roll_frames, render_frames = pipeline.render_hdri_rotation_video(
+    hdri_roll_frames, render_frames = PIPELINE.render_hdri_rotation_video(
         slat, hdri_np,
         num_frames=int(num_frames), yaw_deg=yaw, pitch_deg=pitch,
         fov=fov, radius=radius, full_video=full_video, shadow_video=shadow_video,
@@ -710,7 +444,6 @@ def render_hdri_video(
 def export_glb(
     asset_state: Dict[str, Any],
     hdri_file_obj: Any,
-    tone_mapper_name: str,
     hdri_rot: float,
     simplify: float,
     texture_size: int,
@@ -718,12 +451,12 @@ def export_glb(
     progress=gr.Progress(track_tqdm=True),
 ):
     t0 = time.time()
-    session_dir = ensure_session_dir(req)
+    session_dir = CACHE_DIR / str(req.session_hash)
     progress(0.1, desc="Loading SLaT and HDRI")
-    pipeline, slat, hdri_np = load_asset_and_hdri(asset_state, hdri_file_obj, tone_mapper_name)
+    slat, hdri_np = load_asset_and_hdri(asset_state, hdri_file_obj, req)
 
     progress(0.6, desc="Baking PBR textures")
-    glb = pipeline.export_glb_from_slat(
+    glb = PIPELINE.export_glb_from_slat(
         slat, hdri_np,
         hdri_rot_deg=hdri_rot, base_mesh=None,
         simplify=simplify, texture_size=int(texture_size), fill_holes=True,
@@ -1018,7 +751,7 @@ def build_app() -> gr.Blocks:
                                 value=str(DEFAULT_IMAGE) if DEFAULT_IMAGE.exists() else None,
                                 height=400,
                             )
-                            seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed (SLaT)")
+                            seed = gr.Slider(0, MAX_SEED, value=43, step=1, label="Seed (SLaT)")
                             mesh_button = gr.Button("① Generate Mesh", variant="primary", min_width=100)
 
                         with gr.Tab("SLaT", id=1):
@@ -1074,11 +807,10 @@ def build_app() -> gr.Blocks:
                     gr.HTML("<div class='ctrl-strip-title'>Camera &amp; HDRI</div>")
                     with gr.Row():
                         tone_mapper_name = gr.Dropdown(
-                            choices=TONE_MAPPER_CHOICES,
-                            value=TONE_MAPPER_CHOICES[0] if TONE_MAPPER_CHOICES else None,
+                            choices=AVAILABLE_TONE_MAPPERS,
+                            value="AgX",
                             label="Tone Mapper",
                             min_width=120,
-                            allow_custom_value=True,
                         )
                         hdri_rot   = gr.Slider(0, 360,   value=0,   step=1,   label="HDRI Rotation °")
                         resolution = gr.Slider(256, 1024, value=512, step=256, label="Preview Res")
@@ -1088,6 +820,12 @@ def build_app() -> gr.Blocks:
                         fov    = gr.Slider(10, 70,   value=40,  step=1,    label="FoV")
                         radius = gr.Slider(1.0, 4.0, value=2.0, step=0.05, label="Radius")
 
+                    tone_mapper_name.change(
+                        set_tone_mapper,
+                        inputs=[tone_mapper_name],
+                        outputs=[],
+                    )
+
                 with gr.Tabs(elem_classes=["main-output-tabs"]):
 
                     with gr.Tab("Geometry", id=0):
@@ -1103,10 +841,6 @@ def build_app() -> gr.Blocks:
                             export_glb_button = gr.Button("Export PBR GLB", variant="primary", min_width=140)
 
                     with gr.Tab("Preview", id=1):
-                        gr.HTML(
-                            "<p style='font-size:0.78rem;color:#9ca3af;margin:0 0 0.35rem 0;'>"
-                            "Use <b>Camera &amp; HDRI</b> under the tabs, then render.</p>"
-                        )
                         preview_button = gr.Button("Render Preview", variant="primary", min_width=100)
                         gr.HTML("<hr class='divider'>")
                         with gr.Row():
@@ -1181,6 +915,7 @@ def build_app() -> gr.Blocks:
                     else:
                         gr.Markdown("*No `.exr` examples in `assets/hdris`*")
 
+        demo.load(start_session)
         demo.unload(end_session)
 
         source_mode.change(switch_asset_source, inputs=[source_mode], outputs=[source_tabs])
@@ -1193,10 +928,10 @@ def build_app() -> gr.Blocks:
             outputs=[col_img_examples, col_slat_examples],
         )
 
-        for _trigger in (hdri_file.upload, hdri_file.change, tone_mapper_name.change):
+        for _trigger in (hdri_file.upload, hdri_file.change):
             _trigger(
                 preview_hdri,
-                inputs=[hdri_file, tone_mapper_name],
+                inputs=[hdri_file],
                 outputs=[hdri_preview, status_md],
             )
 
@@ -1220,7 +955,7 @@ def build_app() -> gr.Blocks:
 
         preview_button.click(
             render_preview,
-            inputs=[asset_state, hdri_file, tone_mapper_name, hdri_rot,
+            inputs=[asset_state, hdri_file, hdri_rot,
                     yaw, pitch, fov, radius, resolution],
             outputs=[
                 color_output,
@@ -1234,87 +969,36 @@ def build_app() -> gr.Blocks:
 
         camera_video_button.click(
             render_camera_video,
-            inputs=[asset_state, hdri_file, tone_mapper_name, hdri_rot,
+            inputs=[asset_state, hdri_file, hdri_rot,
                     fps, num_views, fov, radius, full_video, shadow_video],
             outputs=[camera_video_output, status_md],
         )
 
         hdri_video_button.click(
             render_hdri_video,
-            inputs=[asset_state, hdri_file, tone_mapper_name,
+            inputs=[asset_state, hdri_file,
                     fps, num_frames, yaw, pitch, fov, radius, full_video, shadow_video],
             outputs=[hdri_roll_video_output, hdri_render_video_output, status_md],
         )
 
         export_glb_button.click(
             export_glb,
-            inputs=[asset_state, hdri_file, tone_mapper_name, hdri_rot, simplify, texture_size],
+            inputs=[asset_state, hdri_file, hdri_rot, simplify, texture_size],
             outputs=[pbr_viewer, status_md],
         )
-
-        clear_button.click(
-            clear_session_dir,
-            outputs=[status_md],
-        ).then(
-            lambda: ({}, None, None, None, None, None, None, None, None, None, None),
-            outputs=[
-                asset_state,
-                mesh_viewer,
-                pbr_viewer,
-                color_output,
-                base_color_output,
-                metallic_output,
-                roughness_output,
-                shadow_output,
-                camera_video_output,
-                hdri_roll_video_output,
-                hdri_render_video_output,
-            ],
-        )
-
     return demo
 
 
 demo = build_app()
-demo.queue(max_size=8)
-
-# Gradio 6: theme/css must be passed to launch(); HF Spaces calls demo.launch() without our __main__ block.
-_orig_blocks_launch = demo.launch
-
-
-def _near_launch(*args: Any, **kwargs: Any):
-    kwargs.setdefault("theme", NEAR_GRADIO_THEME)
-    kwargs.setdefault("css", CUSTOM_CSS)
-    if _CPU_PRELOAD_AT_START:
-        run_model_cpu_preload_blocking()
-    return _orig_blocks_launch(*args, **kwargs)
-
-
-demo.launch = _near_launch  # type: ignore[method-assign]
-
-start_tmp_gradio_pruner()
 
 if __name__ == "__main__":
-    import argparse
 
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--host",
-        type=str,
-        default=os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0"),
-    )
-    parser.add_argument(
-        "--port",
-        type=int,
-        default=int(
-            os.environ.get("PORT", os.environ.get("GRADIO_SERVER_PORT", str(DEFAULT_PORT)))
-        ),
-    )
-    parser.add_argument("--share", action="store_true")
-    args = parser.parse_args()
+    PIPELINE = NeARImageToRelightable3DPipeline.from_pretrained("luh0502/NeAR")
+    GEOMETRY_PIPELINE = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained("tencent/Hunyuan3D-2.1")
+
+    PIPELINE.to("cuda")
+    GEOMETRY_PIPELINE.to("cuda")
 
     demo.launch(
-        server_name=args.host,
-        server_port=args.port,
-        share=args.share,
+        mcp_server=True
     )

trellis/modules/sparse/basic.py

@@ -258,12 +258,8 @@ class SparseTensor:
             )
             new_data._caches = self.data._caches
         elif BACKEND == 'spconv':
-            # Must store `feats` in the tensor that backs `.features`; only setting
-            # `_features` leaves `self.data.features` stale (breaks replace() callers
-            # that refresh feats, e.g. detach().clone() for inference tensors).
-            feats_storage = feats.reshape(feats.shape[0], -1).contiguous()
             new_data = SparseTensorData(
-                feats_storage,
+                self.data.features.reshape(self.data.features.shape[0], -1),
                 self.data.indices,
                 self.data.spatial_shape,
                 self.data.batch_size,

trellis/pipelines/near_image_to_relightable_3d.py

@@ -75,8 +75,9 @@ class NeARImageToRelightable3DPipeline(Pipeline):
 
     def setup_tone_mapper(self, view: str = "AgX") -> None:
         """Initialize the tone mapper used for HDR-to-LDR conversion."""
-        self.tone_mapper = ToneMapper()
-        self.tone_mapper.view = view
+        # Construct with target view in one step so we never `del self.cpu` via the
+        # property setter without a successful rebuild (avoids missing `.cpu`).
+        self.tone_mapper = ToneMapper(view=view)
 
     @staticmethod
     def from_pretrained(path: str) -> "NeARImageToRelightable3DPipeline":
@@ -319,7 +320,7 @@ class NeARImageToRelightable3DPipeline(Pipeline):
     @torch.no_grad()
     def run_with_coords(
         self,
-        image: Image.Image,
+        image: ImageInput,
         coords: torch.Tensor,
         seed: int = 42,
         preprocess_image: bool = True,
@@ -539,8 +540,6 @@ class NeARImageToRelightable3DPipeline(Pipeline):
             hdri_roll_frames.append(ldr)
 
         extr, intr = self.generate_camera(yaw_deg, pitch_deg, radius, fov)
-        extr = extr.to(self.device)
-        intr = intr.to(self.device)
         hs, rfs = self.decoder_pbr_feats(slat)
         render_frames: List[np.ndarray] = []
 
@@ -590,20 +589,17 @@ class NeARImageToRelightable3DPipeline(Pipeline):
         fill_holes: bool = True,
     ) -> trimesh.Trimesh:
         """Export a textured PBR GLB, preferring an externally provided base mesh."""
-        # SLaT feats may be inference tensors (e.g. ZeroGPU). Do not wrap the whole export in
-        # inference_mode: bake_texture(mode='opt') needs loss.backward(). Cloning feats yields
-        # normal tensors so mesh decoder / linears work without the inference-tensor autograd error.
-        slat = slat.replace(slat.feats.detach().clone())
-        if base_mesh is None and "slat_decoder_mesh" in self.models:
-            mesh_out = self.models["slat_decoder_mesh"](slat)
-            base_mesh = mesh_out[0] if isinstance(mesh_out, (list, tuple)) else mesh_out
-        if base_mesh is None:
-            raise ValueError(
-                "export_glb_from_slat requires `base_mesh` or pipeline model `slat_decoder_mesh`"
-            )
+        with torch.inference_mode():
+            if base_mesh is None and "slat_decoder_mesh" in self.models:
+                mesh_out = self.models["slat_decoder_mesh"](slat)
+                base_mesh = mesh_out[0] if isinstance(mesh_out, (list, tuple)) else mesh_out
+            if base_mesh is None:
+                raise ValueError(
+                    "export_glb_from_slat requires `base_mesh` or pipeline model `slat_decoder_mesh`"
+                )
 
-        hdri_cond = self.encode_hdri(hdri_np, hdri_rot_deg)
-        hs, rfs = self.decoder_pbr_feats(slat)
+            hdri_cond = self.encode_hdri(hdri_np, hdri_rot_deg)
+            hs, rfs = self.decoder_pbr_feats(slat)
         return render_utils_rl.to_glb(
             self.models["renderer"],
             hs,