Upload 4 files

src/quota.py

@@ -0,0 +1,148 @@
+"""
+ZeroGPU quota tracking.
+
+HF Pro gives 1500s (25 min) of H200 time per day. There's no official Python
+API to query remaining quota directly — we track it locally per-call.
+
+This module:
+- Records GPU time consumed via the `@spaces.GPU` decorated functions
+- Provides estimates for upcoming operations
+- Resets daily (per-day file on disk)
+
+Note: this is approximate. The authoritative source is HF's quota error if
+you go over. Our tracking is for UX (showing "~18/25 min used today").
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+from .workspace import WORKSPACE
+
+QUOTA_FILE = WORKSPACE / "quota_log.json"
+DAILY_QUOTA_SECONDS = 1500   # 25 minutes for HF Pro
+OVERAGE_RATE_PER_SECOND = 1.0 / 600.0   # $1 per 600s (10 min)
+
+
+@dataclass
+class QuotaState:
+    date: str                # YYYY-MM-DD (UTC)
+    used_seconds: float
+    operations: list[dict]   # log of recent operations
+
+    def remaining_seconds(self) -> float:
+        return max(0.0, DAILY_QUOTA_SECONDS - self.used_seconds)
+
+    def usage_fraction(self) -> float:
+        return min(1.0, self.used_seconds / DAILY_QUOTA_SECONDS)
+
+    def overage_cost_usd(self) -> float:
+        overage = max(0.0, self.used_seconds - DAILY_QUOTA_SECONDS)
+        return overage * OVERAGE_RATE_PER_SECOND
+
+
+def _today_utc() -> str:
+    return time.strftime("%Y-%m-%d", time.gmtime())
+
+
+def _load_or_new() -> QuotaState:
+    today = _today_utc()
+    if QUOTA_FILE.exists():
+        try:
+            with QUOTA_FILE.open("r") as f:
+                data = json.load(f)
+            if data.get("date") == today:
+                return QuotaState(
+                    date=data["date"],
+                    used_seconds=float(data["used_seconds"]),
+                    operations=data.get("operations", []),
+                )
+        except (json.JSONDecodeError, KeyError, ValueError):
+            pass
+    # Fresh day or corrupted file
+    return QuotaState(date=today, used_seconds=0.0, operations=[])
+
+
+def _save(state: QuotaState) -> None:
+    with QUOTA_FILE.open("w") as f:
+        json.dump(
+            {
+                "date": state.date,
+                "used_seconds": state.used_seconds,
+                "operations": state.operations[-50:],  # keep last 50
+            },
+            f,
+            indent=2,
+        )
+
+
+def get_state() -> QuotaState:
+    """Get current quota state, refreshed for today."""
+    return _load_or_new()
+
+
+def record_usage(operation: str, seconds: float) -> QuotaState:
+    """Record a completed GPU operation. Returns updated state."""
+    state = _load_or_new()
+    state.used_seconds += seconds
+    state.operations.append({
+        "op": operation,
+        "seconds": round(seconds, 2),
+        "timestamp": time.time(),
+    })
+    _save(state)
+    return state
+
+
+# ---------------------------------------------------------------------------
+# Per-operation estimates (used by UI to warn before expensive operations)
+# ---------------------------------------------------------------------------
+
+ESTIMATES = {
+    # Stage 1
+    "generate_trellis2_fast": 30,
+    "generate_trellis2_balanced": 60,
+    "generate_trellis2_hero": 90,
+    "generate_hunyuan3d": 60,
+    # Stage 2 — baking (nvdiffrast is fast)
+    "bake_normal_2k": 5,
+    "bake_normal_4k": 12,
+    "bake_albedo": 2,
+    "bake_materials": 3,
+    "bake_ao_fast": 3,
+    "bake_ao_standard": 10,
+    "bake_ao_high": 30,
+    # Stage 2 — optional
+    "inpaint_sdxl": 30,
+    # Stage 3
+    "auto_rig": 40,
+}
+
+
+def estimate(operation: str) -> int:
+    """Get the typical GPU duration for an operation, in seconds.
+
+    Used to:
+      - set `@spaces.GPU(duration=N)` correctly
+      - show cost warnings in the UI before triggering an operation
+    """
+    return ESTIMATES.get(operation, 60)
+
+
+def format_status() -> str:
+    """One-line quota summary for the UI status bar."""
+    s = get_state()
+    used_min = s.used_seconds / 60
+    total_min = DAILY_QUOTA_SECONDS / 60
+    remaining_min = s.remaining_seconds() / 60
+
+    if s.used_seconds >= DAILY_QUOTA_SECONDS:
+        cost = s.overage_cost_usd()
+        return f"⚠️ Quota: {used_min:.1f}/{total_min:.0f} min (overage: ${cost:.2f})"
+    elif s.usage_fraction() > 0.8:
+        return f"⚠️ Quota: {used_min:.1f}/{total_min:.0f} min ({remaining_min:.1f} min left)"
+    else:
+        return f"Quota: {used_min:.1f}/{total_min:.0f} min H200 today"

src/ui_helpers.py

@@ -0,0 +1,119 @@
+"""
+UI helpers — shared functions across Gradio tabs.
+
+Keeps `app.py` focused on layout, not formatting logic.
+"""
+
+from __future__ import annotations
+
+from . import quota, workspace
+
+
+def get_status_bar() -> str:
+    """Build the global status bar shown at the bottom of every tab.
+
+    Markdown-formatted single line with quota + workspace info.
+    """
+    state = workspace.get_state()
+    parts = [
+        quota.format_status(),
+        f"📁 Workspace: {workspace.current_size_mb():.1f} MB",
+        f"📦 Exports: {workspace.export_count()}",
+    ]
+    if state.model_used:
+        parts.append(f"🔧 Last: {state.model_used}")
+    return " · ".join(parts)
+
+
+def get_asset_summary() -> str:
+    """Multi-line summary of what's currently in `workspace/current/`.
+
+    Used in the viewer panel to show pipeline progress.
+    """
+    state = workspace.get_state()
+    lines = [f"**Asset:** `{state.asset_name}`"]
+
+    if state.face_count:
+        lines.append(f"- Faces: {state.face_count:,}")
+        lines.append(f"- Vertices: {state.vertex_count:,}")
+
+    stages = []
+    if state.high_poly_glb:
+        stages.append("✓ Generated")
+    if state.cleaned_glb:
+        stages.append("✓ Cleaned")
+    if state.low_poly_glb:
+        stages.append("✓ Decimated")
+    if state.unwrapped_glb:
+        stages.append("✓ UV unwrapped")
+    if state.normal_dx_png or state.normal_gl_png:
+        stages.append("✓ Normal baked")
+    if state.albedo_png:
+        stages.append("✓ Albedo baked")
+    if state.orm_png or (state.roughness_png and state.metallic_png):
+        stages.append("✓ PBR maps")
+    if state.lod_glbs:
+        stages.append(f"✓ LODs ({len(state.lod_glbs)})")
+    if state.collision_glb:
+        stages.append("✓ Collision")
+    if state.rigged_glb or state.rigged_fbx:
+        stages.append("✓ Rigged")
+
+    if stages:
+        lines.append("")
+        lines.append("**Progress:**")
+        for s in stages:
+            lines.append(f"- {s}")
+    else:
+        lines.append("")
+        lines.append("*No asset loaded yet. Start at the Generate tab.*")
+
+    return "\n".join(lines)
+
+
+def get_viewer_model_path() -> str | None:
+    """Pick the best GLB to show in the 3D viewer right now.
+
+    Prefers the most-processed version available.
+    """
+    state = workspace.get_state()
+    # Order: most processed → least processed
+    candidates = (
+        state.rigged_glb,
+        state.final_glb,
+        state.unwrapped_glb,
+        state.low_poly_glb,
+        state.cleaned_glb,
+        state.repaired_glb,
+        state.raw_gen_glb,
+        state.high_poly_glb,
+    )
+    for path in candidates:
+        if path and path.exists():
+            return str(path)
+    return None
+
+
+def quota_warning(operation: str) -> str:
+    """Generate a warning if an operation would exceed the daily quota.
+
+    Returns an empty string if the operation fits comfortably.
+    """
+    estimated = quota.estimate(operation)
+    state = quota.get_state()
+    remaining = state.remaining_seconds()
+
+    if estimated > remaining:
+        overage = estimated - remaining
+        cost = overage * quota.OVERAGE_RATE_PER_SECOND
+        return (
+            f"⚠️ **Quota warning:** This will use ~{estimated}s but you only "
+            f"have {remaining:.0f}s left today. "
+            f"Overage cost: ~${cost:.2f}"
+        )
+    elif estimated > remaining * 0.5:
+        return (
+            f"ℹ️ Estimated GPU time: ~{estimated}s "
+            f"({remaining:.0f}s remaining today)"
+        )
+    return ""

src/workspace.py

@@ -0,0 +1,219 @@
+"""
+Workspace management for Open3DForge.
+
+Single-user pattern: one persistent workspace folder.
+No session IDs, no multi-tenancy.
+
+Layout:
+    workspace/
+        current/    -- active work, overwritten per generation
+        exports/    -- finished asset zips, kept indefinitely
+        presets/    -- saved parameter configurations (JSON)
+        history/    -- thumbnails + metadata of past assets
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+# ---------------------------------------------------------------------------
+# Paths
+# ---------------------------------------------------------------------------
+
+ROOT = Path(__file__).resolve().parent.parent
+WORKSPACE = ROOT / "workspace"
+CURRENT = WORKSPACE / "current"
+EXPORTS = WORKSPACE / "exports"
+PRESETS = WORKSPACE / "presets"
+HISTORY = WORKSPACE / "history"
+
+# Subdirectories of `current/` created on each generation
+CURRENT_TEXTURES = CURRENT / "textures"
+CURRENT_LODS = CURRENT / "lods"
+
+
+def ensure_dirs() -> None:
+    """Create all workspace directories if missing. Safe to call repeatedly."""
+    for p in (WORKSPACE, CURRENT, EXPORTS, PRESETS, HISTORY,
+              CURRENT_TEXTURES, CURRENT_LODS):
+        p.mkdir(parents=True, exist_ok=True)
+
+
+def reset_current() -> None:
+    """Clear `current/` for a fresh asset. Called at the start of generation."""
+    if CURRENT.exists():
+        shutil.rmtree(CURRENT)
+    CURRENT.mkdir(parents=True, exist_ok=True)
+    CURRENT_TEXTURES.mkdir(parents=True, exist_ok=True)
+    CURRENT_LODS.mkdir(parents=True, exist_ok=True)
+
+
+# ---------------------------------------------------------------------------
+# Asset state (in-memory representation of what's in `current/`)
+# ---------------------------------------------------------------------------
+
+@dataclass
+class AssetState:
+    """Tracks what files exist in `current/` and the pipeline progress.
+
+    Updated by each stage as it completes. Used by the UI to enable/disable
+    buttons and show status indicators.
+    """
+
+    # Input
+    input_images: list[Path] = field(default_factory=list)
+
+    # Stage 1 outputs
+    high_poly_glb: Path | None = None         # raw TRELLIS.2 output, kept for baking
+    raw_gen_glb: Path | None = None           # decimated base from generation
+
+    # Stage 2 outputs
+    repaired_glb: Path | None = None
+    cleaned_glb: Path | None = None
+    low_poly_glb: Path | None = None          # post-decimation working mesh
+    unwrapped_glb: Path | None = None         # has UV coordinates
+    final_glb: Path | None = None             # all post-processing done
+
+    # Textures (Stage 2 baking)
+    albedo_png: Path | None = None
+    normal_gl_png: Path | None = None
+    normal_dx_png: Path | None = None
+    roughness_png: Path | None = None
+    metallic_png: Path | None = None
+    ao_png: Path | None = None
+    orm_png: Path | None = None               # UE5-packed AO/Rough/Metal
+    metallic_smoothness_png: Path | None = None  # Unity-packed
+
+    # LODs
+    lod_glbs: list[Path] = field(default_factory=list)
+
+    # Collision
+    collision_glb: Path | None = None
+
+    # Rigging (Stage 3)
+    rigged_glb: Path | None = None
+    rigged_fbx: Path | None = None
+
+    # Metadata
+    asset_name: str = "untitled"
+    generated_at: float = field(default_factory=time.time)
+    model_used: str = ""                      # "TRELLIS.2" or "Hunyuan3D-2"
+    face_count: int = 0
+    vertex_count: int = 0
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialise for status display / debug."""
+        out: dict[str, Any] = {}
+        for k, v in self.__dict__.items():
+            if isinstance(v, Path):
+                out[k] = str(v) if v else None
+            elif isinstance(v, list):
+                out[k] = [str(p) for p in v]
+            else:
+                out[k] = v
+        return out
+
+
+# Module-level singleton: in a single-user app, this is fine.
+# In Gradio, requests are serialized so there's no race condition at the
+# typical interaction level. Real concurrent generation would need locking.
+_state: AssetState = AssetState()
+
+
+def get_state() -> AssetState:
+    return _state
+
+
+def reset_state() -> AssetState:
+    """Replace the global state with a fresh one. Returns the new state."""
+    global _state
+    _state = AssetState()
+    return _state
+
+
+# ---------------------------------------------------------------------------
+# Presets (JSON-on-disk, loaded as dicts)
+# ---------------------------------------------------------------------------
+
+def list_presets() -> list[str]:
+    """Return preset names (filenames without .json), sorted."""
+    if not PRESETS.exists():
+        return []
+    return sorted(p.stem for p in PRESETS.glob("*.json"))
+
+
+def load_preset(name: str) -> dict[str, Any]:
+    """Load a preset by name. Raises FileNotFoundError if missing."""
+    path = PRESETS / f"{name}.json"
+    if not path.exists():
+        raise FileNotFoundError(f"Preset not found: {name}")
+    with path.open("r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def save_preset(name: str, config: dict[str, Any]) -> Path:
+    """Save a preset as JSON. Overwrites if exists."""
+    safe_name = _sanitize_filename(name)
+    path = PRESETS / f"{safe_name}.json"
+    with path.open("w", encoding="utf-8") as f:
+        json.dump(config, f, indent=2, sort_keys=True)
+    return path
+
+
+def delete_preset(name: str) -> bool:
+    """Delete a preset. Returns True if deleted, False if it didn't exist."""
+    path = PRESETS / f"{name}.json"
+    if path.exists():
+        path.unlink()
+        return True
+    return False
+
+
+def _sanitize_filename(name: str) -> str:
+    """Strip path separators and unsafe chars from a filename stem."""
+    return "".join(c for c in name if c.isalnum() or c in "_-").strip("_-") or "preset"
+
+
+# ---------------------------------------------------------------------------
+# Workspace stats (for UI status bar)
+# ---------------------------------------------------------------------------
+
+def workspace_size_mb() -> float:
+    """Total size of the workspace in MB."""
+    total = 0
+    if WORKSPACE.exists():
+        for path in WORKSPACE.rglob("*"):
+            if path.is_file():
+                total += path.stat().st_size
+    return total / (1024 * 1024)
+
+
+def current_size_mb() -> float:
+    """Size of `current/` only (active work)."""
+    total = 0
+    if CURRENT.exists():
+        for path in CURRENT.rglob("*"):
+            if path.is_file():
+                total += path.stat().st_size
+    return total / (1024 * 1024)
+
+
+def export_count() -> int:
+    """How many exported zips exist."""
+    if not EXPORTS.exists():
+        return 0
+    return len(list(EXPORTS.glob("*.zip")))
+
+
+# ---------------------------------------------------------------------------
+# Module init
+# ---------------------------------------------------------------------------
+
+# Auto-create folders on import so the app never crashes on a fresh checkout
+ensure_dirs()