Reverb/open3dforge

1"""2Workspace management for Open3DForge.3 4Single-user pattern: one persistent workspace folder.5No session IDs, no multi-tenancy.6 7Layout:8    workspace/9        current/    -- active work, overwritten per generation10        exports/    -- finished asset zips, kept indefinitely11        presets/    -- saved parameter configurations (JSON)12        history/    -- thumbnails + metadata of past assets13"""14 15from __future__ import annotations16 17import json18import os19import shutil20import time21from dataclasses import dataclass, field22from pathlib import Path23from typing import Any24 25# ---------------------------------------------------------------------------26# Paths27# ---------------------------------------------------------------------------28 29ROOT = Path(__file__).resolve().parent.parent30WORKSPACE = ROOT / "workspace"31CURRENT = WORKSPACE / "current"32EXPORTS = WORKSPACE / "exports"33PRESETS = WORKSPACE / "presets"34HISTORY = WORKSPACE / "history"35 36# Subdirectories of `current/` created on each generation37CURRENT_TEXTURES = CURRENT / "textures"38CURRENT_LODS = CURRENT / "lods"39 40 41def ensure_dirs() -> None:42    """Create all workspace directories if missing. Safe to call repeatedly."""43    for p in (WORKSPACE, CURRENT, EXPORTS, PRESETS, HISTORY,44              CURRENT_TEXTURES, CURRENT_LODS):45        p.mkdir(parents=True, exist_ok=True)46 47 48def reset_current() -> None:49    """Clear `current/` for a fresh asset. Called at the start of generation."""50    if CURRENT.exists():51        shutil.rmtree(CURRENT)52    CURRENT.mkdir(parents=True, exist_ok=True)53    CURRENT_TEXTURES.mkdir(parents=True, exist_ok=True)54    CURRENT_LODS.mkdir(parents=True, exist_ok=True)55 56 57# ---------------------------------------------------------------------------58# Asset state (in-memory representation of what's in `current/`)59# ---------------------------------------------------------------------------60 61@dataclass62class AssetState:63    """Tracks what files exist in `current/` and the pipeline progress.64 65    Updated by each stage as it completes. Used by the UI to enable/disable66    buttons and show status indicators.67    """68 69    # Input70    input_images: list[Path] = field(default_factory=list)71 72    # Stage 1 outputs73    high_poly_glb: Path | None = None         # raw TRELLIS.2 output, kept for baking74    raw_gen_glb: Path | None = None           # decimated base from generation75 76    # Stage 2 outputs77    repaired_glb: Path | None = None78    cleaned_glb: Path | None = None79    low_poly_glb: Path | None = None          # post-decimation working mesh80    unwrapped_glb: Path | None = None         # has UV coordinates81    final_glb: Path | None = None             # all post-processing done82 83    # Textures (Stage 2 baking)84    albedo_png: Path | None = None85    normal_gl_png: Path | None = None86    normal_dx_png: Path | None = None87    roughness_png: Path | None = None88    metallic_png: Path | None = None89    ao_png: Path | None = None90    orm_png: Path | None = None               # UE5-packed AO/Rough/Metal91    metallic_smoothness_png: Path | None = None  # Unity-packed92 93    # LODs94    lod_glbs: list[Path] = field(default_factory=list)95 96    # Collision97    collision_glb: Path | None = None98 99    # Rigging (Stage 3)100    rigged_glb: Path | None = None101    rigged_fbx: Path | None = None102 103    # Metadata104    asset_name: str = "untitled"105    generated_at: float = field(default_factory=time.time)106    model_used: str = ""                      # "TRELLIS.2" or "Hunyuan3D-2"107    face_count: int = 0108    vertex_count: int = 0109 110    def to_dict(self) -> dict[str, Any]:111        """Serialise for status display / debug."""112        out: dict[str, Any] = {}113        for k, v in self.__dict__.items():114            if isinstance(v, Path):115                out[k] = str(v) if v else None116            elif isinstance(v, list):117                out[k] = [str(p) for p in v]118            else:119                out[k] = v120        return out121 122 123# ---------------------------------------------------------------------------124# Filesystem-based state persistence125#126# ZeroGPU runs @spaces.GPU functions in a forked subprocess.  Any writes to127# module-level variables (like _state) happen in the subprocess and are128# invisible to the parent Gradio process.  To survive the process boundary we129# write state to a JSON file inside CURRENT/ and always read back from there.130# ---------------------------------------------------------------------------131 132_META_FILE = CURRENT / ".meta.json"133 134# File names that map to AssetState path attributes (order = preference)135_PATH_ATTRS: list[tuple[str, Path]] = [136    ("rigged_fbx",     CURRENT / "rigged.fbx"),137    ("rigged_glb",     CURRENT / "rigged.glb"),138    ("final_glb",      CURRENT / "scaled.glb"),139    ("final_glb",      CURRENT / "pivoted.glb"),140    ("unwrapped_glb",  CURRENT / "unwrapped.glb"),141    ("low_poly_glb",   CURRENT / "low_poly.glb"),142    ("cleaned_glb",    CURRENT / "cleaned.glb"),143    ("repaired_glb",   CURRENT / "repaired.glb"),144    ("raw_gen_glb",    CURRENT / "raw_gen.glb"),145    ("high_poly_glb",  CURRENT / "high_poly.glb"),146    ("normal_dx_png",  CURRENT / "textures" / "normal_dx.png"),147    ("normal_gl_png",  CURRENT / "textures" / "normal_gl.png"),148    ("albedo_png",     CURRENT / "textures" / "albedo.png"),149    ("roughness_png",  CURRENT / "textures" / "roughness.png"),150    ("metallic_png",   CURRENT / "textures" / "metallic.png"),151    ("ao_png",         CURRENT / "textures" / "ao.png"),152    ("orm_png",        CURRENT / "textures" / "orm.png"),153    ("collision_glb",  CURRENT / "collision.glb"),154]155 156 157def _build_state_from_disk() -> AssetState:158    """Reconstruct AssetState by scanning CURRENT/ and reading .meta.json."""159    state = AssetState()160 161    # Read persisted metadata (face count, model name, etc.)162    if _META_FILE.exists():163        try:164            meta = json.loads(_META_FILE.read_text())165            state.asset_name  = meta.get("asset_name", "untitled")166            state.model_used  = meta.get("model_used", "")167            state.face_count  = meta.get("face_count", 0)168            state.vertex_count = meta.get("vertex_count", 0)169        except Exception:170            pass171 172    # Populate path attributes from filesystem173    seen_attrs: set[str] = set()174    for attr, path in _PATH_ATTRS:175        if path.exists() and attr not in seen_attrs:176            setattr(state, attr, path)177            seen_attrs.add(attr)178 179    # LODs180    lod_dir = CURRENT / "lods"181    if lod_dir.exists():182        state.lod_glbs = sorted(lod_dir.glob("LOD*.glb"))183 184    return state185 186 187def flush_meta(state: AssetState) -> None:188    """Write lightweight metadata to disk so the parent process can read it."""189    try:190        _META_FILE.write_text(json.dumps({191            "asset_name":   state.asset_name,192            "model_used":   state.model_used,193            "face_count":   state.face_count,194            "vertex_count": state.vertex_count,195        }))196    except Exception:197        pass198 199 200# Module-level singleton — kept for in-process use (e.g. stage2 steps that201# run in the same process as Gradio).  Always prefer get_state() which syncs202# from disk first, making it safe across the ZeroGPU process boundary.203_state: AssetState = AssetState()204 205 206def get_state() -> AssetState:207    """Return current state, rebuilding from disk to handle ZeroGPU isolation."""208    global _state209    _state = _build_state_from_disk()210    return _state211 212 213def reset_state() -> AssetState:214    """Replace the global state with a fresh one. Returns the new state."""215    global _state216    _state = AssetState()217    return _state218 219 220# ---------------------------------------------------------------------------221# Presets (JSON-on-disk, loaded as dicts)222# ---------------------------------------------------------------------------223 224def list_presets() -> list[str]:225    """Return preset names (filenames without .json), sorted."""226    if not PRESETS.exists():227        return []228    return sorted(p.stem for p in PRESETS.glob("*.json"))229 230 231def load_preset(name: str) -> dict[str, Any]:232    """Load a preset by name. Raises FileNotFoundError if missing."""233    path = PRESETS / f"{name}.json"234    if not path.exists():235        raise FileNotFoundError(f"Preset not found: {name}")236    with path.open("r", encoding="utf-8") as f:237        return json.load(f)238 239 240def save_preset(name: str, config: dict[str, Any]) -> Path:241    """Save a preset as JSON. Overwrites if exists."""242    safe_name = _sanitize_filename(name)243    path = PRESETS / f"{safe_name}.json"244    with path.open("w", encoding="utf-8") as f:245        json.dump(config, f, indent=2, sort_keys=True)246    return path247 248 249def delete_preset(name: str) -> bool:250    """Delete a preset. Returns True if deleted, False if it didn't exist."""251    path = PRESETS / f"{name}.json"252    if path.exists():253        path.unlink()254        return True255    return False256 257 258def _sanitize_filename(name: str) -> str:259    """Strip path separators and unsafe chars from a filename stem."""260    return "".join(c for c in name if c.isalnum() or c in "_-").strip("_-") or "preset"261 262 263# ---------------------------------------------------------------------------264# Workspace stats (for UI status bar)265# ---------------------------------------------------------------------------266 267def workspace_size_mb() -> float:268    """Total size of the workspace in MB."""269    total = 0270    if WORKSPACE.exists():271        for path in WORKSPACE.rglob("*"):272            if path.is_file():273                total += path.stat().st_size274    return total / (1024 * 1024)275 276 277def current_size_mb() -> float:278    """Size of `current/` only (active work)."""279    total = 0280    if CURRENT.exists():281        for path in CURRENT.rglob("*"):282            if path.is_file():283                total += path.stat().st_size284    return total / (1024 * 1024)285 286 287def export_count() -> int:288    """How many exported zips exist."""289    if not EXPORTS.exists():290        return 0291    return len(list(EXPORTS.glob("*.zip")))292 293 294# ---------------------------------------------------------------------------295# Module init296# ---------------------------------------------------------------------------297 298# Auto-create folders on import so the app never crashes on a fresh checkout299ensure_dirs()300