refactor: merge ③④ GPU callbacks for ZeroGPU stability

- Merged generate_videos() and export_pbr_glb() into single generate_renderings() callback
- Avoids CUDA context resets between steps by keeping renderer/tone_mapper valid
- Improved CUDA resource cleanup: explicit del() instead of just None assignment
- Enhanced _ensure_near_on_cuda() with better old object deletion
- Enhanced _teardown_near() with exception handling for robust cleanup
- Updated UI: combined button shows videos + GLB in one action
- Duration increased to 600s to accommodate merged workload

app.py

@@ -1,18 +1,20 @@
 """
-NeAR Gradio Space — ZeroGPU-compatible pipeline.
+NeAR Gradio Space — ZeroGPU-compatible pipeline (REFACTORED).
 
 ZeroGPU contract: every @spaces.GPU call gets a **fresh** CUDA context.
 Any CUDA-backed object (nvdiffrast renderer, tone_mapper) from a previous
-call is invalid in the next call.  Fix: always teardown renderer/tone_mapper
-before returning, re-init at the top of each GPU callback.
+call is invalid in the next call.
 
-UI: 4 linear steps
-  ① Generate Geometry  (Hunyuan3D)
-  ② Generate SLaT      (NeAR — needs mesh from ①)
-  ③ Generate Videos    (NeAR — camera orbit + HDRI rotation)
-  ④ Export PBR GLB     (NeAR — baked PBR mesh)
+REFACTORED ARCHITECTURE:
+- ① Generate Geometry (Hunyuan3D) — independent GPU callback
+- ② Generate SLaT (NeAR) — independent GPU callback  
+- ③④ MERGED: Generate Videos + Export PBR GLB — SINGLE GPU callback
+  * Avoids cross-callback CUDA context resets
+  * Renderer/tone_mapper stay valid throughout rendering and export
+
+Session state = file paths only. No in-memory objects retained between
+steps (except within the unified ③④ callback).
 
-Session state = file paths only. No in-memory SLAT/mesh between callbacks.
 CPU preload runs in a background daemon thread at Space startup (no GPU lease).
 """
 from __future__ import annotations
@@ -170,22 +172,69 @@ def _ensure_near_on_cuda() -> NeARImageToRelightable3DPipeline:
     # ZeroGPU runs one GPU callback at a time so no lock is needed.
     _load_near_cpu_locked()
     assert PIPELINE is not None
+    
+    # ── Cleanup old CUDA objects before re-init ─────────────────────────
+    # Each @GPU callback has a NEW CUDA context; old renderer/tone_mapper
+    # pointers are invalid and must be fully deleted, not just set to None.
+    if hasattr(PIPELINE, 'renderer') and PIPELINE.renderer is not None:
+        try:
+            del PIPELINE.renderer
+        except Exception as e:
+            print(f"[NeAR] warning: failed to delete old renderer: {e}", flush=True)
+        PIPELINE.renderer = None
+    
+    if hasattr(PIPELINE, 'tone_mapper') and PIPELINE.tone_mapper is not None:
+        try:
+            del PIPELINE.tone_mapper
+        except Exception as e:
+            print(f"[NeAR] warning: failed to delete old tone_mapper: {e}", flush=True)
+        PIPELINE.tone_mapper = None
+    
+    _free_cuda()  # Clear any lingering GPU memory
+    
+    # ── Move to new CUDA context and re-init fresh objects ──────────────
     if torch.cuda.is_available():
         PIPELINE.to("cuda")
-    # Always re-init: previous callback's CUDA context is gone.
+    
+    # Initialize fresh renderer/tone_mapper in the new CUDA context
+    print("[NeAR] initializing renderer/tone_mapper in fresh CUDA context…", flush=True)
     PIPELINE.setup_renderer()
     PIPELINE.setup_tone_mapper("AgX")
     return PIPELINE
 
 
 def _teardown_near() -> None:
-    """Release CUDA-backed objects; move weights back to CPU."""
+    """Release CUDA-backed objects; move weights back to CPU.
+    
+    ZeroGPU contract: always teardown CUDA objects before exiting @GPU callback.
+    The next callback will get a completely fresh CUDA context.
+    """
     if PIPELINE is None:
         return
-    PIPELINE.renderer = None
-    PIPELINE.tone_mapper = None
+    
+    print("[NeAR] tearing down renderer/tone_mapper…", flush=True)
+    
+    # Explicitly delete CUDA objects (not just None assignment)
+    if hasattr(PIPELINE, 'renderer'):
+        try:
+            if PIPELINE.renderer is not None:
+                del PIPELINE.renderer
+        except Exception as e:
+            print(f"[NeAR] warning: failed to delete renderer in teardown: {e}", flush=True)
+        PIPELINE.renderer = None
+    
+    if hasattr(PIPELINE, 'tone_mapper'):
+        try:
+            if PIPELINE.tone_mapper is not None:
+                del PIPELINE.tone_mapper
+        except Exception as e:
+            print(f"[NeAR] warning: failed to delete tone_mapper in teardown: {e}", flush=True)
+        PIPELINE.tone_mapper = None
+    
+    # Move model weights back to CPU
     if torch.cuda.is_available():
         PIPELINE.to("cpu")
+    
     _free_cuda()
 
 
@@ -320,7 +369,7 @@ def end_session(req: gr.Request) -> None:
 # ── GPU callbacks ─────────────────────────────────────────────────────────────
 
 @_gpu(duration=240)
-@torch.inference_mode()
+@torch.no_grad()
 def generate_geometry(
     image_input: Optional[Image.Image],
     req: gr.Request,
@@ -356,7 +405,7 @@ def generate_geometry(
 
 
 @_gpu(duration=240)
-@torch.inference_mode()
+@torch.no_grad()
 def generate_slat(
     asset_state: Dict[str, Any],
     image_input: Optional[Image.Image],
@@ -421,9 +470,9 @@ def load_slat_file(
     return state, f"SLaT loaded: `{Path(resolved).name}`"
 
 
-@_gpu(duration=360)
-@torch.inference_mode()
-def generate_videos(
+@_gpu(duration=600)
+@torch.no_grad()
+def generate_renderings(
     asset_state: Dict[str, Any],
     hdri_file_obj: Any,
     hdri_rot: float,
@@ -434,22 +483,42 @@ def generate_videos(
     pitch: float,
     fov: float,
     radius: float,
+    simplify: float,
+    texture_size: int,
     req: gr.Request,
     progress: gr.Progress = gr.Progress(track_tqdm=True),
-) -> tuple[str, str, str]:
-    """③ Camera-orbit video + HDRI-rotation video (one GPU call)."""
+) -> tuple[str, str, str, str, str]:
+    """③④ Unified rendering pipeline: videos + PBR GLB in ONE CUDA context.
+    
+    This merged callback avoids ZeroGPU CUDA context resets between steps.
+    Returns: (cam_video, hdri_video, roll_video, pbr_glb, status_msg)
+    """
     slat_path = _require_slat(asset_state)
     hdri_path = _require_hdri(hdri_file_obj)
     session_dir = CACHE_DIR / str(req.session_hash)
 
+    # ── Load NeAR once (fresh CUDA context) ───────────────────────────
     progress(0.05, desc="Loading NeAR on GPU")
     pipe = _ensure_near_on_cuda()
 
-    progress(0.1, desc="Loading SLaT / HDRI")
+    progress(0.08, desc="Loading SLaT / HDRI")
     slat = pipe.load_slat(slat_path)
     hdri_np = _load_hdri_resized(pipe, hdri_path)
 
-    progress(0.15, desc="Rendering camera-orbit video…")
+    # ── Prepare base mesh (optional, for GLB export) ────────────────────
+    mesh_path = asset_state.get("mesh_path")
+    base_mesh: Optional[trimesh.Trimesh] = None
+    if mesh_path and os.path.isfile(mesh_path):
+        print(f"[NeAR] loading Hunyuan mesh from {mesh_path}…", flush=True)
+        raw = trimesh.load(mesh_path, force="mesh")
+        print(f"[NeAR] raw mesh bounds: {raw.bounds[0].round(3)} → {raw.bounds[1].round(3)}", flush=True)
+        base_mesh = _hunyuan_mesh_to_renderer_space(raw)
+        del raw
+    else:
+        print("[NeAR] no mesh_path — will use SLaT decoder mesh for GLB", flush=True)
+
+    # ── Render videos ──────────────────────────────────────────────────
+    progress(0.12, desc="Rendering camera-orbit video…")
     cam_frames = pipe.render_camera_path_video(
         slat, hdri_np,
         num_views=int(num_cam),
@@ -463,7 +532,7 @@ def generate_videos(
     p_cam = session_dir / "video_camera_orbit.mp4"
     imageio.mimsave(p_cam, cam_frames, fps=int(fps))
     del cam_frames
-    _free_cuda()
+    _free_cuda()  # Free GPU mem temporarily while still in same CUDA context
 
     progress(0.55, desc="Rendering HDRI-rotation video…")
     roll_frames, hdri_frames = pipe.render_hdri_rotation_video(
@@ -480,10 +549,36 @@ def generate_videos(
     p_roll = session_dir / "video_env_roll.mp4"
     imageio.mimsave(p_hdri, hdri_frames, fps=int(fps))
     imageio.mimsave(p_roll, roll_frames, fps=int(fps))
-    del hdri_frames, roll_frames, slat, hdri_np
+    del hdri_frames, roll_frames
+    _free_cuda()
 
+    # ── Export PBR GLB ────────────────────────────────────────────────
+    # (renderer/tone_mapper still valid in same CUDA context)
+    progress(0.85, desc="Baking PBR GLB…")
+    glb = pipe.export_glb_from_slat(
+        slat, hdri_np,
+        hdri_rot_deg=float(hdri_rot),
+        base_mesh=base_mesh,
+        simplify=float(simplify),
+        texture_size=int(texture_size),
+        fill_holes=True,
+    )
+    del slat, hdri_np, base_mesh
+    _free_cuda()
+
+    out_path = session_dir / "near_pbr.glb"
+    glb.export(out_path)
+    del glb
+
+    # ── Cleanup: tear down NeAR before exiting CUDA callback ────────────
     _teardown_near()
-    return str(p_cam), str(p_hdri), str(p_roll)
+
+    msg = (
+        f"**③ Videos ready** → `{Path(p_cam).name}`, "
+        f"`{Path(p_hdri).name}`, `{Path(p_roll).name}` \n\n"
+        f"**④ PBR GLB ready** → `{Path(out_path).name}`"
+    )
+    return str(p_cam), str(p_hdri), str(p_roll), str(out_path), msg
 
 
 def _hunyuan_mesh_to_renderer_space(mesh: trimesh.Trimesh) -> trimesh.Trimesh:
@@ -508,65 +603,6 @@ def _hunyuan_mesh_to_renderer_space(mesh: trimesh.Trimesh) -> trimesh.Trimesh:
         flush=True,
     )
     return mesh
-
-
-@_gpu(duration=180)
-def export_pbr_glb(
-    asset_state: Dict[str, Any],
-    hdri_file_obj: Any,
-    hdri_rot: float,
-    simplify: float,
-    texture_size: int,
-    req: gr.Request,
-    progress: gr.Progress = gr.Progress(track_tqdm=True),
-) -> tuple[str, str]:
-    """④ Export PBR GLB with baked textures, using Hunyuan mesh as base."""
-    slat_path = _require_slat(asset_state)
-    hdri_path = _require_hdri(hdri_file_obj)
-    session_dir = CACHE_DIR / str(req.session_hash)
-
-    # Use Hunyuan geometry as base mesh when available; fall back to TRELLIS decoder.
-    mesh_path = asset_state.get("mesh_path")
-    base_mesh: Optional[trimesh.Trimesh] = None
-    if mesh_path and os.path.isfile(mesh_path):
-        raw = trimesh.load(mesh_path, force="mesh")
-        print(
-            f"[NeAR] raw Hunyuan mesh: bounds={raw.bounds[0].round(3)} → {raw.bounds[1].round(3)}",
-            flush=True,
-        )
-        base_mesh = _hunyuan_mesh_to_renderer_space(raw)
-        del raw
-    else:
-        print("[NeAR] no mesh_path in state — falling back to slat_decoder_mesh", flush=True)
-
-    progress(0.1, desc="Loading NeAR on GPU")
-    pipe = _ensure_near_on_cuda()
-
-    progress(0.2, desc="Loading SLaT / HDRI")
-    slat = pipe.load_slat(slat_path)
-    hdri_np = _load_hdri_resized(pipe, hdri_path)
-
-    progress(0.4, desc="Baking PBR GLB…")
-    glb = pipe.export_glb_from_slat(
-        slat, hdri_np,
-        hdri_rot_deg=float(hdri_rot),
-        base_mesh=base_mesh,
-        simplify=float(simplify),
-        texture_size=int(texture_size),
-        fill_holes=True,
-    )
-    del slat, hdri_np
-    _free_cuda()
-
-    out_path = session_dir / "near_pbr.glb"
-    glb.export(out_path)
-    del glb
-
-    _teardown_near()
-    return str(out_path), f"**④ PBR GLB** → `{out_path.name}`"
-
-
-def clear_cache(req: gr.Request) -> str:
     session_dir = CACHE_DIR / str(req.session_hash)
     shutil.rmtree(session_dir, ignore_errors=True)
     session_dir.mkdir(parents=True, exist_ok=True)
@@ -675,8 +711,8 @@ def build_app() -> gr.Blocks:
         asset_state = gr.State({})
 
         gr.Markdown(
-            "## NeAR — Relightable 3D (ZeroGPU)\n"
-            "**① Geometry** → **② SLaT** → **③ Videos** and/or **④ PBR GLB**\n\n"
+            "## NeAR — Relightable 3D (ZeroGPU Optimized)\n"
+            "**① Geometry** → **② SLaT** → **③④ Videos + PBR GLB** (merged for stability)\n\n"
             "Tip: after generating geometry, swap the input image before running **② SLaT** "
             "to apply a different texture style to the same shape — geometry and appearance are decoupled."
         )
@@ -717,8 +753,8 @@ def build_app() -> gr.Blocks:
                 hdri_rot = gr.Slider(0, 360, value=0, step=1, label="HDRI rotation °")
 
                 gr.HTML('<p class="section-kicker" style="margin:2px 0 2px;padding:0">Actions</p>')
-                btn_videos = gr.Button("③ Generate Videos", variant="primary")
-                btn_glb = gr.Button("④ Export PBR GLB", variant="primary")
+                btn_videos = gr.Button("③④ Generate Videos & GLB", variant="primary")
+                btn_glb = gr.Button("④ Export PBR GLB", variant="primary", visible=False)
                 with gr.Accordion("Video / export settings", open=False):
                     fps = gr.Slider(8, 48, value=24, step=1, label="FPS")
                     num_cam = gr.Slider(8, 96, value=36, step=4, label="Camera-orbit frames")
@@ -784,12 +820,10 @@ def build_app() -> gr.Blocks:
                        [asset_state, status])
         btn_load_slat.click(load_slat_file, [slat_upload, slat_path_txt],
                             [asset_state, status])
-        btn_videos.click(generate_videos,
-                         [asset_state, hdri_file, hdri_rot, fps, num_cam, num_hdri, yaw, pitch, fov, radius],
-                         [vid_cam, vid_hdri, vid_roll])
-        btn_glb.click(export_pbr_glb,
-                      [asset_state, hdri_file, hdri_rot, simplify, tex_size],
-                      [glb_view, status])
+        # ③④ Unified render callback: videos + GLB in ONE CUDA context
+        btn_videos.click(generate_renderings,
+                         [asset_state, hdri_file, hdri_rot, fps, num_cam, num_hdri, yaw, pitch, fov, radius, simplify, tex_size],
+                         [vid_cam, vid_hdri, vid_roll, glb_view, status])
         btn_clear.click(clear_cache, [], [status])
 
     return demo